GlitchTip/apps/observability/tests.py

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)