Add test for the COgniteMetrics Config

Author: Hmnt39

tests/test_unstable/test_base.py

@@ -10,7 +10,7 @@
 
 from cognite.client import CogniteClient
 from cognite.client.exceptions import CogniteNotFoundError
-from cognite.extractorutils.metrics import PrometheusPusher
+from cognite.extractorutils.metrics import CognitePusher, PrometheusPusher
 from cognite.extractorutils.statestore.watermark import LocalStateStore, RawStateStore
 from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io
 from cognite.extractorutils.unstable.configuration.models import (
@@ -24,6 +24,7 @@
     RawStateStoreConfig,
     StateStoreConfig,
     TimeIntervalConfig,
+    _CogniteMetricsConfig,
     _PromServerConfig,
     _PushGatewayConfig,
 )
@@ -245,31 +246,52 @@ def test_raw_state_store_integration(
         assert new_extractor.state_store.get_state("my-test-id") == (1, 5)
 
 
-def test_extractor_with_metrics(
-    connection_config: ConnectionConfig,
-    override_level: str | None = None,
-) -> None:
-    app_config = TestConfig(
-        parameter_one=1,
-        parameter_two="a",
-    )
-    metrics_config = MetricsConfig(
-        server=_PromServerConfig(
-            host="localhost",
-            port=9090,
-        ),
-        cognite=None,
-        push_gateways=[
-            _PushGatewayConfig(
-                host="localhost",
-                job_name="test-job",
-                username=None,
-                password=None,
-                clear_after=None,
-                push_interval=TimeIntervalConfig("30s"),
-            )
-        ],
-    )
+@pytest.mark.parametrize("metrics_type", ["prometheus", "cognite"])
+def test_extractor_with_metrics_pushers(connection_config: ConnectionConfig, metrics_type: str) -> None:
+    override_level = "INFO"
+    app_config = TestConfig(parameter_one=1, parameter_two="a")
+    call_count = {"count": 0}
+
+    if metrics_type == "prometheus":
+        metrics_config = MetricsConfig(
+            server=_PromServerConfig(host="localhost", port=9090),
+            cognite=None,
+            push_gateways=[
+                _PushGatewayConfig(
+                    host="localhost",
+                    job_name="test-job",
+                    username=None,
+                    password=None,
+                    clear_after=None,
+                    push_interval=TimeIntervalConfig("30s"),
+                )
+            ],
+        )
+        pusher_cls = PrometheusPusher
+
+        def counting_push(self: PrometheusPusher) -> None:
+            call_count["count"] += 1
+            return original_push(self)
+
+        original_push = pusher_cls._push_to_server
+    else:
+        metrics_config = MetricsConfig(
+            server=None,
+            cognite=_CogniteMetricsConfig(
+                external_id_prefix="extractor_test",
+                asset_name="Extractor Test Metrics",
+                asset_external_id="extractor_testcognite_assets",
+                data_set=None,
+            ),
+            push_gateways=None,
+        )
+        pusher_cls = CognitePusher
+
+        def counting_push(self: CognitePusher) -> None:
+            call_count["count"] += 1
+            return None
+
+        original_push = pusher_cls._push_to_server
 
     full_config = FullConfig(
         connection_config=connection_config,
@@ -282,27 +304,16 @@ def test_extractor_with_metrics(
     extractor = TestExtractor(full_config, worker, metrics=TestMetrics)
     assert isinstance(extractor._metrics, TestMetrics) or extractor._metrics == TestMetrics
 
-    # The metrics instance should be a singleton
-    another_extractor = TestExtractor(full_config, worker, metrics=TestMetrics)
-    assert another_extractor._metrics is extractor._metrics
-    assert isinstance(another_extractor._metrics, TestMetrics) or another_extractor._metrics == TestMetrics
-
-    call_count = {"count": 0}
-    original_push = PrometheusPusher._push_to_server
-
-    def counting_push(self: "PrometheusPusher") -> None:
-        call_count["count"] += 1
-        return original_push(self)
-
-    PrometheusPusher._push_to_server = counting_push
-
-    try:
-        with extractor:
-            for pusher in extractor.metrics_push_manager.pushers:
-                assert pusher.thread is not None
-                assert pusher.thread.is_alive()
-    finally:
-        PrometheusPusher._push_to_server = original_push
+    with contextlib.ExitStack() as stack:
+        stack.enter_context(contextlib.suppress(Exception))
+        pusher_cls._push_to_server = counting_push
+        try:
+            with extractor:
+                for pusher in extractor.metrics_push_manager.pushers:
+                    assert pusher.thread is not None
+                    assert pusher.thread.is_alive()
+        finally:
+            pusher_cls._push_to_server = original_push
     assert call_count["count"] > 0