SIENTIAPDE-1712

Update environment variables in values.yaml and enhance metric labels in metrics.py

- Changed POSTGRES_USER and POSTGRES_PASSWORD values in values.yaml for improved security.
- Added 'runtime' label to metrics in metrics.py for better environment identification.
- Updated CORE_LABELS to include 'runtime' for consistency across metrics.
- Modified type hint for data parameter in PredictionProcess to use a dictionary for better clarity.
- Adjusted tests to reflect changes in core labels and MinIO configuration.
This commit is contained in:
vitor-aignosi
2026-03-20 16:59:04 -03:00
parent e35e27bfb2
commit 2979f2dd5a
11 changed files with 67 additions and 134 deletions

View File

@@ -45,6 +45,7 @@ def test___init__(
'secret_key': 'minio123',
'default_bucket': 'test',
'retention_hours': 24,
'secure': False,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
@@ -144,13 +145,14 @@ def test___init__(
)
mock_minio_repository.assert_called_once_with(
endpoint_url=minio_config['endpoint_url'],
endpoint=minio_config['endpoint_url'],
access_key=minio_config['access_key'],
secret_key=minio_config['secret_key'],
bucket=minio_config['default_bucket'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
secure=minio_config['secure'],
)
@@ -188,6 +190,7 @@ async def test_shutdown(
'secret_key': 'minio123',
'default_bucket': 'test',
'retention_hours': 24,
'secure': False,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}

View File

@@ -78,7 +78,9 @@ def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_clien
'get_core_labels',
return_value={
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'operation_type': '-',
},
):
@@ -86,6 +88,12 @@ def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_clien
metadata=metadata['metadata'], operation_type=None
)
assert 'operation_type' not in labels
assert labels == {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@patch('laborious.activities.api.PIWebAPIClient')
@@ -105,7 +113,9 @@ def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
'get_core_labels',
return_value={
'pod_id': 'test_pod',
'runtime': 'k8s',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'operation_type': 'write',
},
):
@@ -113,6 +123,7 @@ def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
metadata=metadata['metadata'], operation_type='write'
)
assert labels['operation_type'] == 'write'
assert labels['runtime'] == 'k8s'
def test__init__():
@@ -152,6 +163,7 @@ def api(mock_pi_web_api_client):
api_instance.get_core_labels = MagicMock(
return_value={
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@@ -334,6 +346,7 @@ async def test_process_pi_web_api_response_success(api):
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@@ -367,6 +380,7 @@ async def test_process_pi_web_api_response_with_errors(api):
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@@ -395,6 +409,7 @@ async def test_process_pi_web_api_response_missing_tags(api):
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@@ -427,6 +442,7 @@ async def test_process_pi_web_api_response_missing_webid(api):
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
@@ -455,6 +471,7 @@ async def test_process_pi_web_api_response_missing_tag_name(api):
tags = {'tag1': 'web_id_1'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}

View File

@@ -820,15 +820,17 @@ async def test_write_metrics(mock_metrics, gates_activity):
'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
}
await gates_activity.write_metrics(input_data)
core_tags = {
'pod_id': gates_activity.pod_id,
'runtime': gates_activity.runtime,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
}
gates_activity.emit_metric.assert_has_calls(
[
call(
metric_object=mock_metrics.PREDICTIONS_WRITTEN_COUNT,
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
},
tags=core_tags,
),
]
)
@@ -837,11 +839,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
call(
metric_object=mock_metrics.PREDICTION_CONFIDENCE_MONITOR,
method='set',
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
},
tags=core_tags,
value=0.9,
),
]
@@ -851,11 +849,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
call(
metric_object=mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR,
method='observe',
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
},
tags=core_tags,
value=0.1,
),
]
@@ -865,9 +859,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
call(
metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT,
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
**core_tags,
'opc_server_id': 'server1',
'tag': 'tag1',
},
@@ -880,9 +872,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
method='observe',
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
**core_tags,
'opc_server_id': 'server1',
'tag': 'tag1',
},
@@ -895,9 +885,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
call(
metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT,
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
**core_tags,
'opc_server_id': 'server1',
'tag': 'tag2',
},
@@ -910,9 +898,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
method='observe',
tags={
'pod_id': gates_activity.pod_id,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
**core_tags,
'opc_server_id': 'server1',
'tag': 'tag2',
},
@@ -943,6 +929,7 @@ async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_act
method='observe',
tags={
'pod_id': gates_activity.pod_id,
'runtime': gates_activity.runtime,
'model_name': metadata['metadata']['model_name'],
'workflow_name': metadata['metadata']['workflow_name'],
'opc_server_id': 'server1',

View File

@@ -117,80 +117,6 @@ def test___init___done_repository(mock_minio_repository, storage):
assert storage.minio_repository is not None
@mark.asyncio
async def test_query_to_minio_minio_repository_not_initialized(storage):
storage.minio_repository = None
with raises(ValueError) as e:
await storage.query_to_minio({})
assert str(e.value) == 'Minio repository not initialized'
@mark.asyncio
async def test_query_to_minio_not_data(storage):
storage.load_custom_query = AsyncMock(return_value=None)
result = await storage.query_to_minio({})
storage.load_custom_query.assert_called_once_with({})
assert result['success'] is False
assert result['message'] == 'No data returned from query'
@mark.asyncio
@patch('laborious.activities.storage.pd.DataFrame')
@patch('laborious.activities.storage.now')
async def test_query_to_minio_success(now, dataframe, storage):
data = [{'a': 1}, {'a': 2}, {'a': 3}]
storage.load_custom_query = AsyncMock(return_value=data)
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
storage.minio_repository.upload_file = AsyncMock(
return_value={
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
}
)
storage.minio_repository.bucket = 'test'
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
dataframe.assert_called_once_with(data)
storage.minio_repository.upload_file.assert_called_once_with(
file_bytes=ANY,
relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet',
metadata=metadata['metadata'],
)
assert result['success'] is True
assert (
result['object_key']
== 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
)
assert (
result['uri']
== 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
)
@mark.asyncio
async def test_query_to_minio_error(storage):
storage.send_notification = MagicMock()
storage.send_notification_async = AsyncMock()
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
assert result['success'] is False
assert result['message'] == 'test'
storage.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='ERROR_STORING_QUERY_TO_MINIO',
message='Error storing query to MinIO: test',
block='query_to_minio',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
def test_close(storage):
storage.minio_repository = MagicMock()
@@ -275,8 +201,11 @@ async def test_cleanup_minio_objects_expired(mock_now, storage):
storage.minio_repository.delete_file = AsyncMock()
storage.send_notification_async = AsyncMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = await storage.cleanup_minio_objects_expired(
{**metadata, 'prefix': 'training_datasets/m'}
{**metadata, 'data': data_mock}
)
assert result['deleted_count'] == 1
@@ -326,8 +255,10 @@ async def test_export_payload_to_postgres(storage):
async def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
storage.minio_repository = None
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'test'
with raises(ValueError, match='Minio repository not initialized'):
await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'})
await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
@mark.asyncio
@@ -340,7 +271,9 @@ async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage):
storage.minio_repository.delete_file = AsyncMock()
storage.send_notification_async = AsyncMock()
result = await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'})
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'test'
result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
assert result['deleted_count'] == 0
assert result['failed_count'] == 0
@@ -356,8 +289,10 @@ async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage):
storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error'))
storage.send_notification_async = AsyncMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = await storage.cleanup_minio_objects_expired(
{**metadata, 'prefix': 'training_datasets/m'}
{**metadata, 'data': data_mock}
)
assert result['deleted_count'] == 0
@@ -375,8 +310,10 @@ async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storag
storage.send_notification_async = AsyncMock()
storage.error = MagicMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = await storage.cleanup_minio_objects_expired(
{**metadata, 'prefix': 'training_datasets/m'}
{**metadata, 'data': data_mock}
)
assert result['deleted_count'] == 0

View File

@@ -102,6 +102,7 @@ def test_build_minio_config_with_env_vars():
'secret_key': 'test-secret',
'default_bucket': 'test-bucket',
'retention_hours': 24,
'secure': False,
}
@@ -117,4 +118,5 @@ def test_build_minio_config_with_defaults():
'secret_key': 'minioadmin',
'default_bucket': 'laborious',
'retention_hours': 24,
'secure': False,
}

View File

@@ -38,7 +38,7 @@ async def test_run(workflow_mock, prediction_process):
# Arrange
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -199,7 +199,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
# Arrange
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -251,7 +251,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
# Arrange
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -336,7 +336,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
# Arrange
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -441,7 +441,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
# Arrange
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -784,7 +784,7 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
data_payload = MagicMock()
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
data_payload.last_timestamp = '2024-01-01'
data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
input_data = {
'metadata': metadata,
'data': data_payload,
@@ -816,7 +816,7 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
workflow_mock.execute_activity_method.assert_any_call(
Activities.cleanup_minio_objects_expired,
{**metadata, 'prefix': 'training_datasets/test'},
{**metadata, 'data': data_payload},
retry_policy=ANY,
start_to_close_timeout=ANY,
)

View File

@@ -6,15 +6,6 @@ from laborious.activities.activities import Activities
from laborious.workflows.minimal_retrain import MinimalRetrain
@fixture(autouse=True)
def _passthrough_from_dict():
with patch(
'laborious.workflows.minimal_retrain.MinioDataFramePayload.from_dict',
side_effect=lambda x: x,
):
yield
@fixture
def minimal_retrain() -> MinimalRetrain:
return MinimalRetrain()

View File

@@ -22,13 +22,9 @@ metadata = {
@mark.asyncio
@patch(
'laborious.workflows.predictions_batch.MinioDataFramePayload.from_dict', side_effect=lambda x: x
)
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch: PredictionsBatch):
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
activity_return = MagicMock()
activity_return.cleanup_prefix.return_value = None
workflow_mock.execute_activity_method.return_value = activity_return
input_data = {
@@ -66,7 +62,6 @@ async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch:
prediction_input = {
'metadata': metadata,
'data': activity_return,
'cleanup_prefix': activity_return.cleanup_prefix(),
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],