118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
from collections.abc import Iterable, Mapping
|
|
from typing import Optional
|
|
|
|
from django.urls import reverse
|
|
from model_bakery import baker
|
|
from prometheus_client import Metric
|
|
from prometheus_client.parser import text_string_to_metric_families
|
|
from rest_framework import status
|
|
from rest_framework.test import APIClient, APITestCase
|
|
|
|
from .metrics import clear_metrics_cache, organizations_metric, projects_metric
|
|
|
|
|
|
def get_sample_value(
|
|
metric_families: Iterable[Metric],
|
|
metric_name: str,
|
|
metric_type: str,
|
|
labels: Mapping[str, str],
|
|
) -> Optional[float]:
|
|
for metric_family in metric_families:
|
|
if metric_family.name != metric_name or metric_family.type != metric_type:
|
|
continue
|
|
for metric in metric_family.samples:
|
|
if metric[1] != labels:
|
|
continue
|
|
return metric.value
|
|
return None
|
|
|
|
|
|
def parse_prometheus_text(text: str) -> list[Metric]:
|
|
parser = text_string_to_metric_families(text)
|
|
return list(parser)
|
|
|
|
|
|
class ObservabilityAPITestCase(APITestCase):
|
|
def setUp(self):
|
|
self.client = APIClient()
|
|
self.user = baker.make("users.user", is_staff=True)
|
|
self.client.force_login(self.user)
|
|
self.url = reverse("prometheus-django-metrics")
|
|
|
|
def _get_metrics(self) -> list[Metric]:
|
|
resp = self.client.get(self.url)
|
|
return parse_prometheus_text(resp.content.decode("utf-8"))
|
|
|
|
def test_get_metrics_and_cache(self):
|
|
clear_metrics_cache()
|
|
with self.assertNumQueries(2):
|
|
resp = self.client.get(self.url)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
|
|
with self.assertNumQueries(1):
|
|
resp = self.client.get(self.url)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
|
|
def test_org_metric(self):
|
|
before_orgs_metric = get_sample_value(
|
|
self._get_metrics(),
|
|
organizations_metric._name,
|
|
organizations_metric._type,
|
|
{},
|
|
)
|
|
|
|
# create new org; must invalidate the cache
|
|
org = baker.make("organizations_ext.Organization")
|
|
metrics = self._get_metrics()
|
|
orgs_metric = get_sample_value(
|
|
metrics, organizations_metric._name, organizations_metric._type, {}
|
|
)
|
|
self.assertEqual(orgs_metric, before_orgs_metric + 1)
|
|
|
|
# delete org and test again
|
|
org.delete()
|
|
metrics = self._get_metrics()
|
|
orgs_metric = get_sample_value(
|
|
metrics, organizations_metric._name, organizations_metric._type, {}
|
|
)
|
|
self.assertEqual(orgs_metric, before_orgs_metric)
|
|
|
|
def test_project_metric(self):
|
|
# create new org
|
|
org = baker.make("organizations_ext.Organization")
|
|
|
|
# no projects yet
|
|
metrics = self._get_metrics()
|
|
projs_metric = get_sample_value(
|
|
metrics,
|
|
projects_metric._name,
|
|
projects_metric._type,
|
|
{"organization": org.slug},
|
|
)
|
|
self.assertEqual(projs_metric, 0)
|
|
|
|
# create new project
|
|
proj = baker.make("projects.Project", organization=org)
|
|
# test
|
|
metrics = self._get_metrics()
|
|
projs_metric = get_sample_value(
|
|
metrics,
|
|
projects_metric._name,
|
|
projects_metric._type,
|
|
{"organization": org.slug},
|
|
)
|
|
self.assertEqual(projs_metric, 1)
|
|
|
|
# delete project
|
|
proj.force_delete()
|
|
|
|
# test
|
|
metrics = self._get_metrics()
|
|
projs_metric = get_sample_value(
|
|
metrics,
|
|
projects_metric._name,
|
|
projects_metric._type,
|
|
{"organization": org.slug},
|
|
)
|
|
self.assertEqual(projs_metric, 0)
|