From c34551014f2b401bd4908ab3038a8d39fd5cb930 Mon Sep 17 00:00:00 2001 From: nhnn Date: Mon, 8 Jul 2024 18:00:49 +0300 Subject: [PATCH] fix: refactor utils/prometheus.py --- selfprivacy_api/utils/prometheus.py | 56 +++++++++++++---------- tests/test_graphql/test_api_monitoring.py | 54 ++++++++++++---------- 2 files changed, 61 insertions(+), 49 deletions(-) diff --git a/selfprivacy_api/utils/prometheus.py b/selfprivacy_api/utils/prometheus.py index 4f25859..da057ee 100644 --- a/selfprivacy_api/utils/prometheus.py +++ b/selfprivacy_api/utils/prometheus.py @@ -2,42 +2,52 @@ # pylint: disable=too-few-public-methods import requests -from typing import Optional, Dict +import strawberry +from dataclasses import dataclass +from typing import Optional +from strawberry.scalars import JSON from datetime import datetime, timedelta PROMETHEUS_URL = "http://localhost:9001" -class PrometheusInfo: - """ """ - - http_response: int - output: Optional[Dict] = None - error: Optional[str] = None +@strawberry.type +@dataclass +class PrometheusQueryResult: + result_type: str + result: JSON class PrometheusQueries: @staticmethod - def _send_request(endpoint="/api/v1/query", params=None) -> PrometheusInfo: + def _send_query(query: str, start: int, end: int, step: int): try: - response = requests.get(f"{PROMETHEUS_URL}{endpoint}", params=params) - return PrometheusInfo( - http_response=response.status_code, - output=(response.json() if response.status_code == 200 else None), + response = requests.get( + f"{PROMETHEUS_URL}/api/v1/query", + params={ + "query": query, + "start": start, + "end": end, + "step": step, + }, ) - except requests.RequestException as error: - return PrometheusInfo( - http_response=500, - error=str(error), + if response.status_code != 200: + raise Exception("Prometheus returned unexpected HTTP status code") + json = response.json() + return PrometheusQueryResult( + result_type=json.result_type, result=json.result ) + except Exception as error: + raise Exception("Prometheus request failed! " + str(error)) @staticmethod def cpu_usage( start: Optional[int] = None, end: Optional[int] = None, step: int = 60, # seconds - ) -> PrometheusInfo: - """Get CPU information,. + ) -> PrometheusQueryResult: + """ + Get CPU information. Args: start (int, optional): Unix timestamp indicating the start time. @@ -55,16 +65,14 @@ class PrometheusQueries: query = '100 - (avg by (instance) (rate(node_cpu_seconds_total{mode="idle"}[5m])) * 100)' - params = {"query": query, "start": start, "end": end, "step": step} - - return PrometheusQueries._send_request(params=params) + return PrometheusQueries._send_query(query, start, end, step) @staticmethod def disk_usage( start: Optional[int] = None, end: Optional[int] = None, step: int = 60, # seconds - ) -> PrometheusInfo: + ) -> PrometheusQueryResult: """ Get disk usage information. @@ -84,6 +92,4 @@ class PrometheusQueries: query = '100 - (100 * ((node_filesystem_avail_bytes{mountpoint="/",fstype!="rootfs"} ) / (node_filesystem_size_bytes{mountpoint="/",fstype!="rootfs"}) ))' - params = {"query": query, "start": start, "end": end, "step": step} - - return PrometheusQueries._send_request(params=params) + return PrometheusQueries._send_query(query, start, end, step) diff --git a/tests/test_graphql/test_api_monitoring.py b/tests/test_graphql/test_api_monitoring.py index 8420fcd..b57434b 100644 --- a/tests/test_graphql/test_api_monitoring.py +++ b/tests/test_graphql/test_api_monitoring.py @@ -3,14 +3,10 @@ # pylint: disable=missing-function-docstring import pytest -from datetime import datetime, timezone - from tests.test_graphql.common import ( assert_empty, get_data, assert_ok, - assert_errorcode, - assert_original, ) @@ -53,14 +49,6 @@ MOCK_CPU_USAGE_RESPONSE = { } -@pytest.fixture -def mock_post(): - with patch( - "path.to.authorized_client.post", return_value=mock_response - ) as mock_method: - yield mock_method - - MOCK_DISKS_USAGE_RESPONSE = { "data": { "monitoring": { @@ -148,37 +136,54 @@ MOCK_MEMORY_USAGE_RESPONSE = { DISKS_USAGE = """ -query DefShouldSleepEnough { +query TestDisksUsage { monitoring { - disks_usage { - <ты не сделал как я просил вывод данных из Prometheus, по этому тут будет ошибка> - } + disks_usage } } """ DISKS_USAGE_WITH_OPTIONS = """ -query DefShouldSleepEnough($start: int, $end: int, $step: int) { +query TestDisksUsageWithOptions($start: int, $end: int, $step: int) { monitoring { - disks_usage(start: $start, end: $end, step: $step) { - <ты не сделал как я просил вывод данных из Prometheus, по этому тут будет ошибка> - } + disks_usage(start: $start, end: $end, step: $step) } } """ -def test_graphql_get_disks_usage(client, authorized_client): +@pytest.fixture +def mock_cpu_usage(mocker): + mock = mocker.patch( + "selfprivacy_api.utils.PrometheusQueries._sent_query", + return_value=MOCK_CPU_USAGE_RESPONSE["data"]["monitoring"]["cpu_usage"], + ) + return mock + + +@pytest.fixture +def mock_disk_usage(mocker): + mock = mocker.patch( + "selfprivacy_api.utils.PrometheusQueries._sent_query", + return_value=MOCK_DISKS_USAGE_RESPONSE["data"]["monitoring"]["disk_usage"], + ) + return mock + + +def test_graphql_get_disks_usage(client, authorized_client, mock_disk_usage): response = authorized_client.post( "/graphql", json={"query": DISKS_USAGE}, ) - data = get_data(response)["monitoring"]["disks_usage"] + data = get_data(response) assert_ok(data) + assert data["data"] == MOCK_DISKS_USAGE_RESPONSE["data"] -def test_graphql_get_disks_usage_with_options(client, authorized_client): +def test_graphql_get_disks_usage_with_options( + client, authorized_client, mock_disk_usage +): response = authorized_client.post( "/graphql", json={ @@ -191,8 +196,9 @@ def test_graphql_get_disks_usage_with_options(client, authorized_client): }, ) - data = get_data(response)["monitoring"]["disks_usage"] + data = get_data(response) assert_ok(data) + assert data["data"] == MOCK_DISKS_USAGE_RESPONSE["data"] def test_graphql_get_disks_usage_unauthorized(client):