SIENTIAPDE-1478

Refactor Activities and API Integration for PI Web API - Reintroduced the API import in the Activities class for proper integration. - Cleaned up whitespace and formatting in the API class and related tests for improved readability. - Updated test cases to ensure consistent formatting in error messages and configuration structures for PI Web API. - Enhanced connectors_config.py with additional whitespace for better organization.
2026-01-09 09:00:21 -03:00
parent 892823df11
commit 1bddde17f4
8 changed files with 63 additions and 38 deletions
--- a/tests/laborious/activities/test_api.py
+++ b/tests/laborious/activities/test_api.py
@@ -1,7 +1,6 @@
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch

 import pytest_asyncio
-from pandas import DataFrame
 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel

@@ -21,7 +20,7 @@ def _create_mock_dataframe(to_dict_return=None):
    """Helper function to create a mocked DataFrame for testing."""
    mock_df = MagicMock()
    mock_head = MagicMock()
-    
+
    def get_column_values(key):
        if key == 'prediction':
            return MagicMock(values=[0.75])
@@ -29,10 +28,10 @@ def _create_mock_dataframe(to_dict_return=None):
            return MagicMock(values=[0.95])
        else:
            return MagicMock(values=['2024-01-01T00:00:00+00:00'])
-    
+
    mock_head.__getitem__.side_effect = get_column_values
    mock_df.head.return_value = mock_head
-    
+
    if to_dict_return is None:
        to_dict_return = {
            'prediction': [0.75],
@@ -40,7 +39,7 @@ def _create_mock_dataframe(to_dict_return=None):
            'timestamp': ['2024-01-01T00:00:00+00:00'],
        }
    mock_df.to_dict.return_value = to_dict_return
-    
+
    return mock_df


@@ -147,11 +146,13 @@ async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_dat
@mark.asyncio
@patch('laborious.activities.api.DataFrame')
 async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
-    mock_dataframe.return_value = _create_mock_dataframe({
-        'prediction': [0.75],
-        'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
-        'timestamp': ['2024-01-01T00:00:00+00:00'],
-    })
+    mock_dataframe.return_value = _create_mock_dataframe(
+        {
+            'prediction': [0.75],
+            'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        }
+    )

    api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')

@@ -160,7 +161,7 @@ async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_
    api.send_notification_async.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
-        message='Error writing prediction data to PI Web API: Prediction write failed\n Tags: {\'tag1\': \'web_id_1\'}',
+        message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
        block='write_pi_web_api_data',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
@@ -185,7 +186,7 @@ async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_
    api.send_notification_async.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
-        message='Error writing confidence data to PI Web API: Confidence write failed\n Tags: {\'tag2\': \'web_id_2\'}',
+        message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
        block='write_pi_web_api_data',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
@@ -199,8 +200,6 @@ async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_
    assert api.pi_web_api_client.write_value.call_count == 2


-
-
@mark.asyncio
@patch('laborious.activities.api.DataFrame')
 async def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
--- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
+++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
@@ -392,7 +392,11 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo
        'prediction_confidence': 0,
        'schema': 'test_schema',
        'table_name': 'test_table',
-        'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
        'prediction_store_policy': 'erl:1',
    }

@@ -483,7 +487,9 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo
    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
    new_callable=AsyncMock,
 )
-async def test_run_none_path_flag_with_pi_web_api_and_opc(workflow_mock, format_and_export_prediction):
+async def test_run_none_path_flag_with_pi_web_api_and_opc(
+    workflow_mock, format_and_export_prediction
+):
    input_data = {
        'metadata': metadata,
        'path_flag': None,
@@ -494,7 +500,11 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc(workflow_mock, format_
        'schema': 'test_schema',
        'table_name': 'test_table',
        'opc_output_config': {'test': 'config'},
-        'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
        'prediction_store_policy': 'erl:1',
    }

@@ -532,7 +542,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc(workflow_mock, format_
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY,
-            )
+            ),
        ]
    )

@@ -590,7 +600,11 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e
        'prediction_confidence': 0,
        'schema': 'test_schema',
        'table_name': 'test_table',
-        'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
        'comment': 'test_comment',
    }

--- a/tests/laborious/workflows/subworkflows/test_prediction_process.py
+++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py
@@ -731,7 +731,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
            'model_name': model_name,
            'model_config': model_config,
            'opc_output_config': {'test': 'config'},
-            'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
            'prediction_store_policy': prediction_store_policy,
        },
        confidence,
@@ -758,7 +762,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
            'transform_table_name': 'test_transform_table',
            'comment': 'Prediction Process',
            'opc_output_config': {'test': 'config'},
-            'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
            'prediction_store_policy': prediction_store_policy,
        },
    )
@@ -792,7 +800,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
            'model_name': model_name,
            'model_config': model_config,
            'opc_output_config': {'test': 'config'},
-            'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
            'prediction_store_policy': prediction_store_policy,
        },
        confidence,