SIENTIAPDE-1102

Refactor OPC activity to enhance data writing and confidence processing; update tests accordingly

tests/laborious/activities/test_opc.py

@@ -1,4 +1,5 @@
 from unittest.mock import patch, MagicMock, ANY, call
+from pandas import DataFrame
 from pytest import fixture, mark
 from laborious.activities.opc import NotificationLevel
 
@@ -75,7 +76,7 @@ def test___init__(mock_opc_repository):
 
 @fixture
 @patch("laborious.activities.opc.OpcRepository")
-def opc(_mock_opc_repository):
+def opc(mock_opc_repository):
     servers = {
         'server1': {
             'url': 'http://localhost:8080',
@@ -86,6 +87,9 @@ def opc(_mock_opc_repository):
             'reconnection_interval': 60,
         }
     }
+    mock_opc_repository.write_data = MagicMock(
+        return_value=True
+    )
     return OPC(
         opc_servers=servers,
         logger=MagicMock(),
@@ -103,8 +107,8 @@ WRITE_DATA_CASES = [
 
 @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
 def test_write_data_success(opc, tag, data_type, data):
-    opc.write_data(server='server1', tag=tag, data=data,
-                   data_type=data_type, tag_type='prediction')
+    assert opc.write_data(server='server1', tag=tag, data=data,
+                          data_type=data_type, tag_type='prediction')
     opc.opc_repository['server1'].write_data.assert_called_once_with(
         tag, data, data_type)
 
@@ -152,9 +156,11 @@ async def test_write_opc_data_success(opc):
 
     # Act
     opc.write_data = MagicMock()
-    await opc.write_opc_data(input_data)
+    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='server1',
@@ -197,6 +203,18 @@ async def test_write_opc_data_empty_config(opc):
     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)
+
+    # Assert
+    assert result['prediction_confidence'][0] == expected
+
+
 def test_shutdown(opc):
     opc.shutdown()
     opc.opc_repository['server1'].disconnect.assert_called_once()

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -40,17 +40,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
             retry_policy=ANY,
             start_to_close_timeout=ANY
         )])
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_local_activity_method.return_value
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )])
+
     workflow_mock.execute_activity_method.assert_has_calls([
         call(
             Activities.write_opc_data,
@@ -63,6 +53,18 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
         )
     ])
 
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.export_data_to_postgres,
+            {
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name'],
+                'data': workflow_mock.execute_activity_method.return_value
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )])
+
     assert workflow_mock.execute_activity_method.call_count == 2
     assert workflow_mock.execute_local_activity_method.call_count == 1
 
@@ -99,18 +101,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
             start_to_close_timeout=ANY
         )
     ])
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_local_activity_method.return_value
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+
     workflow_mock.execute_activity_method.assert_has_calls([
         call(
             Activities.write_opc_data,
@@ -123,5 +114,18 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
         )
     ])
 
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.export_data_to_postgres,
+            {
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name'],
+                'data': workflow_mock.execute_activity_method.return_value
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])
+
     assert workflow_mock.execute_activity_method.call_count == 2
     assert workflow_mock.execute_local_activity_method.call_count == 1