diff --git a/.vscode/launch.json b/.vscode/launch.json index b7b24a4..ebc0538 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -5,7 +5,7 @@ "version": "0.2.0", "configurations": [ { - "name": "Python: Current File", + "name": "Current File", "type": "python", "request": "launch", "program": "${file}", @@ -13,18 +13,7 @@ "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", diff --git a/.vscode/settings.json b/.vscode/settings.json index 9b38853..d7338ad 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -3,5 +3,5 @@ "tests" ], "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true + "python.testing.pytestEnabled": true, } \ No newline at end of file diff --git a/estimenergy/collectors/collector.py b/estimenergy/collectors/collector.py index 48027db..7e3da31 100644 --- a/estimenergy/collectors/collector.py +++ b/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 diff --git a/estimenergy/collectors/glow_collector.py b/estimenergy/collectors/glow_collector.py index 161d702..2e9a81d 100644 --- a/estimenergy/collectors/glow_collector.py +++ b/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}") diff --git a/estimenergy/common.py b/estimenergy/common.py index 6cef4e6..366ead7 100644 --- a/estimenergy/common.py +++ b/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, ) \ No newline at end of file diff --git a/estimenergy/const.py b/estimenergy/const.py index 5543953..9827d4a 100644 --- a/estimenergy/const.py +++ b/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 = [ diff --git a/estimenergy/helpers.py b/estimenergy/helpers.py index 31118a6..68ab0f0 100644 --- a/estimenergy/helpers.py +++ b/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 \ No newline at end of file + return 31 + + +def get_current_datetime(): + """Get current datetime. + + Returns + ------- + datetime + Current datetime. + + """ + return datetime.now() diff --git a/estimenergy/metrics/collector_metrics.py b/estimenergy/metrics/collector_metrics.py index d17bcf6..3185438 100644 --- a/estimenergy/metrics/collector_metrics.py +++ b/estimenergy/metrics/collector_metrics.py @@ -1,14 +1,17 @@ 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 } @@ -16,6 +19,8 @@ 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 \ No newline at end of file diff --git a/estimenergy/models/collector_data.py b/estimenergy/models/collector_data.py index afcb684..bab9aa5 100644 --- a/estimenergy/models/collector_data.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index f3fe528..5a5c2c1 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/conftest.py b/tests/conftest.py index 643c0de..e96661c 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/integration_tests/__init__.py b/tests/integration_tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/integration_tests/test_metrics.py b/tests/integration_tests/test_metrics.py deleted file mode 100644 index d691c3a..0000000 --- a/tests/integration_tests/test_metrics.py +++ /dev/null @@ -1,35 +0,0 @@ -import os -from httpx import AsyncClient -import pytest -from estimenergy.main import app -from fastapi.testclient import TestClient - -from estimenergy.models.collector_data import CollectorData - -client = TestClient(app) - -@pytest.mark.asyncio -async def test_metrics(client: AsyncClient): - # collector_data = await CollectorData.create( - # id=1, - # 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() - # collector_data = await CollectorData.filter(name="glow_test", id=1) - # assert collector_data is not None - - # response = client.get("/metrics") - # assert response.status_code == 200 - # assert os.environ["DB_PATH"] == "test.db" - # print(response.content) - - collector_datas = await CollectorData.filter().count() - assert collector_datas == 0 diff --git a/tests/test_app.py b/tests/test_app.py deleted file mode 100644 index daa8021..0000000 --- a/tests/test_app.py +++ /dev/null @@ -1,5 +0,0 @@ - -from estimenergy.main import app - -def test_app(): - assert app is not None diff --git a/tests/test_glow_collector.py b/tests/test_glow_collector.py new file mode 100644 index 0000000..1c94d49 --- /dev/null +++ b/tests/test_glow_collector.py @@ -0,0 +1,71 @@ + +from freezegun import freeze_time +import pytest + +from estimenergy.collectors.glow_collector import GlowCollector +from estimenergy.models.energy_data import EnergyData + + +@pytest.mark.anyio +async def test_fixtures(glow_collector, collector_data): + assert glow_collector.collector_data == collector_data + + +@pytest.mark.anyio +async def test_kwh_changed(glow_collector: GlowCollector, get_metric_value): + kwh = 0 + for _ in range(10): + kwh += 1 + await glow_collector.update_kwh(kwh) + await glow_collector.metrics.update_metrics() + value = await get_metric_value("estimenergy_day_kwh", "glow_test") + assert value == kwh + + +@freeze_time("2021-11-01") +@pytest.mark.anyio +async def test_create_energy_data(glow_collector: GlowCollector): + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is None + + await glow_collector.update_kwh(123.321) + + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is not None + assert energy_data.kwh == 123.321 + +@freeze_time("2021-11-01") +@pytest.mark.anyio +async def test_update_energy_data(glow_collector: GlowCollector): + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is None + + await glow_collector.update_kwh(123.321) + await glow_collector.update_kwh(321.123) + + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is not None + assert energy_data.kwh == 321.123 + +@freeze_time("2021-11-01") +@pytest.mark.anyio +async def test_create_energy_data_next_day(glow_collector: GlowCollector): + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is None + + await glow_collector.update_kwh(123.321) + + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is not None + assert energy_data.kwh == 123.321 + + with freeze_time("2021-11-02"): + await glow_collector.update_kwh(321.123) + + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=1).first() + assert energy_data is not None + assert energy_data.kwh == 123.321 + + energy_data = await EnergyData.filter(collector=glow_collector.collector_data, year=2021, month=11, day=2).first() + assert energy_data is not None + assert energy_data.kwh == 321.123 \ No newline at end of file diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..2ac17bf --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,143 @@ + +from freezegun import freeze_time +from httpx import AsyncClient +import pytest +from estimenergy.const import Metric, MetricPeriod, MetricType +from estimenergy.main import app +from fastapi.testclient import TestClient + +from estimenergy.models.collector_data import CollectorData +from estimenergy.models.energy_data import EnergyData + + +@pytest.mark.anyio +@freeze_time("2023-06-05") +async def test_day_kwh(client: AsyncClient, get_metric_value, create_collector_metrics): + 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() + + energy_data = await EnergyData.create( + collector_id=1, + year=2023, + month=6, + day=5, + kwh=10, + hour_created=0, + hour_updated=23, + is_completed=True, + ) + await energy_data.save() + + energy_data = await EnergyData.filter(collector=collector_data, year=2023, month=6, day=5).first() + + assert energy_data.kwh == 10 + + await create_collector_metrics(collector_data) + value = await get_metric_value( + Metric( + MetricType.ENERGY, + MetricPeriod.DAY, + is_predicted=False, + is_raw=False + ).json_key, + collector_data.name + ) + + assert value == 10 + + +@pytest.mark.anyio +@freeze_time("2023-06-05") +async def test_month_kwh(client: AsyncClient, get_metric_value, create_collector_metrics): + 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() + + for day in range(1, 6): + energy_data = await EnergyData.create( + collector_id=1, + year=2023, + month=6, + day=day, + kwh=10, + hour_created=0, + hour_updated=23, + is_completed=True, + ) + await energy_data.save() + + await create_collector_metrics(collector_data) + value = await get_metric_value( + Metric( + MetricType.ENERGY, + MetricPeriod.MONTH, + is_predicted=False, + is_raw=False + ).json_key, + collector_data.name + ) + + assert value == 50 + + +@pytest.mark.anyio +@freeze_time("2023-06-05") +async def test_year_kwh(client: AsyncClient, get_metric_value, create_collector_metrics): + 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() + + for month in range(1, 6): + for day in range(1, 6): + energy_data = await EnergyData.create( + collector_id=1, + year=2023, + month=month, + day=day, + kwh=10, + hour_created=0, + hour_updated=23, + is_completed=True, + ) + await energy_data.save() + + await create_collector_metrics(collector_data) + value = await get_metric_value( + Metric( + MetricType.ENERGY, + MetricPeriod.YEAR, + is_predicted=False, + is_raw=False + ).json_key, + collector_data.name + ) + + assert value == 250