From f937eca1518ce43e2e9594cad73b17299606cbde Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 28 Aug 2025 10:21:20 -0300 Subject: [PATCH] SIENTIAPDE-1205 Refactor OpcRepository and update tests for async handling - Removed the shutdown method from OpcRepository and adjusted the disconnect logic. - Updated tests in test_activities.py and test_opc.py to support async shutdown functionality. - Enhanced test cases in test_opc_repository.py to ensure proper async behavior and error handling in OpcRepository methods. --- laborious/utils/repository/opc_repository.py | 10 - run_coverage.sh | 11 + tests/laborious/activities/test_activities.py | 7 +- tests/laborious/activities/test_opc.py | 115 ++++++-- .../utils/repository/test_opc_repository.py | 248 +++++++++++------- 5 files changed, 260 insertions(+), 131 deletions(-) create mode 100755 run_coverage.sh diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 1315c44..3439a19 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -158,16 +158,6 @@ class OpcRepository(): f"Failed to disconnect from OPC server: {e}", self.metadata) self.client = None - def shutdown(self): - """ - Disconnects from the OPC server when the object is destroyed. - """ - try: - asyncio.run(self.disconnect()) - except Exception as e: - self.logger.custom_error( - f"Error in destructor: {e}", self.metadata) - async def validate_connection(self) -> tuple[bool, dict[str, Any]]: """ Validates the connection to the OPC server using protocol state checking. diff --git a/run_coverage.sh b/run_coverage.sh new file mode 100755 index 0000000..f9af4cb --- /dev/null +++ b/run_coverage.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Exit on any error +set -e + +echo "Activating virtual environment..." +source ./venv/bin/activate + +pytest --cov=laborious --cov-report=html + +xdg-open htmlcov/index.html \ No newline at end of file diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index e482ca8..de418be 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -90,11 +90,12 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre ) +@mark.asyncio @patch('laborious.activities.activities.Postgres', return_value=MagicMock()) @patch('laborious.activities.activities.MLFlow', return_value=MagicMock()) @patch('laborious.activities.activities.OPC', return_value=MagicMock()) -def test_shutdown(mock_opc_init, - _mock_mlflow_init, mock_postgres_init): +async def test_shutdown(mock_opc_init, + _mock_mlflow_init, mock_postgres_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -129,6 +130,6 @@ def test_shutdown(mock_opc_init, notification_handler=notification_handler ) - activities.shutdown() + await activities.shutdown() mock_opc_init.shutdown.assert_called_once() mock_postgres_init.close.assert_called_once() diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index 37ae391..07f00b8 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -1,7 +1,8 @@ -from unittest.mock import patch, MagicMock, ANY, call +from unittest.mock import patch, MagicMock, ANY, call, AsyncMock from pandas import DataFrame from pytest import fixture, mark -from laborious.activities.opc import NotificationLevel +import pytest_asyncio +from sientia_do.notifications.models import NotificationLevel from laborious.activities.opc import OPC @@ -15,27 +16,42 @@ metadata = { } +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("laborious.activities.opc.OpcRepository") @patch("laborious.activities.opc.OPC.send_notification") -def test___init__(mock_send_notification, mock_opc_repository): +async def test_init_opc(mock_send_notification, mock_opc_repository): mock_logger = MagicMock() server1 = MagicMock( - connect=MagicMock(return_value=(True, {})), - write_data=MagicMock(return_value=(True, {})) + connect=AsyncMock(return_value=(True, {})), + write_data=AsyncMock(return_value=(True, {})) ) server2 = MagicMock( - connect=MagicMock(return_value=(True, {})), - write_data=MagicMock(return_value=(True, {})) + connect=AsyncMock(return_value=(True, {})), + write_data=AsyncMock(return_value=(True, {})) ) server3 = MagicMock( - connect=MagicMock(return_value=(False, { + 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=MagicMock(return_value=(True, {})) + write_data=AsyncMock(return_value=(True, {})) ) mock_opc_repository.side_effect = [server1, server2, server3] mock_notification_handler = MagicMock() @@ -73,6 +89,7 @@ def test___init__(mock_send_notification, mock_opc_repository): logger=mock_logger, notification_handler=mock_notification_handler ) + await opc.init_opc() assert opc.opc_servers == servers assert opc.logger == mock_logger @@ -129,9 +146,9 @@ def test___init__(mock_send_notification, mock_opc_repository): ]) -@fixture +@pytest_asyncio.fixture @patch("laborious.activities.opc.OpcRepository") -def opc(mock_opc_repository): +async def opc(mock_opc_repository): servers = { 'server1': { 'id': 'server1', @@ -144,10 +161,10 @@ def opc(mock_opc_repository): } } - mock_opc_repository.return_value.write_data = MagicMock( + mock_opc_repository.return_value.write_data = AsyncMock( return_value=(True, {}) ) - mock_opc_repository.return_value.connect = MagicMock( + mock_opc_repository.return_value.connect = AsyncMock( return_value=(True, {}) ) opc = OPC( @@ -155,7 +172,7 @@ def opc(mock_opc_repository): logger=MagicMock(), notification_handler=MagicMock() ) - + await opc.init_opc() opc.send_notification = MagicMock() return opc @@ -169,14 +186,17 @@ WRITE_DATA_CASES = [ @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES) -def test_write_data_success(opc, tag, data_type, data): - assert opc.write_data(server_id='server1', tag=tag, data=data, - data_type=data_type, tag_type='prediction', metadata=metadata) +@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) -def test_write_data_failed(opc): +@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', @@ -185,8 +205,9 @@ def test_write_data_failed(opc): 'attachment_content': 'Test error' }) - assert opc.write_data(server_id='server1', tag='tag1', data=50, - data_type='int', tag_type='prediction', metadata=metadata) is False + 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, @@ -198,13 +219,14 @@ def test_write_data_failed(opc): ) -def test_write_data_exception(opc): +@mark.asyncio +async def test_write_data_exception(opc): opc.opc_repository['server1'].write_data.side_effect = Exception( "Test error") try: - opc.write_data(server_id='server1', tag='tag1', data=50, - data_type='int', tag_type='prediction', metadata=metadata) + 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( @@ -242,7 +264,7 @@ async def test_write_opc_data_success(opc): } # Act - opc.write_data = MagicMock() + opc.write_data = AsyncMock(return_value=True) opc.process_confidence = MagicMock(return_value={'data': 'data'}) output = await opc.write_opc_data(input_data) @@ -281,8 +303,38 @@ async def test_write_opc_data_empty_config(opc): }, 'opc_servers': ['server1'], 'opc_output_config': { - 'prediction_tags': {}, - 'confidence_tags': {} + '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'} + } + } } } @@ -305,6 +357,13 @@ def test_process_confidence(opc, data, success, expected): assert result['prediction_confidence'][0] == expected -def test_shutdown(opc): - opc.shutdown() +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() diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index cf9b493..bc84db8 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -1,17 +1,17 @@ -from unittest.mock import Mock, patch, MagicMock, ANY, call +import pytest +from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call from asyncua.crypto.security_policies import SecurityPolicyBasic256 -from pytest import fixture from laborious.utils.repository.opc_repository import OpcRepository from sientia_do.notifications.models import NotificationLevel from datetime import datetime -@fixture +@pytest.fixture def mock_logger(): return Mock() -@fixture +@pytest.fixture def opc_repository(mock_logger): return OpcRepository( id="test_repo", @@ -26,10 +26,10 @@ def opc_repository(mock_logger): ) -@fixture +@pytest.fixture def mock_client(): with patch('laborious.utils.repository.opc_repository.Client') as mock: - client_instance = MagicMock() + client_instance = AsyncMock() mock.return_value = client_instance yield client_instance @@ -57,9 +57,10 @@ def test_init(opc_repository): assert opc_repository.error_count == 0 -def test_set_security(opc_repository, mock_client): +@pytest.mark.asyncio +async def test_set_security(opc_repository, mock_client): opc_repository.client = mock_client - opc_repository.set_security() + await opc_repository.set_security() mock_client.application_uri = "urn:test:server" mock_client.set_security.assert_called_once_with( @@ -72,50 +73,59 @@ def test_set_security(opc_repository, mock_client): assert mock_client.session_timeout == 10000000 -def test_set_security_missing_certificates(opc_repository): +@pytest.mark.asyncio +async def test_set_security_missing_certificates(opc_repository): opc_repository.cert_path = None opc_repository.private_key_path = None try: - opc_repository.set_security() + await opc_repository.set_security() except ValueError as e: assert str( e) == "Certificate and private key paths must be provided for secure connection." -def test_connect_with_security(opc_repository, mock_client): - opc_repository.try_connect = MagicMock() - opc_repository.connect() +@pytest.mark.asyncio +async def test_connect_with_security(opc_repository, mock_client): + opc_repository.try_connect = AsyncMock(return_value=(True, {})) + result = await opc_repository.connect() opc_repository.try_connect.assert_called_once() assert opc_repository.client == mock_client + assert result == (True, {}) -def test_connect_without_security(opc_repository, mock_client): +@pytest.mark.asyncio +async def test_connect_without_security(opc_repository, mock_client): opc_repository.cert_path = None - opc_repository.try_connect = MagicMock() - opc_repository.set_security = MagicMock() - opc_repository.connect() + opc_repository.try_connect = AsyncMock(return_value=(True, {})) + opc_repository.set_security = AsyncMock() + result = await opc_repository.connect() opc_repository.try_connect.assert_called_once() opc_repository.set_security.assert_not_called() assert opc_repository.client == mock_client + assert result == (True, {}) -def test_try_connect_sucess(opc_repository): +@pytest.mark.asyncio +async def test_try_connect_success(opc_repository): opc_repository.last_reconnection_time = None - opc_repository.client = MagicMock() - opc_repository.try_connect() + opc_repository.client = AsyncMock() + result = await opc_repository.try_connect() + opc_repository.client.connect.assert_called_once() assert opc_repository.last_reconnection_time is not None + assert result == (True, {}) -def test_try_connect_fail(opc_repository): +@pytest.mark.asyncio +async def test_try_connect_fail(opc_repository): opc_repository.last_reconnection_time = None opc_repository.client = MagicMock() opc_repository.client.connect.side_effect = Exception("Test error") - is_connected, error_data = opc_repository.try_connect() + is_connected, error_data = await opc_repository.try_connect() opc_repository.client.connect.assert_called_once() assert is_connected is False @@ -126,18 +136,26 @@ def test_try_connect_fail(opc_repository): assert error_data['attachment_content'] is not None -def test_disconnect(opc_repository, mock_client): +@pytest.mark.asyncio +async def test_disconnect(opc_repository, mock_client): opc_repository.client = mock_client - opc_repository.disconnect() + await opc_repository.disconnect() mock_client.disconnect.assert_called_once() assert opc_repository.client is None -def test_disconnect_error(opc_repository, mock_client): +@pytest.mark.asyncio +async def test_disconnect_no_client(opc_repository): + opc_repository.client = None + assert await opc_repository.disconnect() is None + + +@pytest.mark.asyncio +async def test_disconnect_error(opc_repository, mock_client): opc_repository.client = mock_client mock_client.disconnect.side_effect = Exception("Test error") - opc_repository.disconnect() + await opc_repository.disconnect() opc_repository.logger.custom_error.assert_called_once_with( "Failed to disconnect from OPC server: Test error", @@ -146,21 +164,25 @@ def test_disconnect_error(opc_repository, mock_client): assert opc_repository.client is None -def test_validate_connection_none_client(opc_repository): +@pytest.mark.asyncio +async def test_validate_connection_none_client(opc_repository): opc_repository.client = None - opc_repository.connect = MagicMock() - response = opc_repository.validate_connection() - assert response + opc_repository.connect = AsyncMock(return_value=(True, {})) + response = await opc_repository.validate_connection() + assert response == (True, {}) opc_repository.connect.assert_called_once() -def test_validate_connection_error_count_disconnect_error(opc_repository): +@pytest.mark.asyncio +async def test_validate_connection_error_count_disconnect_error(opc_repository): opc_repository.error_count = 6 - opc_repository.client = MagicMock() - opc_repository.disconnect = MagicMock(side_effect=Exception("Test error")) - opc_repository.connect = MagicMock() + opc_repository.client = AsyncMock() + opc_repository.disconnect = AsyncMock( + side_effect=Exception("Test error") + ) + opc_repository.connect = AsyncMock(return_value=(True, {})) - response = opc_repository.validate_connection() + response = await opc_repository.validate_connection() assert response == opc_repository.connect.return_value opc_repository.disconnect.assert_called_once() opc_repository.connect.assert_called_once() @@ -171,18 +193,37 @@ def test_validate_connection_error_count_disconnect_error(opc_repository): ) -@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True) -@patch('laborious.utils.repository.opc_repository.datetime', - MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0)))) -def test_validate_connection_lost_not_time_to_reconect(_mock_datetime, opc_repository): +@pytest.mark.asyncio +async def test_validate_connection_error_validate_connection_error(opc_repository): + opc_repository.client = MagicMock( + uaclient=Exception("Test error") + ) + opc_repository.error_count = 0 + + response = await opc_repository.validate_connection() + + assert response == (False, { + "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}", + "message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'", + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": ANY + }) + + +@pytest.mark.asyncio +@patch('laborious.utils.repository.opc_repository.datetime') +async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository): + _mock_datetime.now = MagicMock( + return_value=datetime(2025, 1, 1, 0, 0, 0)) opc_repository.error_count = 0 opc_repository.client = MagicMock() - opc_repository.client.aio_obj.uaclient.protocol = None + opc_repository.client.uaclient.protocol = None opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) - opc_repository.try_connect = MagicMock() + opc_repository.connect = MagicMock(return_value=(True, {})) - response = opc_repository.validate_connection() - opc_repository.try_connect.assert_not_called() + response = await opc_repository.validate_connection() + opc_repository.connect.assert_not_called() assert response == (False, { "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}", "message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...", @@ -191,96 +232,120 @@ def test_validate_connection_lost_not_time_to_reconect(_mock_datetime, opc_repos }) -@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True) -@patch('laborious.utils.repository.opc_repository.datetime', - MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0)))) -def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repository): +@pytest.mark.asyncio +@patch('laborious.utils.repository.opc_repository.datetime') +async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository): + mock_datetime.now = MagicMock( + return_value=datetime(2025, 1, 1, 1, 0, 0)) opc_repository.error_count = 0 - opc_repository.client = MagicMock() - opc_repository.client.aio_obj.uaclient.protocol = None + opc_repository.client = AsyncMock() + opc_repository.client.uaclient.protocol = None opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) - opc_repository.connect = MagicMock() + opc_repository.connect = AsyncMock(return_value=(True, {})) - response = opc_repository.validate_connection() + response = await opc_repository.validate_connection() opc_repository.connect.assert_called_once() assert response == opc_repository.connect.return_value -def test_validate_connection_failed(opc_repository): +@pytest.mark.asyncio +async def test_validate_connection_success(opc_repository): opc_repository.client = MagicMock() opc_repository.error_count = 0 + opc_repository.client.uaclient.protocol = MagicMock() + opc_repository.client.uaclient.protocol.state = "open" - output = opc_repository.validate_connection() + output = await opc_repository.validate_connection() assert output == (True, {}) -def test_write_data_validate_connection_do_nothing(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=(True, {})) - opc_repository.client = MagicMock() - opc_repository.write_data("ns=2;s=TestNode", 42.0, - "float", opc_repository.logger, metadata) +@pytest.mark.asyncio +async def test_write_data_validate_connection_do_nothing(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = AsyncMock( + get_node=MagicMock() + ) + mock_node = AsyncMock() + opc_repository.client.get_node.return_value = mock_node + + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") + assert result == (True, {}) -def test_write_data_validate_connection_failed(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=(False, {})) - opc_repository.client = MagicMock() +@pytest.mark.asyncio +async def test_write_data_validate_connection_failed(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(False, {})) + opc_repository.client = AsyncMock() opc_repository.error_count = 0 - opc_repository.write_data("ns=2;s=TestNode", 42.0, - "float", opc_repository.logger, metadata) + + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_not_called() + assert result == (False, {}) -def test_write_data_get_node_failed(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=(True, {})) - opc_repository.client = MagicMock() +@pytest.mark.asyncio +async def test_write_data_get_node_failed(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = AsyncMock() opc_repository.error_count = 0 - opc_repository.client.get_node.side_effect = Exception("Test error") - is_success, error_data = opc_repository.write_data("ns=2;s=TestNode", 42.0, - "float", opc_repository.logger, metadata) + opc_repository.client.get_node = MagicMock( + side_effect=Exception("Test error")) + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") assert is_success is False assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}" - assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}" + assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" assert error_data['block'] == "opc_repository" assert error_data['level'] == NotificationLevel.ERROR assert error_data['attachment_content'] is not None -def test_write_data_invalid_data_type(opc_repository, mock_client): - opc_repository.validate_connection = MagicMock(return_value=(True, {})) +@pytest.mark.asyncio +async def test_write_data_invalid_data_type(opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) opc_repository.client = mock_client - mock_node = MagicMock() - mock_client.get_node.return_value = mock_node + mock_node = AsyncMock() + mock_client.get_node = MagicMock(return_value=mock_node) + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "invalid_type", opc_repository.logger, metadata['metadata']) - is_success, error_data = opc_repository.write_data("ns=2;s=TestNode", 42.0, - "invalid_type", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") assert is_success is False assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}" - assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}" + assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" assert error_data['block'] == "opc_repository" assert error_data['level'] == NotificationLevel.ERROR assert error_data.get('attachment_content') is None +@pytest.mark.asyncio @patch('laborious.utils.repository.opc_repository.metrics') -def test_write_data(mock_metrics, opc_repository, mock_client): - opc_repository.validate_connection = MagicMock(return_value=(True, {})) +async def test_write_data(mock_metrics, opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) opc_repository.client = mock_client - mock_node = MagicMock() - mock_client.get_node.return_value = mock_node + mock_node = AsyncMock() + mock_client.get_node = MagicMock(return_value=mock_node) - opc_repository.write_data("ns=2;s=TestNode", 42.0, - "float", opc_repository.logger, metadata['metadata']) + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") mock_node.write_value.assert_called_once() + assert result == (True, {}) mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with( pod_id=opc_repository.pod_id, @@ -300,21 +365,24 @@ def test_write_data(mock_metrics, opc_repository, mock_client): ANY) -def test_write_data_write_value_failed(opc_repository, mock_client): - opc_repository.validate_connection = MagicMock(return_value=(True, {})) +@pytest.mark.asyncio +async def test_write_data_write_value_failed(opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) opc_repository.client = mock_client - mock_node = MagicMock() + mock_node = AsyncMock() opc_repository.error_count = 0 - mock_client.get_node.return_value = mock_node + mock_client.get_node = MagicMock(return_value=mock_node) mock_node.write_value.side_effect = Exception("Test error") - is_success, error_data = opc_repository.write_data("ns=2;s=TestNode", 42.0, - "float", opc_repository.logger, metadata) + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + opc_repository.validate_connection.assert_called_once() mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") mock_node.write_value.assert_called_once() assert is_success is False assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}" - assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}" + assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" assert error_data['block'] == "opc_repository" assert error_data['level'] == NotificationLevel.ERROR assert error_data['attachment_content'] is not None