Improving coverage

Author: angelosantos4

tests/unit/cli/operations/test_run_pipeline.py

@@ -141,3 +141,37 @@ def test_json_progress_indicator(mocker):
         mocker.call("Starting Pipeline"),
         mocker.call("Pipeline Completed"),
     ]
+
+
+def test_json_progress_indicator_on_fatal_error(mocker):
+    indicator = JsonProgressIndicator(mocker.Mock(), "pipeline_name")
+    indicator.logger.error = mocker.Mock()
+    exception = Exception("Boom")
+    indicator.on_fatal_error(exception)
+    indicator.logger.error.assert_called_once_with(
+        "Pipeline Failed", exc_info=exception
+    )
+
+
+def test_json_progress_indicator_on_progress(mocker):
+    indicator = JsonProgressIndicator(mocker.Mock(), "pipeline_name")
+    indicator.logger.info = mocker.Mock()
+    indicator.progress_callback(1000, Metrics())
+    indicator.logger.info.assert_called_once_with(
+        "Processing Record", extra={"index": 1000}
+    )
+
+
+def test_json_progress_indicator_on_finish_with_exception(mocker):
+    indicator = JsonProgressIndicator(mocker.Mock(), "pipeline_name")
+    indicator.logger.info = mocker.Mock()
+    indicator.logger.error = mocker.Mock()
+    exception = Exception("Boom")
+    indicator.on_fatal_error(exception)
+    with pytest.raises(Exception):
+        indicator.on_finish(Metrics())
+    indicator.logger.error.assert_called_once_with(
+        "Pipeline Failed", exc_info=exception
+    )
+    indicator.logger.info.assert_called_once_with("Pipeline Completed")
+    assert indicator.exception is exception

tests/unit/test_metrics.py

@@ -136,3 +136,12 @@ def test_prometheus_metric_handler_import_error(mocker):
         "The prometheus_client library is required to use the PrometheusMetricHandler."
         in str(excinfo.value)
     )
+
+
+def test_aggregate_handler_tick_calls_all_handlers(mocker):
+    mock_handler1 = mocker.Mock()
+    mock_handler2 = mocker.Mock()
+    handler = AggregateHandler([mock_handler1, mock_handler2])
+    handler.tick()
+    mock_handler1.tick.assert_called_once()
+    mock_handler2.tick.assert_called_once()