gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-994

Browse Source 
Refactor and enhance the laborious workflow and utilities

- Removed outdated test file `test_predictions_batch.py` from workflows.
- Added `input_sample.json` for standardized input configuration.
- Introduced `connectors_config.py` to manage database and service configurations.
- Implemented a logging utility in `logger.py` for consistent logging across the application.
- Created `policies.py` to define retry policies for workflows.
- Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`.
- Added extensive tests for `OpcRepository` in `test_opc_repository.py`.
- Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
This commit is contained in:
vitor-aignosi
2025-05-23 17:34:47 -03:00
parent 5fe552410b
commit 67fe4afaa6
30 changed files with 1385 additions and 765 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/laborious/activities/test_base.py
View File

@@ -1,6 +1,6 @@
from unittest.mock import MagicMock
from laborious.activities.base import BaseActivity
from pytest import fixture
from pytest import fixture, mark
from sientia_do.notifications.models import Notification
@@ -12,7 +12,8 @@ def base_activity():
)
def test_prepare_activity(base_activity):
@mark.asyncio
async def test_prepare_activity(base_activity):
base_activity.notification_handler.base_notification = Notification(
project="project",
pipeline="pipeline",
@@ -21,12 +22,14 @@ def test_prepare_activity(base_activity):
model_id="-",
)
base_activity.prepare_activity(
schedule_name="test_schedule",
model_name="test_model",
model_id="test_model_id",
)
await base_activity.prepare_activity({
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id'
})
assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule"
assert base_activity.notification_handler.base_notification.model_name == "test_model"
assert base_activity.notification_handler.base_notification.model_id == "test_model_id"
assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow"

26
tests/laborious/activities/test_gates.py
View File

@@ -51,7 +51,7 @@ async def test_input_gate_specific_variables_null_values_with_stop_policy_only(
}
result = await gates.input_gate(input_data)
assert result == ('stop', -1)
assert result == ('stop', -1, 'Input data with bad quality')
input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame(
@@ -98,7 +98,7 @@ async def test_input_gate_specific_variables_null_values_with_continue_policy_on
}
result = await gates.input_gate(input_data)
assert result == ('continue', 2)
assert result == ('continue', 2, 'Input data with bad quality')
input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame(
@@ -145,7 +145,7 @@ async def test_input_gate_specific_variables_null_values_no_filtered(
}
result = await gates.input_gate(input_data)
assert result == (None, 0)
assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
@@ -197,7 +197,7 @@ async def test_input_gate_one_stop_policy(
}
result = await gates.input_gate(input_data)
assert result == ('stop', -1)
assert result == ('stop', -1, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -251,7 +251,7 @@ async def test_input_gate_one_continue_policy(
}
result = await gates.input_gate(input_data)
assert result == ('continue', 2)
assert result == ('continue', 2, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -305,7 +305,7 @@ async def test_input_gate_no_filtered(
}
result = await gates.input_gate(input_data)
assert result == (None, 0)
assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -340,7 +340,7 @@ async def test_input_gate_error(
}
result = await gates.input_gate(input_data)
assert result == (None, 0)
assert result == (None, 0, '')
gates.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES',
@@ -389,7 +389,7 @@ async def test_mlflow_response_gate_no_filtered(
}
result = await gates.mlflow_response_gate(input_data)
assert result == (None, 0)
assert result == (None, 0, '')
api_error_filter_mock.assert_called_once_with(
input_data['data'],
@@ -433,7 +433,7 @@ async def test_mlflow_response_gate_filtered(
}
result = await gates.mlflow_response_gate(input_data)
assert result == ('continue', 255)
assert result == ('continue', 255, "Error")
api_error_filter_mock.assert_called_once_with(
input_data['data'],
@@ -480,7 +480,7 @@ async def test_mlflow_content_gate_no_filtered(
}
result = await gates.mlflow_content_gate(input_data)
assert result == (None, 0)
assert result == (None, 0, '')
nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame(
@@ -504,7 +504,8 @@ async def test_mlflow_content_gate_filtered(
if x == 'path_confidence':
return transform_filter_path_confidence
mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
mlflow_content_filter_functions_mock.__getitem__.side_effect = \
transform_filter_functions_side_effect
input_data = {
'filters': {
@@ -521,7 +522,8 @@ async def test_mlflow_content_gate_filtered(
}
result = await gates.mlflow_content_gate(input_data)
assert result == ('repeat', -1)
assert result == (
'repeat', -1, "Transformed data not passed the content filter")
nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame(

3
tests/laborious/activities/test_mlflow.py
View File

@@ -61,7 +61,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
mlflow.model_monitoring_repository.transform.return_value = expected_response
# Call the method
response_data, timestamp = await mlflow.request_transform(input_data)
response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data'])
@@ -75,7 +75,6 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
# Verify the response
assert response_data == expected_response
assert timestamp == '2024-01-02'
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(

227
tests/laborious/activities/test_opc.py
View File

@@ -1,111 +1,120 @@
from unittest.mock import patch, MagicMock
from unittest.mock import patch, MagicMock, ANY, call
from pytest import fixture, mark
from laborious.activities.opc import NotificationLevel
from laborious.activities.opc import OPC
from sientia_do.notifications.models import NotificationLevel
from unittest.mock import ANY
@patch("laborious.activities.opc.OpcRepository")
def test___init__(mock_opc_repository):
mock_logger = MagicMock()
server1 = MagicMock()
server2 = MagicMock()
mock_opc_repository.side_effect = [server1, server2]
mock_notification_handler = MagicMock()
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
'server2': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
opc = OPC(
name="test",
url="http://localhost:8080",
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
logger=MagicMock(),
notification_handler=MagicMock()
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
)
assert opc.name == "test"
assert opc.url == "http://localhost:8080"
assert opc.server_uri == "opc.tcp://localhost:4840"
assert opc.cert_path == ""
assert opc.private_key_path == ""
assert opc.server_cert_path == ""
assert opc.opc_repository == mock_opc_repository.return_value
assert opc.opc_servers == servers
assert opc.logger == mock_logger
assert opc.notification_handler == mock_notification_handler
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_called_once_with(
name="test",
url="http://localhost:8080",
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
logger=opc.logger,
)
mock_opc_repository.assert_has_calls([
call(
name="server1",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
),
])
mock_opc_repository.assert_has_calls([
call(
name="server2",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
)
])
opc.opc_repository.connect.assert_called_once()
server1.connect.assert_called_once()
server2.connect.assert_called_once()
@fixture
@patch("laborious.activities.opc.OpcRepository")
def opc(mock_opc_repository):
def opc(_mock_opc_repository):
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
return OPC(
name="test",
url="http://localhost:8080",
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
async def test_write_opc_data_success(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
# Act
await opc.write_opc_data(input_data)
# Assert
opc.opc_repository.write_data.assert_any_call('tag1', 0.75, 'float')
opc.opc_repository.write_data.assert_any_call('tag2', 0.95, 'float')
assert opc.opc_repository.write_data.call_count == 2
WRITE_DATA_CASES = [
('tag1', 'int', 50),
('tag2', 'float', 50.5),
('tag3', 'bool', True),
('tag4', 'string', 'test'),
]
@mark.asyncio
async def test_write_opc_data_prediction_error(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
}
}
}
@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')
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type)
opc.opc_repository.write_data.side_effect = Exception("Test error")
# Act
await opc.write_opc_data(input_data)
# Assert
opc.notification_handler.build_and_send_notification.assert_called_with(
def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
opc.write_data(server='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction')
opc.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
@@ -116,44 +125,48 @@ async def test_write_opc_data_prediction_error(opc):
@mark.asyncio
async def test_write_opc_data_confidence_error(opc):
async def test_write_opc_data_success(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
}
# Make first call succeed but second fail
def side_effect(*args, **kwargs):
if args[0] == 'tag2':
raise ValueError("Test error")
return None
opc.opc_repository.write_data.side_effect = side_effect
# Act
opc.write_data = MagicMock()
await opc.write_opc_data(input_data)
# Assert
opc.notification_handler.build_and_send_notification.assert_called_with(
notification_id="WRITE_OPC_CONFIDENCE_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
opc.logger.error.assert_called_once()
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction'
)])
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence'
)
])
assert opc.write_data.call_count == 2
@mark.asyncio
@@ -175,4 +188,4 @@ async def test_write_opc_data_empty_config(opc):
await opc.write_opc_data(input_data)
# Assert
opc.opc_repository.write_data.assert_not_called()
opc.opc_repository['server1'].write_data.assert_not_called()

106
tests/laborious/utils/filters/repository/test_opc_repository.py
View File

@@ -1,106 +0,0 @@
from unittest.mock import Mock, patch, MagicMock
from pathlib import Path
from asyncua.sync import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType
from pytest import fixture
from laborious.utils.repository.opc_repository import OpcRepository
@fixture
def mock_logger():
return Mock()
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
)
@fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = MagicMock()
mock.return_value = client_instance
yield client_instance
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.non_receive_count == 0
assert opc_repository.client is None
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
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.connect()
mock_client.connect.assert_called_once()
assert opc_repository.client == mock_client
def test_connect_without_security(opc_repository, mock_client):
opc_repository.cert_path = None
opc_repository.connect()
mock_client.connect.assert_called_once()
assert opc_repository.client == mock_client
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.disconnect()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
def test_write_data(opc_repository, mock_client, mock_logger):
opc_repository.client = mock_client
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float", mock_logger)
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
mock_logger.info.assert_called_once_with(
"Writing 42.0 - <class 'float'> to " + str(mock_node))

6
tests/laborious/utils/filters/test_conditional_filters.py
View File

@@ -7,18 +7,18 @@ def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'VARIABLES': ['variable2']}) == True
config={'VARIABLES': ['variable2']}) is False
def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'VARIABLES': ['variable2']}) == False
config={'VARIABLES': ['variable2']}) is True
def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) == True
assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data():

0
tests/laborious/utils/filters/repository/test_model_repository.py → tests/laborious/utils/repository/test_model_repository.py
View File

242
tests/laborious/utils/repository/test_opc_repository.py Normal file
View File

@@ -0,0 +1,242 @@
from unittest.mock import 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
def mock_logger():
return Mock()
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
notification_handler=Mock(),
reconnection_interval=60,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
)
@fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = MagicMock()
mock.return_value = client_instance
yield client_instance
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.reconnection_interval == 60
assert opc_repository.client is None
assert opc_repository.last_reconnection_time is None
assert opc_repository.error_count == 0
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
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()
opc_repository.try_connect.assert_called_once()
assert opc_repository.client == mock_client
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.assert_called_once()
opc_repository.set_security.assert_not_called()
assert opc_repository.client == mock_client
def test_try_connect_sucess(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
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")
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.disconnect()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
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.assert_called_once()
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()
response = opc_repository.validate_connection()
assert response == opc_repository.connect.return_value
opc_repository.disconnect.assert_called_once()
opc_repository.connect.assert_called_once()
opc_repository.logger.error.assert_has_calls(
[
call("Failed to disconnect from OPC server: Test error"),
]
)
@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):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_not_called()
assert response is False
@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):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_called_once()
assert response == opc_repository.try_connect.return_value
def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=False)
opc_repository.client = MagicMock()
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.get_node.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}",
message="Failed to get node from OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
assert opc_repository.error_count == 1
def test_write_data(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = mock_client
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.logger.info.assert_called_once_with(
"Writing 42.0 - <class 'float'> to " + str(mock_node))
def test_write_data_write_value_failed(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = mock_client
mock_node = MagicMock()
opc_repository.error_count = 0
mock_client.get_node.return_value = mock_node
mock_node.write_value.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
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()
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}",
message="Failed to write data to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
assert opc_repository.error_count == 1

44
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -1,4 +1,4 @@
from unittest.mock import call, patch, AsyncMock
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from laborious.activities.activities import Activities
@@ -28,7 +28,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_prediction,
{
@@ -36,7 +36,9 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
}
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
@@ -44,8 +46,10 @@ 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
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
@@ -53,12 +57,15 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value
}
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@@ -80,7 +87,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_default_prediction,
{
@@ -88,7 +95,9 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment']
}
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
@@ -97,8 +106,10 @@ 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': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
@@ -107,9 +118,12 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value
}
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1

305
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -1,4 +1,4 @@
from unittest.mock import AsyncMock, patch, call
from unittest.mock import AsyncMock, patch, call, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@@ -18,65 +18,78 @@ async def test_run(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict)
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'],
'data': input_data['data']
})])
workflow_mock.execute_activity_method.assert_has_calls([
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
@@ -85,9 +98,10 @@ async def test_run(workflow_mock, prediction_process):
'data': 'predicted_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 'test_model',
'model_id': 1,
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'opc_output_config': input_data['opc_output_config']
}
)
@@ -101,27 +115,34 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95), # input_gate
('stop', 0.95, "Input data with bad quality"), # input_gate
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 2
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']}),
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -135,42 +156,53 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95), # input_gate
('repeat', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 4
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
'model_retention': input_data['model_retention']},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -184,49 +216,61 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 5
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -240,63 +284,75 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict)
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -305,7 +361,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'stop'
path_flag = 'STOP'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
@@ -317,12 +373,12 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()
@@ -343,7 +399,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert
@@ -353,8 +409,10 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
{
'schema': schema,
'table_name': table_name,
'model': model
}
'model_id': model
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
workflow_mock.execute_child_workflow.assert_not_called()
@@ -364,7 +422,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'continue'
path_flag = 'CONTINUE'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
@@ -376,7 +434,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, 'Prediction Process'
)
# Assert
@@ -391,7 +449,10 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process'
}
)
@@ -413,7 +474,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert

48
tests/laborious/workflows/subworkflows/test_predictions_batch.py
View File

@@ -1,48 +0,0 @@
from unittest.mock import AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@fixture
def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_activity_method.return_value = {
'data': 'test_data'
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'query': 'SELECT * FROM test'
}
await predictions_batch.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id']
}
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query']
)
])
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', input_data)
])

68
tests/laborious/workflows/test_predictions_batch.py Normal file
View File

@@ -0,0 +1,68 @@
from unittest.mock import AsyncMock, call, patch, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@fixture
def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': 'test_opc_output_config'
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
prediction_input = {
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])