Updates

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/testutils.py

@@ -0,0 +1,199 @@
+import copy
+
+from prometheus_client import REGISTRY
+
+METRIC_EQUALS_ERR_EXPLANATION = """
+%s%s = %s, expected %s.
+The values for %s are:
+%s"""
+
+METRIC_DIFF_ERR_EXPLANATION = """
+%s%s changed by %f, expected %f.
+Value before: %s
+Value after: %s
+"""
+
+METRIC_COMPARE_ERR_EXPLANATION = """
+The change in value of %s%s didn't match the predicate.
+Value before: %s
+Value after: %s
+"""
+
+METRIC_DIFF_ERR_NONE_EXPLANATION = """
+%s%s was None after.
+Value before: %s
+Value after: %s
+"""
+
+
+"""A collection of utilities that make it easier to write test cases
+that interact with metrics.
+"""
+
+
+def assert_metric_equal(expected_value, metric_name, registry=REGISTRY, **labels):
+    """Asserts that metric_name{**labels} == expected_value."""
+    value = get_metric(metric_name, registry=registry, **labels)
+    assert_err = METRIC_EQUALS_ERR_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        value,
+        expected_value,
+        metric_name,
+        format_vector(get_metrics_vector(metric_name)),
+    )
+    assert expected_value == value, assert_err
+
+
+def assert_metric_diff(frozen_registry, expected_diff, metric_name, registry=REGISTRY, **labels):
+    """Asserts that metric_name{**labels} changed by expected_diff between
+    the frozen registry and now. A frozen registry can be obtained
+    by calling save_registry, typically at the beginning of a test
+    case.
+    """
+    saved_value = get_metric_from_frozen_registry(metric_name, frozen_registry, **labels)
+    current_value = get_metric(metric_name, registry=registry, **labels)
+    assert current_value is not None, METRIC_DIFF_ERR_NONE_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        saved_value,
+        current_value,
+    )
+    diff = current_value - (saved_value or 0.0)
+    assert_err = METRIC_DIFF_ERR_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        diff,
+        expected_diff,
+        saved_value,
+        current_value,
+    )
+    assert expected_diff == diff, assert_err
+
+
+def assert_metric_no_diff(frozen_registry, expected_diff, metric_name, registry=REGISTRY, **labels):
+    """Asserts that metric_name{**labels} isn't changed by expected_diff between
+    the frozen registry and now. A frozen registry can be obtained
+    by calling save_registry, typically at the beginning of a test
+    case.
+    """
+    saved_value = get_metric_from_frozen_registry(metric_name, frozen_registry, **labels)
+    current_value = get_metric(metric_name, registry=registry, **labels)
+    assert current_value is not None, METRIC_DIFF_ERR_NONE_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        saved_value,
+        current_value,
+    )
+    diff = current_value - (saved_value or 0.0)
+    assert_err = METRIC_DIFF_ERR_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        diff,
+        expected_diff,
+        saved_value,
+        current_value,
+    )
+    assert expected_diff != diff, assert_err
+
+
+def assert_metric_not_equal(expected_value, metric_name, registry=REGISTRY, **labels):
+    """Asserts that metric_name{**labels} == expected_value."""
+    value = get_metric(metric_name, registry=registry, **labels)
+    assert_err = METRIC_EQUALS_ERR_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        value,
+        expected_value,
+        metric_name,
+        format_vector(get_metrics_vector(metric_name)),
+    )
+    assert expected_value != value, assert_err
+
+
+def assert_metric_compare(frozen_registry, predicate, metric_name, registry=REGISTRY, **labels):
+    """Asserts that metric_name{**labels} changed according to a provided
+    predicate function between the frozen registry and now. A
+    frozen registry can be obtained by calling save_registry,
+    typically at the beginning of a test case.
+    """
+    saved_value = get_metric_from_frozen_registry(metric_name, frozen_registry, **labels)
+    current_value = get_metric(metric_name, registry=registry, **labels)
+    assert current_value is not None, METRIC_DIFF_ERR_NONE_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        saved_value,
+        current_value,
+    )
+    assert predicate(saved_value, current_value) is True, METRIC_COMPARE_ERR_EXPLANATION % (
+        metric_name,
+        format_labels(labels),
+        saved_value,
+        current_value,
+    )
+
+
+def save_registry(registry=REGISTRY):
+    """Freezes a registry. This lets a user test changes to a metric
+    instead of testing the absolute value. A typical use case looks like:
+
+        registry = save_registry()
+        doStuff()
+        assert_metric_diff(registry, 1, 'stuff_done_total')
+    """
+    return copy.deepcopy(list(registry.collect()))
+
+
+def get_metric(metric_name, registry=REGISTRY, **labels):
+    """Gets a single metric."""
+    return get_metric_from_frozen_registry(metric_name, registry.collect(), **labels)
+
+
+def get_metrics_vector(metric_name, registry=REGISTRY):
+    """Returns the values for all labels of a given metric.
+
+    The result is returned as a list of (labels, value) tuples,
+    where `labels` is a dict.
+
+    This is quite a hack since it relies on the internal
+    representation of the prometheus_client, and it should
+    probably be provided as a function there instead.
+    """
+    return get_metric_vector_from_frozen_registry(metric_name, registry.collect())
+
+
+def get_metric_vector_from_frozen_registry(metric_name, frozen_registry):
+    """Like get_metrics_vector, but from a frozen registry."""
+    output = []
+    for metric in frozen_registry:
+        for sample in metric.samples:
+            if sample[0] == metric_name:
+                output.append((sample[1], sample[2]))
+    return output
+
+
+def get_metric_from_frozen_registry(metric_name, frozen_registry, **labels):
+    """Gets a single metric from a frozen registry."""
+    for metric in frozen_registry:
+        for sample in metric.samples:
+            if sample[0] == metric_name and sample[1] == labels:
+                return sample[2]
+
+
+def format_labels(labels):
+    """Format a set of labels to Prometheus representation.
+
+    In:
+      {'method': 'GET', 'port': '80'}
+
+    Out:
+      '{method="GET",port="80"}'
+    """
+    return "{{{}}}".format(",".join([f'{k}="{v}"' for k, v in labels.items()]))
+
+
+def format_vector(vector):
+    """Formats a list of (labels, value) where labels is a dict into a
+    human-readable representation.
+    """
+    return "\n".join([f"{format_labels(labels)} = {value}" for labels, value in vector])