SIENTIAPDE-1314

Enhance OPC Metrics Handling and Refactor Write Operations - Updated the OPC class to return response times for write operations, improving metrics tracking. - Refactored the Gates activity to incorporate OPC metrics into the metrics writing process. - Adjusted the manage_output_tags method in OpcRepository to return response times for each tag written. - Modified tests to validate the new metrics structure and ensure correct behavior of the updated methods.
2025-10-23 17:48:43 -03:00
parent 32a76b35ea
commit 0324e2e143
10 changed files with 177 additions and 99 deletions
--- a/tests/laborious/activities/test_gates.py
+++ b/tests/laborious/activities/test_gates.py
@@ -1,4 +1,4 @@
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import ANY, MagicMock, call, patch

 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
@@ -636,6 +636,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
            'prediction_confidence': [0.9, 0.8, 0.7],
            'response_time': [0.1, 0.2, 0.3],
        },
+        'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
    }
    await gates_activity.write_metrics(input_data)
    mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
@@ -660,3 +661,35 @@ async def test_write_metrics(mock_metrics, gates_activity):
    mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
        0.1
    )
+
+    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_has_calls(
+        [
+            call(
+                pod_id=gates_activity.pod_id,
+                model_name=metadata['metadata']['model_name'],
+                pipeline_name=metadata['metadata']['workflow_name'],
+                opc_server_id='server1',
+                tag='tag1',
+            )
+        ]
+    )
+    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_has_calls(
+        [
+            call(
+                pod_id=gates_activity.pod_id,
+                model_name=metadata['metadata']['model_name'],
+                pipeline_name=metadata['metadata']['workflow_name'],
+                opc_server_id='server1',
+                tag='tag1',
+            )
+        ]
+    )
+
+    assert mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.call_count == 2
+    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_has_calls(
+        [
+            call(0.1),
+            call(0.2),
+        ],
+        any_order=True,
+    )
--- a/tests/laborious/activities/test_opc.py
+++ b/tests/laborious/activities/test_opc.py
@@ -180,6 +180,8 @@ WRITE_DATA_CASES = [
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
 async def test_write_data_success(opc, tag, data_type, data):
+    opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
+
    result = await opc.write_data(
        server_id='server1',
        tag=tag,
@@ -188,7 +190,7 @@ async def test_write_data_success(opc, tag, data_type, data):
        tag_type='prediction',
        metadata=metadata,
    )
-    assert result is True
+    assert result == 0.1
    opc.opc_repository['server1'].write_data.assert_called_once_with(
        tag, data, data_type, opc.logger, metadata
    )
@@ -215,7 +217,7 @@ async def test_write_data_failed(opc):
        tag_type='prediction',
        metadata=metadata,
    )
-    assert result is False
+    assert result is None

    opc.send_notification.assert_called_once_with(
        metadata=metadata,
@@ -256,7 +258,8 @@ async def test_write_data_exception(opc):


@mark.asyncio
-async def test_write_opc_data_success(opc):
+@patch('laborious.activities.opc.DataFrame')
+async def test_write_opc_data_success(mock_dataframe, opc):
    # Arrange
    input_data = {
        **metadata,
@@ -270,37 +273,28 @@ async def test_write_opc_data_success(opc):
    }

    # Act
-    opc.write_data = AsyncMock(return_value=True)
+    opc.manage_output_tags = AsyncMock(return_value=(True, {'tag1': 0.1, 'tag2': 0.2}))
+
    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'],
-            )
-        ]
+    assert output == {
+        'data': {'data': 'data'},
+        'metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
+    }
+    opc.manage_output_tags.assert_called_once_with(
+        'server1',
+        input_data['opc_output_config']['server1'],
+        mock_dataframe.return_value,
+        metadata['metadata'],
+        True,
    )
-    opc.write_data.assert_has_calls(
-        [
-            call(
-                server_id='server1',
-                tag='tag2',
-                data=0.95,
-                data_type='float',
-                tag_type='confidence',
-                metadata=metadata['metadata'],
-            )
-        ]
+    opc.process_confidence.assert_called_once_with(
+        mock_dataframe.return_value,
+        True,
+        metadata['metadata'],
    )
-    assert opc.write_data.call_count == 2


@mark.asyncio
--- a/tests/laborious/utils/repository/test_opc_repository.py
+++ b/tests/laborious/utils/repository/test_opc_repository.py
@@ -335,7 +335,7 @@ async def test_write_data_validate_connection_do_nothing(opc_repository):

    opc_repository.validate_connection.assert_called_once()
    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
-    assert result == (True, {})
+    assert result == (True, {'response_time': ANY})


@pytest.mark.asyncio
@@ -403,8 +403,7 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client):


@pytest.mark.asyncio
-@patch('laborious.utils.repository.opc_repository.metrics')
-async def test_write_data(mock_metrics, opc_repository, mock_client):
+async def test_write_data(opc_repository, mock_client):
    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
    opc_repository.client = mock_client
    mock_node = AsyncMock()
@@ -416,25 +415,7 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):

    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,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id,
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
-
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
-        pod_id=opc_repository.pod_id,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id,
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
-        ANY
-    )
+    assert result == (True, {'response_time': ANY})


@pytest.mark.asyncio
--- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
+++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, call, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch

 from pytest import fixture, mark
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@@ -42,6 +42,15 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
        'prediction_store_policy': 'erl:1',
    }

+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
    await format_and_export_prediction.run(input_data)

    workflow_mock.execute_local_activity_method.assert_has_calls(
@@ -84,7 +93,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                {
                    'schema': input_data['schema'],
                    'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
                    **metadata,
                    'timestamp_conversion': {
                        'column': 'timestamp',
@@ -97,6 +106,21 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
        ]
    )

+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
    assert workflow_mock.execute_activity_method.call_count == 3
    assert workflow_mock.execute_local_activity_method.call_count == 1

@@ -121,6 +145,15 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
        'comment': 'test_comment',
    }

+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
    await format_and_export_prediction.run(input_data)

    workflow_mock.execute_local_activity_method.assert_has_calls(
@@ -162,7 +195,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                {
                    'schema': input_data['schema'],
                    'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': prediction_data,
                    **metadata,
                    'timestamp_conversion': {
                        'column': 'timestamp',
@@ -175,5 +208,20 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
        ]
    )

+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
    assert workflow_mock.execute_activity_method.call_count == 3
    assert workflow_mock.execute_local_activity_method.call_count == 1