feat: add tests

.vscode/launch.json

@@ -5,26 +5,15 @@
     "version": "0.2.0",
     "configurations": [
         {
-            "name": "Python: Current File",
+            "name": "Current File",
             "type": "python",
             "request": "launch",
             "program": "${file}",
             "console": "integratedTerminal",
             "justMyCode": true
         },
         {
-            "name": "Python: Module",
-            "type": "python",
-            "request": "launch",
-            "module": "poetry",
-            "args": [
-                "run",
-                "api"
-            ],
-            "justMyCode": true
-        },
-        {
-            "name": "Debug EstimEnergy",
+            "name": "API",
             "type": "python",
             "request": "launch",
             "module": "poetry",

.vscode/settings.json

@@ -3,5 +3,5 @@
         "tests"
     ],
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true
+    "python.testing.pytestEnabled": true,
 }

estimenergy/collectors/collector.py

@@ -11,4 +11,6 @@ def __init__(self):
     async def start(self):
         pass
 
-
+    @abstractmethod
+    async def update_kwh(self, kwh: float):
+        pass

estimenergy/collectors/glow_collector.py

@@ -16,6 +16,7 @@
 from estimenergy.collectors import Collector
 from estimenergy.models import CollectorData, EnergyData
 from estimenergy.metrics import CollectorMetrics
+from estimenergy.helpers import get_current_datetime
 
 
 class GlowCollector(Collector):
@@ -50,6 +51,9 @@ async def start(self):
         await self.__try_login()
         await self.reconnect_logic.start()
 
+    async def update_kwh(self, kwh: float):
+        return await self.__on_kwh_changed(kwh)
+
     async def __try_login(self):
         try:
             await self.api.connect(login=True)
@@ -83,7 +87,7 @@ def __state_changed(self, state: EntityState):
         loop.create_task(self.__on_kwh_changed(current_kwh))
 
     async def __on_kwh_changed(self, current_kwh: float):
-        date = datetime.datetime.now()
+        date = get_current_datetime()
 
         self.logger.info(f"Current KWh: {current_kwh}")
 

estimenergy/common.py

@@ -9,6 +9,7 @@
 load_dotenv()
 
 settings = Settings()
+metric_registry = CollectorRegistry(auto_describe=True)
 instrumentator = Instrumentator(
     should_group_status_codes=False,
     should_ignore_untemplated=True,
@@ -18,4 +19,5 @@
     env_var_name="ENABLE_METRICS",
     inprogress_name="inprogress",
     inprogress_labels=True,
+    registry=metric_registry,
 )

estimenergy/const.py

@@ -36,11 +36,12 @@ def json_key(self) -> str:
     def friendly_name(self) -> str:
         return f"{self.metric_period.value[1]} {self.metric_type.value[1]} {'(Predicted)' if self.is_predicted else ''} {'(Raw)' if self.is_raw else ''}"
 
-    def create_gauge(self) -> Gauge:
+    def create_gauge(self, registry) -> Gauge:
         return Gauge(
             f"{self.json_key}",
             f"EstimEnergy {self.friendly_name}",
             ["name", "id"],
+            registry=registry
         )
 
 METRICS = [

estimenergy/helpers.py

@@ -1,4 +1,6 @@
 
+from datetime import datetime
+
 
 def get_days_in_month(month: int, year: int):
     """Return the number of days in a given month and year."""
@@ -11,4 +13,16 @@ def get_days_in_month(month: int, year: int):
     elif month in {4, 6, 9, 11}:
         return 30
     else:
-        return 31
+        return 31
+
+
+def get_current_datetime():
+    """Get current datetime.
+
+    Returns
+    -------
+    datetime
+        Current datetime.
+
+    """
+    return datetime.now()

estimenergy/metrics/collector_metrics.py

@@ -1,21 +1,26 @@
 
 from prometheus_client.registry import Collector as Metrics
 from estimenergy.const import METRICS
+from estimenergy.helpers import get_current_datetime
 
 from estimenergy.models.collector_data import CollectorData
+from estimenergy.common import metric_registry
+
 
 class CollectorMetrics(Metrics):
     def __init__(self, collector: CollectorData):
         self.collector = collector
         self.metrics = {
-            metric: metric.create_gauge()
+            metric: metric.create_gauge(registry=metric_registry)
             for metric in METRICS
         }
 
     async def collect(self):
         return self.metrics.values()
 
     async def update_metrics(self):
-        data = await self.collector.get_metrics()
+        date = get_current_datetime()
+        data = await self.collector.get_metrics(date)
         for metric in METRICS:
             self.metrics[metric].labels(name=self.collector.name, id=self.collector.id).set(data[metric.json_key])
+        pass

estimenergy/models/collector_data.py

@@ -25,13 +25,13 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.logger = logging.getLogger("energy_collector").getChild(self.name)
 
-    async def get_metrics(self, date: datetime.date = datetime.datetime.now()):
+    async def get_metrics(self, date: datetime.date):
         return {
             metric.json_key: await self.get_metric(metric, date)
             for metric in METRICS
         }
 
-    async def get_metric(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_metric(self, metric: Metric, date: datetime.date):
         if metric.metric_type == MetricType.ENERGY:
             return await self.get_energy(metric, date)
 
@@ -46,7 +46,7 @@ async def get_metric(self, metric: Metric, date: datetime.date = datetime.dateti
 
         raise ValueError(f"Unknown metric type {metric.metric_type}")
 
-    async def get_energy(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_energy(self, metric: Metric, date: datetime.date):
         energy_datas = await self.get_energy_datas(metric.metric_period, date)
 
         if not metric.is_predicted:
@@ -93,7 +93,7 @@ async def get_energy(self, metric: Metric, date: datetime.date = datetime.dateti
 
             return kwh_total / (12 - missing_months) * 12
 
-    async def get_cost(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_cost(self, metric: Metric, date: datetime.date):
         energy = await self.get_energy(metric, date)
         kwh_cost = energy * self.cost_per_kwh
 
@@ -105,20 +105,20 @@ async def get_cost(self, metric: Metric, date: datetime.date = datetime.datetime
 
         return kwh_cost + self.base_cost_per_month * 12
 
-    async def get_cost_difference(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_cost_difference(self, metric: Metric, date: datetime.date):
         cost = await self.get_cost(metric, date)
         payment = await self.get_payment(metric, date)
 
         return payment - cost 
 
-    async def get_accuracy(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_accuracy(self, metric: Metric, date: datetime.date):
         energy_datas = await self.get_energy_datas(metric.metric_period, date)
         day_count = await self.get_day_count(metric.metric_period, date)
         accuracy_total = sum(energy_data.accuracy for energy_data in energy_datas) / day_count
 
         return accuracy_total
 
-    async def get_payment(self, metric: Metric, date: datetime.date = datetime.datetime.now()):
+    async def get_payment(self, metric: Metric, date: datetime.date):
         if metric.metric_period == MetricPeriod.DAY:
             return self.payment_per_month / get_days_in_month(date.month, date.year)
 
@@ -131,7 +131,7 @@ async def get_energy_datas(
             self, 
             metric_period: 
             MetricPeriod, 
-            date: datetime.date = datetime.datetime.now()
+            date: datetime.date
         ) -> list[EnergyData]:
         if metric_period == MetricPeriod.DAY:
             return await EnergyData.filter(collector=self, year=date.year, month=date.month, day=date.day)
@@ -143,8 +143,8 @@ async def get_energy_datas(
             return await EnergyData.filter(
                 Q(collector=self) & 
                 (
-                    Q(year=date.year - 1, month__gte=self.billing_month) | 
-                    Q(year=date.year, month__lt=self.billing_month)
+                    Q(year=date.year - 1, month__lt=self.billing_month) | 
+                    Q(year=date.year, month__gte=self.billing_month)
                 )
             )
 
@@ -153,7 +153,7 @@ async def get_energy_datas(
 
         raise ValueError(f"Unknown metric period {metric_period}")
 
-    async def get_day_count(self, metric_period: MetricPeriod, date: datetime.date = datetime.datetime.now()) -> int:
+    async def get_day_count(self, metric_period: MetricPeriod, date: datetime.date) -> int:
         if metric_period == MetricPeriod.DAY:
             return 1
 

pyproject.toml

@@ -41,6 +41,8 @@ httpx = "~0.23.3"
 PyYAML = "~6.0"
 poetry = "~1.4.0"
 pytest-asyncio = "^0.20.3"
+pytest-mock = "^3.10.0"
+freezegun = "^1.2.2"
 
 [tool.poetry.group.build.dependencies]
 poetry-dynamic-versioning = "~0.21.4"

tests/conftest.py

@@ -1,48 +1,82 @@
+import prometheus_client
 import pytest
 import pytest
 from httpx import AsyncClient
 from tortoise import Tortoise
+from prometheus_client.parser import text_string_to_metric_families
+import estimenergy
+from estimenergy.collectors.glow_collector import GlowCollector
 
 from estimenergy.main import app
+from estimenergy.common import metric_registry
+from estimenergy.models.collector_data import CollectorData
 
 DB_URL = "sqlite://:memory:"
 
-
-async def init_db(db_url, create_db: bool = False, schemas: bool = False) -> None:
-    """Initial database connection"""
-    await Tortoise.init(
-        db_url=db_url, modules={"models": ["estimenergy.models"]}, _create_db=create_db
-    )
-    if create_db:
-        print(f"Database created! {db_url = }")
-    if schemas:
-        await Tortoise.generate_schemas()
-        print("Success to generate schemas")
-
-
-async def init(db_url: str = DB_URL):
-    await init_db(db_url, True, True)
-
-
 @pytest.fixture(scope="session")
 def anyio_backend():
     return "asyncio"
 
-
 @pytest.fixture(scope="session")
 async def client():
     async with AsyncClient(app=app, base_url="http://test") as client:
-        print("Client is ready")
         yield client
 
-
-@pytest.fixture(scope="session", autouse=True)
+@pytest.fixture(scope="function", autouse=True)
 async def initialize_tests():
-    await init()
+    collectors = list(metric_registry._collector_to_names.keys())
+    for collector in collectors:
+        metric_registry.unregister(collector)
+    await Tortoise.init(
+        db_url=DB_URL, modules={"models": ["estimenergy.models"]}, _create_db=True
+    )
+    await Tortoise.generate_schemas()
     yield
     await Tortoise._drop_databases()
 
-# @pytest.fixture(scope="function", autouse=True)
-# def test_config(monkeypatch):
-#     monkeypatch.setenv("DB_PATH", "test.db")
-#     monkeypatch.setenv("CONFIG_PATH", "test.config.yml")
+@pytest.fixture(scope="function")
+async def create_collector_metrics():
+    async def create_collector_metrics(collector_data: CollectorData):
+        from estimenergy.metrics import CollectorMetrics
+        collector_metrics = CollectorMetrics(collector_data)
+        await collector_metrics.update_metrics()
+        return collector_metrics
+
+    return create_collector_metrics
+
+@pytest.fixture(scope="function")
+async def get_metric_value(client: AsyncClient):
+    async def get_metric_value(metric_name: str, collector_name: str):
+        response = await client.get("/metrics")
+        assert response.status_code == 200
+
+        families = list(text_string_to_metric_families(response.text))
+        for family in families:
+            if family.name == metric_name:
+                for sample in family.samples:
+                    if sample.labels["name"] == collector_name:
+                        return sample.value
+
+        return None
+
+    return get_metric_value
+
+@pytest.fixture(scope="function")
+async def collector_data():
+    collector_data = await CollectorData.create(
+        name="glow_test",
+        host="0.0.0.0",
+        port=0,
+        password="",
+        cost_per_kwh=1,
+        base_cost_per_month=1,
+        payment_per_month=100,
+        billing_month=1,
+        min_accuracy=0
+    )
+    await collector_data.save()
+    return collector_data
+
+@pytest.fixture(scope="function")
+async def glow_collector(collector_data):
+    return GlowCollector(collector_data)