diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 19d52cd..e1ba49d 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -402,3 +402,42 @@ async def test_get_last_timestamp_no_data(gates_activity): # Assert assert isinstance(result, str) # Should be a timestamp string assert len(result) > 0 + + +@mark.asyncio +@patch('laborious.activities.gates.metrics') +async def test_write_metrics(mock_metrics, gates_activity): + """Test write_metrics method.""" + input_data = { + **metadata, + 'prediction': { + 'prediction': [1, 2, 3], + 'prediction_confidence': [0.9, 0.8, 0.7], + 'response_time': [0.1, 0.2, 0.3] + } + } + await gates_activity.write_metrics(input_data) + mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with( + pod_id=gates_activity.pod_id, + model_name=metadata['metadata']['model_name'], + pipeline_name=metadata['metadata']['workflow_name'] + ) + mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with() + + mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with( + pod_id=gates_activity.pod_id, + model_name=metadata['metadata']['model_name'], + pipeline_name=metadata['metadata']['workflow_name'] + ) + mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with( + 0.9 + ) + + mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with( + pod_id=gates_activity.pod_id, + model_name=metadata['metadata']['model_name'], + pipeline_name=metadata['metadata']['workflow_name'] + ) + mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with( + 0.1 + )