Merge pull request #6 from Aignosi/SIENTIAPDE-1102-ajustar-escrita-postgres-do-laborious-para-ocorrer-apos-a-escrita-opc

SIENTIAPDE-1102

laborious/activities/opc.py

@@ -11,6 +11,8 @@ with workflow.unsafe.imports_passed_through():
     import traceback
     from pandas import DataFrame
 
+OPC_WRITTING_ERROR_CONFIDENCE = 12
+
 
 class OPC(BaseActivity):
     def __init__(self, opc_servers: dict[str, dict[str, Any]],
@@ -38,7 +40,7 @@ class OPC(BaseActivity):
         BaseActivity.__init__(self, logger, notification_handler)
 
     def write_data(self, server: str, tag: str, data: Any,
-                   data_type: str, tag_type: str):
+                   data_type: str, tag_type: str) -> bool:
         """
         Write data to OPC server.
 
@@ -48,10 +50,13 @@ class OPC(BaseActivity):
             - data (Any): The data to write.
             - data_type (str): The data type.
             - tag_type (str): The tag type.
+
+        Returns:
+            - bool: True if the data was written successfully, False otherwise.
         """
 
         try:
-            self.opc_repository[server].write_data(
+            return self.opc_repository[server].write_data(
                 tag, data, data_type)
             self.logger.debug(f"Wrote {tag_type} to {tag}")
         except Exception as e:
@@ -66,7 +71,7 @@ class OPC(BaseActivity):
             raise e
 
     @activity.defn(name='write_opc_data')
-    async def write_opc_data(self, input_data: dict[str, Any]):
+    async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Write prediction and confidence data to OPC servers. The two writing
         operations are optional and independent of each other.
@@ -80,12 +85,17 @@ class OPC(BaseActivity):
                     - prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
                     - confidence_tags(dict[str, Any]): The tags to write to the OPC servers.
 
+        Returns:
+            - dict[Any, Any]: The data that was written to the OPC servers.
+
         """
         self.logger.debug("Writing data to OPC servers...")
         data = DataFrame(input_data['data'])
         opc_output_config = input_data['opc_output_config']
         self.logger.debug(data)
 
+        success = True
+
         for server, config in opc_output_config.items():
             if self.opc_repository.get(server) is None:
                 self.logger.error(f"OPC server {server} not found")
@@ -93,13 +103,14 @@ class OPC(BaseActivity):
 
             if 'prediction_tags' in config:
                 for tag, tag_config in config['prediction_tags'].items():
-                    self.write_data(
+                    success = success and self.write_data(
                         server=server,
                         tag=tag,
                         data=data.head(1)['prediction'].values[0],
                         data_type=tag_config['data_type'],
                         tag_type='prediction'
                     )
+
             if 'confidence_tags' in config:
                 for tag, tag_config in config['confidence_tags'].items():
                     self.write_data(
@@ -110,6 +121,34 @@ class OPC(BaseActivity):
                         tag_type='confidence'
                     )
 
+        return self.process_confidence(data, success)
+
+    def process_confidence(self, data: DataFrame, success: bool) -> dict[Any, Any]:
+        """
+        Processes the confidence of OPC server write operations and updates the DataFrame accordingly.
+
+        If the write operation was not successful, sets the 'prediction_confidence' column in the DataFrame
+        to a predefined error confidence value and logs a debug message. Otherwise, logs a success message.
+
+        Args:
+            data (DataFrame): The DataFrame containing the data to be processed.
+            success (bool): Indicates whether the data was successfully written to the OPC servers.
+
+        Returns:
+            dict[Any, Any]: The processed data as a dictionary.
+        """
+        if not success:
+            data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
+            self.logger.debug(
+                "Some data could not be written to OPC servers, setting confidence to "
+                f"{OPC_WRITTING_ERROR_CONFIDENCE}."
+            )
+
+        else:
+            self.logger.info("Data written to OPC servers successfully.")
+
+        return data.to_dict()
+
     def shutdown(self):
         for opc in self.opc_repository.values():
             opc.disconnect()

laborious/utils/repository/opc_repository.py

@@ -2,9 +2,11 @@ import traceback
 from logging import Logger
 from datetime import datetime
 from pathlib import Path
+from typing import Any
 from asyncua.sync import Client
 from asyncua.crypto.security_policies import SecurityPolicyBasic256
 from asyncua.ua import DataValue, Variant, VariantType
+from regex import F
 from sientia_do.notifications.handlers import NotificationHandler
 from sientia_do.notifications.models import NotificationLevel
 
@@ -190,7 +192,7 @@ class OpcRepository():
 
         return True
 
-    def write_data(self, node, value, data_type):
+    def write_data(self, node: str, value: Any, data_type: str) -> bool:
         """
         Writes data to the OPC server.
         If the connection is not established, it attempts to reconnect.
@@ -202,7 +204,7 @@ class OpcRepository():
         If the client is connected, it returns True.
         """
         if not self.validate_connection():
-            return
+            return False
         try:
             node = self.client.get_node(node)
         except Exception as e:
@@ -216,7 +218,7 @@ class OpcRepository():
             )
             self.logger.error(trace)
             self.error_count += 1
-            return
+            return False
 
         if data_type not in data_type_map:
             self.notification_handler.build_and_send_notification(
@@ -225,7 +227,7 @@ class OpcRepository():
                 block="opc_repository",
                 level=NotificationLevel.ERROR
             )
-            return
+            return False
 
         data = data_type_map[data_type]['converter'](value)
         self.logger.info(f'Writing {data} - {type(data)} to {node}')
@@ -245,5 +247,7 @@ class OpcRepository():
             )
             self.logger.error(trace)
             self.error_count += 1
-            return
+            return False
         self.error_count = 0
+
+        return True

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -68,20 +68,8 @@ class FormatAndExportPrediction():
                 start_to_close_timeout=timedelta(seconds=60)
             )
 
-        # write to postgres
-        postgres_holder = workflow.execute_activity_method(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': prediction
-            },
-            retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
-        )
-
         # write to opc
-        opc_holder = workflow.execute_activity_method(
+        prediction = await workflow.execute_activity_method(
             Activities.write_opc_data,
             {
                 'opc_output_config': input_data['opc_output_config'],
@@ -91,5 +79,14 @@ class FormatAndExportPrediction():
             start_to_close_timeout=timedelta(seconds=60)
         )
 
-        await postgres_holder
+        # write to postgres
-        await opc_holder
+        await workflow.execute_activity_method(
+            Activities.export_data_to_postgres,
+            {
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name'],
+                'data': prediction,
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )

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,
+    assert opc.write_data(server='server1', tag=tag, data=data,
-                   data_type=data_type, tag_type='prediction')
+                          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