fix: refactor utils/prometheus.py

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)

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):