sientia-dataops-model-manager/tests/laborious/activities/test_opc.py

from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
from pandas import DataFrame
from pytest import fixture, mark
import pytest_asyncio
from sientia_do.notifications.models import NotificationLevel

from model_manager.activities.opc import OPC

metadata = {
    "metadata": {
        "model_id": "test_model",
        "model_name": "test_model",
        "workflow_name": "test_workflow",
        "schema_name": "test_schedule",
    },
}


def test__init__():
    servers = {
        'server1': 'config'
    }
    opc = OPC(
        opc_servers=servers,
        logger=MagicMock(),
        notification_handler=MagicMock()
    )

    assert opc.opc_servers == servers
    assert opc.opc_repository == {}


@mark.asyncio
@patch("model_manager.activities.opc.OpcRepository")
@patch("model_manager.activities.opc.OPC.send_notification")
async def test_init_opc(mock_send_notification, mock_opc_repository):
    mock_logger = MagicMock()
    server1 = MagicMock(
        connect=AsyncMock(return_value=(True, {})),
        write_data=AsyncMock(return_value=(True, {}))
    )
    server2 = MagicMock(
        connect=AsyncMock(return_value=(True, {})),
        write_data=AsyncMock(return_value=(True, {}))
    )
    server3 = MagicMock(
        connect=AsyncMock(return_value=(False, {
            'notification_id': 'OPC_CONNECTION_ERROR_server3',
            'message': 'Failed to connect to OPC server: Test error',
            'block': 'opc_repository',
            'level': NotificationLevel.ERROR,
            'attachment_content': 'Test error'
        })),
        write_data=AsyncMock(return_value=(True, {}))
    )
    mock_opc_repository.side_effect = [server1, server2, server3]
    mock_notification_handler = MagicMock()
    servers = {
        'server1': {
            'id': 'server1',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        },
        'server2': {
            'id': 'server2',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        },
        'server3': {
            'id': 'server3',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        }
    }
    opc = OPC(
        opc_servers=servers,
        logger=mock_logger,
        notification_handler=mock_notification_handler
    )
    await opc.init_opc()

    assert opc.opc_servers == servers
    assert opc.logger == mock_logger
    assert opc.notification_handler == mock_notification_handler
    assert opc.opc_repository['server1'] == server1
    assert opc.opc_repository['server2'] == server2

    mock_opc_repository.assert_has_calls([
        call(
            id="server1",
            url="http://localhost:8080",
            logger=mock_logger,
            server_uri="opc.tcp://localhost:4840",
            cert_path="",
            private_key_path="",
            server_cert_path="",
            notification_handler=mock_notification_handler,
            reconnection_interval=60,
            pod_id='localhost'
        ),
    ])
    mock_opc_repository.assert_has_calls([
        call(
            id="server2",
            url="http://localhost:8080",
            logger=mock_logger,
            server_uri="opc.tcp://localhost:4840",
            cert_path="",
            private_key_path="",
            server_cert_path="",
            notification_handler=mock_notification_handler,
            reconnection_interval=60,
            pod_id='localhost'
        )
    ])

    server1.connect.assert_called_once()
    server2.connect.assert_called_once()

    mock_send_notification.assert_has_calls([
        call(
            metadata={
                'model_id': '-',
                'model_name': '-',
                'workflow_name': '-',
                'schedule_name': 'INITIALIZATION'
            },
            notification_id="OPC_CONNECTION_ERROR_server3",
            message="Failed to connect to OPC server: Test error",
            block="opc_repository",
            level=NotificationLevel.ERROR,
            attachment_content=ANY
        )
    ])


@pytest_asyncio.fixture
@patch("model_manager.activities.opc.OpcRepository")
async def opc(mock_opc_repository):
    servers = {
        'server1': {
            'id': 'server1',
            'url': 'http://localhost:8080',
            'server_uri': 'opc.tcp://localhost:4840',
            'cert_path': '',
            'private_key_path': '',
            'server_cert_path': '',
            'reconnection_interval': 60,
        }
    }

    mock_opc_repository.return_value.write_data = AsyncMock(
        return_value=(True, {})
    )
    mock_opc_repository.return_value.connect = AsyncMock(
        return_value=(True, {})
    )
    opc = OPC(
        opc_servers=servers,
        logger=MagicMock(),
        notification_handler=MagicMock()
    )
    await opc.init_opc()
    opc.send_notification = MagicMock()
    return opc


WRITE_DATA_CASES = [
    ('tag1', 'int', 50),
    ('tag2', 'float', 50.5),
    ('tag3', 'bool', True),
    ('tag4', 'string', 'test'),
]


@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
async def test_write_data_success(opc, tag, data_type, data):
    result = await opc.write_data(server_id='server1', tag=tag, data=data,
                                  data_type=data_type, tag_type='prediction', metadata=metadata)
    assert result is True
    opc.opc_repository['server1'].write_data.assert_called_once_with(
        tag, data, data_type, opc.logger, metadata)


@mark.asyncio
async def test_write_data_failed(opc):
    opc.opc_repository['server1'].write_data.return_value = (False, {
        'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
        'message': 'Failed to write data to OPC server: Test error',
        'block': 'opc_repository',
        'level': NotificationLevel.ERROR,
        'attachment_content': 'Test error'
    })

    result = await opc.write_data(server_id='server1', tag='tag1', data=50,
                                  data_type='int', tag_type='prediction', metadata=metadata)
    assert result is False

    opc.send_notification.assert_called_once_with(
        metadata=metadata,
        notification_id="OPC_WRITE_DATA_ERROR_server1",
        message="Failed to write data to OPC server: Test error",
        block="opc_repository",
        level=NotificationLevel.ERROR,
        attachment_content=ANY
    )


@mark.asyncio
async def test_write_data_exception(opc):
    opc.opc_repository['server1'].write_data.side_effect = Exception(
        "Test error")

    try:
        await opc.write_data(server_id='server1', tag='tag1', data=50,
                             data_type='int', tag_type='prediction', metadata=metadata)

    except Exception:
        opc.send_notification.assert_called_once_with(
            metadata=metadata,
            notification_id="WRITE_OPC_PREDICTION_ERROR",
            message="Error writing data to OPC server: Test error",
            block="write_opc_data",
            level=NotificationLevel.ERROR,
            attachment_content=ANY
        )

    else:
        assert False, "Expected an exception to be raised"


@mark.asyncio
async def test_write_opc_data_success(opc):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': [0.75],
            'prediction_confidence': [0.95]
        },
        'opc_output_config': {
            'server1': {
                'prediction_tags': {
                    'tag1': {'data_type': 'float'}
                },
                'confidence_tags': {
                    'tag2': {'data_type': 'float'}
                }
            }
        }
    }

    # Act
    opc.write_data = AsyncMock(return_value=True)
    opc.process_confidence = MagicMock(return_value={'data': 'data'})
    output = await opc.write_opc_data(input_data)

    # Assert
    assert output == {'data': 'data'}
    opc.write_data.assert_has_calls([
        call(
            server_id='server1',
            tag='tag1',
            data=0.75,
            data_type='float',
            tag_type='prediction',
            metadata=metadata['metadata']
        )])
    opc.write_data.assert_has_calls([
        call(
            server_id='server1',
            tag='tag2',
            data=0.95,
            data_type='float',
            tag_type='confidence',
            metadata=metadata['metadata']
        )
    ])
    assert opc.write_data.call_count == 2


@mark.asyncio
async def test_write_opc_data_empty_config(opc):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': [0.75],
            'prediction_confidence': [0.95]
        },
        'opc_servers': ['server1'],
        'opc_output_config': {
            'server1': {
                'prediction_tags': {},
                'confidence_tags': {}
            }
        }
    }

    # Act
    await opc.write_opc_data(input_data)

    # Assert
    opc.opc_repository['server1'].write_data.assert_not_called()


@mark.asyncio
async def test_write_opc_data_no_validate_server(opc):
    opc.validate_server = MagicMock(return_value=False)
    input_data = {
        **metadata,
        'data': {
            'prediction': [0.75],
            'prediction_confidence': [0.95]
        },
        'opc_output_config': {
            'server1': {
                'prediction_tags': {
                    'tag1': {'data_type': 'float'}
                },
                'confidence_tags': {
                    'tag2': {'data_type': 'float'}
                }
            }
        }
    }

    # Act
    await opc.write_opc_data(input_data)

    # Assert
    opc.opc_repository['server1'].write_data.assert_not_called()


@mark.parametrize('data,success,expected', [
    (DataFrame({'prediction_confidence': [0]}), True, 0),
    (DataFrame({'prediction_confidence': [0]}), False, 12),
])
def test_process_confidence(opc, data, success, expected):
    # Act
    result = opc.process_confidence(data, success, metadata)

    # Assert
    assert result['prediction_confidence'][0] == expected


def test_validate_server(opc):
    assert opc.validate_server('server1', metadata) is True
    assert opc.validate_server('server2', metadata) is False


@mark.asyncio
async def test_shutdown(opc):
    opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
    await opc.shutdown()
    opc.opc_repository['server1'].disconnect.assert_called_once()