Created tests for the change.

Author: angelosantos4

nodestream/metrics.py

@@ -23,8 +23,8 @@ def decrement_on(self, handler: "MetricHandler", value: Number = 1):
         """Decrement this metric on the given handler."""
         handler.decrement(self, value)
 
-    def register(self):
-        Metrics().get().increment(self, 0)
+    def register(self, handler: "MetricHandler"):
+        handler.increment(self, 0)
         return self
 
 
@@ -175,7 +175,7 @@ def stop(self):
             self.thread.join()
             self.logger.info("Prometheus metrics server shut down successfully")
 
-        def get_guage(self, metric: Metric) -> Gauge:
+        def get_gauge(self, metric: Metric) -> Gauge:
             """Get the Gauge for the given metric, creating it if it doesn't exist."""
             if metric not in self.instruments_by_metric:
                 self.instruments_by_metric[metric] = Gauge(
@@ -186,10 +186,11 @@ def get_guage(self, metric: Metric) -> Gauge:
             return self.instruments_by_metric[metric]
 
         def increment(self, metric: Metric, value: Number):
-            self.get_guage(metric).inc(value)
+            print(f"Incrementing {metric.name} by {value}")
+            self.get_gauge(metric).inc(value)
 
         def decrement(self, metric: Metric, value: Number):
-            self.get_guage(metric).dec(value)
+            self.get_gauge(metric).dec(value)
 
         def tick(self):
             pass

nodestream/pipeline/pipeline.py

@@ -58,6 +58,7 @@ async def emit_record(self, record):
     async def drive_step(self):
         try:
             while (next_record := await self.input.get()) is not None:
+                print(f"Driving step with record {next_record}")
                 results = self.step.process_record(next_record, self.context)
                 async for record in results:
                     if not await self.emit_record(record):
@@ -172,6 +173,7 @@ async def run(self, reporter: PipelineProgressReporter):
         # the flow of records between the steps.
         executors: List[StepExecutor] = []
         current_input, current_output = channel(self.step_outbox_size)
+        print(current_input, current_output)
         pipeline_output = PipelineOutput(current_input, reporter)
 
         # Create the executors for the steps in the pipeline. The executors

tests/unit/cli/operations/test_run_pipeline.py

@@ -6,6 +6,7 @@
     ProgressIndicator,
     RunPipeline,
     SpinnerProgressIndicator,
+    JsonProgressIndicator,
 )
 from nodestream.metrics import Metrics
 from nodestream.project import PipelineConfiguration, PipelineDefinition, Project
@@ -59,18 +60,21 @@ def test_spinner_on_finish(mocker):
 def test_spinner_progress_callback(mocker):
     spinner = SpinnerProgressIndicator(mocker.Mock(), "pipeline_name")
     spinner.on_start()
-    spinner.progress_callback(1000, None)
+    mock_metrics = mocker.Mock()
+    spinner.progress_callback(1000, mock_metrics)
     spinner.progress.set_message.assert_called_once()
+    mock_metrics.tick.assert_called_once()
 
 
 def test_spinner_error_condition(mocker):
     spinner = SpinnerProgressIndicator(mocker.Mock(), "pipeline_name")
     spinner.on_start()
     spinner.on_fatal_error(Exception())
     spinner.progress.set_message.assert_called_once()
-
+    mock_metrics = mocker.Mock()
     with pytest.raises(Exception):
-        spinner.on_finish(Metrics())
+        spinner.on_finish(mock_metrics)
+    mock_metrics.tick.assert_called_once()
 
 
 @pytest.mark.parametrize(
@@ -126,3 +130,14 @@ def test_progress_indicator_error(mocker):
     indicator = ProgressIndicator(mocker.Mock(), "pipeline_name")
     indicator.on_fatal_error(Exception("Boom"))
     indicator.command.line.assert_called_with("<error>Boom</error>")
+
+
+def test_json_progress_indicator(mocker):
+    indicator = JsonProgressIndicator(mocker.Mock(), "pipeline_name")
+    indicator.logger.info = mocker.Mock()
+    indicator.on_start()
+    indicator.on_finish(Metrics())
+    assert indicator.logger.info.call_args_list == [
+        mocker.call("Starting Pipeline"),
+        mocker.call("Pipeline Completed"),
+    ]

tests/unit/databases/test_query_executor_with_statistics.py

@@ -1,9 +1,11 @@
 import pytest
+from unittest.mock import call
 
 from nodestream.databases.query_executor_with_statistics import (
     QueryExecutorWithStatistics,
 )
-from nodestream.metrics import Metric, Metrics
+from nodestream.metrics import Metrics, NodestreamMetricRegistry
+from nodestream.model import Node, RelationshipWithNodes, Relationship
 
 
 @pytest.fixture
@@ -18,13 +20,23 @@ async def test_upsert_nodes_in_bulk_with_same_operation_increments_counter_by_si
     with Metrics.capture() as metrics:
         metrics.increment = mocker.Mock()
         await query_executor_with_statistics.upsert_nodes_in_bulk_with_same_operation(
-            "operation", ["node1", "node2"]
+            "operation",
+            [Node("node_type", "node1", "id1"), Node("node_type", "node2", "id2")],
         )
         query_executor_with_statistics.inner.upsert_nodes_in_bulk_with_same_operation.assert_awaited_once_with(
             "operation",
-            ["node1", "node2"],
+            [Node("node_type", "node1", "id1"), Node("node_type", "node2", "id2")],
+        )
+
+        assert "node_type" in query_executor_with_statistics.node_metric_by_type
+        assert (
+            call(query_executor_with_statistics.node_metric_by_type["node_type"], 2)
+            in metrics.increment.call_args_list
+        )
+        assert (
+            call(NodestreamMetricRegistry.NODES_UPSERTED, 2)
+            in metrics.increment.call_args_list
         )
-        metrics.increment.assert_called_once_with(Metric.NODES_UPSERTED, 2)
 
 
 @pytest.mark.asyncio
@@ -34,13 +46,44 @@ async def test_upsert_relationships_in_bulk_of_same_operation_increments_counter
     with Metrics.capture() as metrics:
         metrics.increment = mocker.Mock()
         await query_executor_with_statistics.upsert_relationships_in_bulk_of_same_operation(
-            "operation", ["relationship1", "relationship2"]
+            "operation",
+            [
+                RelationshipWithNodes(
+                    "node1", "node2", Relationship("relationship_type")
+                ),
+                RelationshipWithNodes(
+                    "node3", "node4", Relationship("relationship_type")
+                ),
+            ],
         )
         query_executor_with_statistics.inner.upsert_relationships_in_bulk_of_same_operation.assert_awaited_once_with(
             "operation",
-            ["relationship1", "relationship2"],
+            [
+                RelationshipWithNodes(
+                    "node1", "node2", Relationship("relationship_type")
+                ),
+                RelationshipWithNodes(
+                    "node3", "node4", Relationship("relationship_type")
+                ),
+            ],
+        )
+        assert (
+            "relationship_type"
+            in query_executor_with_statistics.relationship_metric_by_relationship_type
+        )
+        assert (
+            call(
+                query_executor_with_statistics.relationship_metric_by_relationship_type[
+                    "relationship_type"
+                ],
+                2,
+            )
+            in metrics.increment.call_args_list
+        )
+        assert (
+            call(NodestreamMetricRegistry.RELATIONSHIPS_UPSERTED, 2)
+            in metrics.increment.call_args_list
         )
-        metrics.increment.assert_called_once_with(Metric.RELATIONSHIPS_UPSERTED, 2)
 
 
 @pytest.mark.asyncio
@@ -53,7 +96,9 @@ async def test_perform_ttl_op_increments_counter_by_one(
         query_executor_with_statistics.inner.perform_ttl_op.assert_awaited_once_with(
             "config"
         )
-        metrics.increment.assert_called_once_with(Metric.TIME_TO_LIVE_OPERATIONS)
+        metrics.increment.assert_called_once_with(
+            NodestreamMetricRegistry.TIME_TO_LIVE_OPERATIONS
+        )
 
 
 @pytest.mark.asyncio
@@ -66,7 +111,9 @@ async def test_execute_hook_increments_counter_by_one(
         query_executor_with_statistics.inner.execute_hook.assert_awaited_once_with(
             "hook"
         )
-        metrics.increment.assert_called_once_with(Metric.INGEST_HOOKS_EXECUTED)
+        metrics.increment.assert_called_once_with(
+            NodestreamMetricRegistry.INGEST_HOOKS_EXECUTED
+        )
 
 
 @pytest.mark.asyncio

tests/unit/pipeline/test_pipeline.py

@@ -1,6 +1,9 @@
 import pytest
 
+from nodestream.pipeline import Extractor
+from nodestream.pipeline.object_storage import NullObjectStore
 from nodestream.pipeline.pipeline import (
+    Pipeline,
     PipelineOutput,
     PipelineProgressReporter,
     Step,
@@ -9,6 +12,7 @@
     StepInput,
     StepOutput,
 )
+from nodestream.pipeline.channel import Channel
 
 
 @pytest.fixture
@@ -136,3 +140,43 @@ def on_start_callback():
 
     output.call_handling_errors(output.reporter.on_start_callback)
     output.reporter.logger.exception.assert_called_once()
+
+
+class TestExtractor(Extractor):
+    async def extract_records(self):
+        return None
+
+
+class TestStep2(Step):
+    async def process_record(self, record, context: StepContext):
+        yield record
+
+
+@pytest.mark.asyncio
+async def test_pipeline_channels_obtain_input_and_output_names_from_steps(mocker):
+    channel = mocker.Mock(Channel)
+    mock_input = mocker.Mock(StepInput, channel=channel)
+    mock_output = mocker.Mock(StepOutput, channel=channel)
+    mock_input.get.return_value = None
+    mock_output.put.return_value = True
+    with mocker.patch(
+        "nodestream.pipeline.pipeline.channel", return_value=(mock_input, mock_output)
+    ):
+        step1 = TestExtractor()
+        step2 = TestStep2()
+        pipeline = Pipeline(
+            steps=[step1, step2], step_outbox_size=1, object_store=NullObjectStore()
+        )
+        await pipeline.run(
+            reporter=PipelineProgressReporter(
+                on_start_callback=lambda: None, logger=mocker.Mock()
+            )
+        )
+        assert mock_input.register.call_args_list == [
+            mocker.call("TestStep2"),
+            mocker.call("TestExtractor"),
+        ]
+        assert mock_output.register.call_args_list == [
+            mocker.call("TestStep2"),
+            mocker.call("TestExtractor"),
+        ]

tests/unit/test_metrics.py

@@ -7,16 +7,43 @@
     ConsoleMetricHandler,
     JsonLogMetricHandler,
     Metric,
+    NodestreamMetricRegistry,
     Metrics,
     NullMetricHandler,
     PrometheusMetricHandler,
 )
 
 
+def test_metric_registry_contains_subclasses_in_all_metrics():
+    for metric in NodestreamMetricRegistry.get_all_metrics().values():
+        assert metric in NodestreamMetricRegistry._subclasses
+
+
+def test_metric_increment_on_handler_increments_metric_on_handler(mocker):
+    handler = mocker.Mock()
+    metric = NodestreamMetricRegistry.RECORDS
+    metric.increment_on(handler, 1)
+    handler.increment.assert_called_once_with(metric, 1)
+
+
+def test_metric_decrement_on_handler_decrements_metric_on_handler(mocker):
+    handler = mocker.Mock()
+    metric = NodestreamMetricRegistry.RECORDS
+    metric.decrement_on(handler, 1)
+    handler.decrement.assert_called_once_with(metric, 1)
+
+
+def test_prometheus_metric_handler_registers_new_metrics(mocker):
+    handler = PrometheusMetricHandler()
+    metric = Metric("new_metric", "New metric")
+    metric.register(handler)
+    assert metric in handler.instruments_by_metric.keys()
+
+
 def test_null_metric_handler():
     handler = NullMetricHandler()
-    handler.increment(Metric.RECORDS, 1)
-    handler.decrement(Metric.RECORDS, 1)
+    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
+    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
     # No assertions needed as NullMetricHandler does nothing
 
 
@@ -26,25 +53,27 @@ def test_prometheus_metric_handler(mocker):
     handler = PrometheusMetricHandler()
     handler.start()
     mock_start_http_server.assert_called_once()
-    handler.increment(Metric.RECORDS, 1)
-    handler.decrement(Metric.RECORDS, 1)
+    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
+    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
     handler.stop()
 
 
 def test_console_metric_handler(mocker):
     mock_command = mocker.Mock()
     handler = ConsoleMetricHandler(mock_command)
-    handler.increment(Metric.RECORDS, 1)
-    handler.decrement(Metric.RECORDS, 1)
+    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
+    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.tick()
     handler.stop()
     mock_command.table.assert_called_once()
 
 
 def test_json_log_metric_handler(mocker):
     mock_logger = mocker.patch("nodestream.metrics.getLogger").return_value
     handler = JsonLogMetricHandler()
-    handler.increment(Metric.RECORDS, 1)
-    handler.decrement(Metric.RECORDS, 1)
+    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
+    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    handler.tick()
     handler.stop()
     mock_logger.info.assert_called_once()
 
@@ -56,12 +85,12 @@ def test_aggregate_handler(mocker):
     handler.start()
     mock_handler1.start.assert_called_once()
     mock_handler2.start.assert_called_once()
-    handler.increment(Metric.RECORDS, 1)
-    mock_handler1.increment.assert_called_once_with(Metric.RECORDS, 1)
-    mock_handler2.increment.assert_called_once_with(Metric.RECORDS, 1)
-    handler.decrement(Metric.RECORDS, 1)
-    mock_handler1.decrement.assert_called_once_with(Metric.RECORDS, 1)
-    mock_handler2.decrement.assert_called_once_with(Metric.RECORDS, 1)
+    handler.increment(NodestreamMetricRegistry.RECORDS, 1)
+    mock_handler1.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
+    mock_handler2.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
+    handler.decrement(NodestreamMetricRegistry.RECORDS, 1)
+    mock_handler1.decrement.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
+    mock_handler2.decrement.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
     handler.stop()
     mock_handler1.stop.assert_called_once()
     mock_handler2.stop.assert_called_once()
@@ -70,8 +99,8 @@ def test_aggregate_handler(mocker):
 def test_metrics_context(mocker):
     handler = mocker.Mock()
     with Metrics.capture(handler) as metrics:
-        metrics.increment(Metric.RECORDS, 1)
-        handler.increment.assert_called_once_with(Metric.RECORDS, 1)
+        metrics.increment(NodestreamMetricRegistry.RECORDS, 1)
+        handler.increment.assert_called_once_with(NodestreamMetricRegistry.RECORDS, 1)
     handler.stop.assert_called_once()