fix: tests

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

src/bentoml/metrics.py

@@ -20,9 +20,9 @@ def __dir__() -> list[str]:
     return dir(metrics_client.prometheus_client)
 
 
-class LazyObject:
+class _LazyMetric:
     def __init__(self, attr: t.Any):
-        self.__bentoml_metrics_attr__ = attr
+        self._attr = attr
         self._proxy = None
         self._initialized = False
         self._args: tuple[t.Any, ...] = ()
@@ -45,16 +45,9 @@ def _load_proxy(
         self,
         metrics_client: PrometheusClient = Provide[BentoMLContainer.metrics_client],
     ) -> None:
-        parent = (
-            metrics_client
-            if self.__bentoml_metrics_attr__ in dir(metrics_client)
-            else metrics_client.prometheus_client
-        )
-        self._proxy = getattr(parent, self.__bentoml_metrics_attr__)(
-            *self._args, **self._kwargs
-        )
+        self._proxy = getattr(metrics_client, self._attr)(*self._args, **self._kwargs)
         self._initialized = True
 
 
 def __getattr__(item: t.Any):
-    return LazyObject(item)
+    return _LazyMetric(item)

tests/unit/test_metrics.py

@@ -1,22 +1,27 @@
 from __future__ import annotations
 
-import sys
-import typing as t
-
-import pytest
+from typing import TYPE_CHECKING
 
 import bentoml
 
+if TYPE_CHECKING:
+    from bentoml._internal.server.metrics.prometheus import PrometheusClient
+
+
+def test_metrics_initialization():
+    o = bentoml.metrics.Gauge(name="test_metrics", documentation="test")
+    assert isinstance(o, bentoml.metrics._LazyMetric)
+    o = bentoml.metrics.Histogram(name="test_metrics", documentation="test")
+    assert isinstance(o, bentoml.metrics._LazyMetric)
+    o = bentoml.metrics.Counter(name="test_metrics", documentation="test")
+    assert isinstance(o, bentoml.metrics._LazyMetric)
+    o = bentoml.metrics.Summary(name="test_metrics", documentation="test")
+    assert isinstance(o, bentoml.metrics._LazyMetric)
+
 
-@pytest.mark.parametrize(
-    "metrics_type",
-    filter(
-        lambda x: isinstance(x, bentoml.metrics.metrics.MetricWrapperBase),
-        dir(bentoml.metrics),
-    ),
-)
-def test_metrics_initialization(
-    metrics_type: t.Type[bentoml.metrics.metrics.MetricWrapperBase],
-):
-    _ = metrics_type(name="test_metrics", documentation="test")
-    assert "prometheus_client" not in sys.modules
+def test_metrics_type(prom_client: PrometheusClient):
+    o = bentoml.metrics.Counter(name="test_metrics", documentation="test")
+    assert o._attr == "Counter"
+    assert o._proxy is None
+    o.inc()
+    assert isinstance(o._proxy, prom_client.prometheus_client.Counter)