infrahub-server 1.2.0b1__py3-none-any.whl → 1.2.1__py3-none-any.whl

infrahub_testcontainers/performance_test.py

@@ -0,0 +1,166 @@
+import hashlib
+import json
+from datetime import UTC, datetime
+from types import TracebackType
+from typing import Any
+
+import httpx
+import pytest
+from typing_extensions import Self
+
+from .constants import PERFORMANCE_TEST_KIND, PERFORMANCE_TEST_VERSION
+from .helpers import InfrahubDockerCompose
+from .host import get_system_stats
+from .models import (
+    ContextUnit,
+    InfrahubActiveMeasurementItem,
+    InfrahubMeasurementItem,
+    InfrahubResultContext,
+    MeasurementDefinition,
+)
+
+
+class InfrahubPerformanceTest:
+    context: dict[str, InfrahubResultContext]
+    measurements: list[InfrahubMeasurementItem]
+    active_measurements: InfrahubActiveMeasurementItem | None = None
+
+    def __init__(self, results_url: str) -> None:
+        self.name: str | None = None
+        self.infrahub_version: str | None = None
+        self.context = {}
+        self.measurements = []
+        self.metrics = {}
+        self.host = get_system_stats()
+        self.env_vars = {}
+        self.project_name = ""
+        self.test_info = {}
+        self.start_time = datetime.now(UTC)
+        self.end_time: datetime | None = None
+        self.results_url = results_url
+        self.scraper_endpoint = ""
+        self.initialized = False
+
+    def initialize(
+        self,
+        name: str,
+        compose: InfrahubDockerCompose | None = None,
+        client: Any | None = None,  # noqa: ANN401
+    ) -> None:
+        self.name = name
+        if client:
+            self.infrahub_version = client.get_version()
+        if compose:
+            self.extract_compose_information(compose)
+
+        self.initialized = True
+
+    def finalize(self, session: pytest.Session) -> None:
+        if self.initialized:
+            self.end_time = datetime.now(UTC)
+            self.extract_test_session_information(session)
+            self.send_results()
+
+    def extract_compose_information(self, compose: InfrahubDockerCompose) -> None:
+        self.env_vars = compose.env_vars
+        self.project_name = compose.project_name
+        self.scraper_endpoint = (
+            f"http://127.0.0.1:{compose.get_service_host_and_port(service_name='scraper')[1]}/api/v1/export"
+        )
+
+    def extract_test_session_information(self, session: pytest.Session) -> None:
+        self.test_info = {
+            "summary": {
+                "exitstatus": session.exitstatus,
+                "nbr_tests": session.testscollected,
+                "nbr_errors": session.testsfailed,
+            }
+        }
+
+    def add_context(self, name: str, value: float | str = 0, unit: ContextUnit = ContextUnit.COUNT) -> None:
+        self.context[name] = InfrahubResultContext(name=name, value=value, unit=unit)
+
+    def add_measurement(
+        self,
+        definition: MeasurementDefinition,
+        value: float | str,
+        **kwargs: str | float,
+    ) -> None:
+        self.measurements.append(
+            InfrahubMeasurementItem(
+                name=definition.name,
+                value=value,
+                unit=definition.unit,
+                context=kwargs or {},
+            )
+        )
+
+    def start_measurement(self, definition: MeasurementDefinition, **kwargs: str | float) -> Self:
+        self.active_measurements = InfrahubActiveMeasurementItem(definition=definition, context=kwargs or {})
+        return self
+
+    def fetch_metrics(self) -> None:
+        with httpx.Client() as client:
+            # Get Infrahub metrics
+            response = client.post(
+                url=self.scraper_endpoint,
+                content='match[]={__name__=~"infrahub.*"}',
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+            self.metrics = [json.loads(line) for line in response.text.splitlines()]
+
+            # Get system metrics, filter by docker project name
+            response = client.post(
+                url=self.scraper_endpoint,
+                content=f'match[]={{__name__=~"container.*", container_label_com_docker_compose_project="{self.project_name}"}}',
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+            self.metrics += [json.loads(line) for line in response.text.splitlines()]
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if not exc_type and self.active_measurements:
+            self.add_measurement(
+                definition=self.active_measurements.definition,
+                value=(datetime.now(UTC) - self.active_measurements.start_time).total_seconds() * 1000,
+                context=self.active_measurements.context,
+            )
+
+        # Don't record the measurement if the test failed
+        self.active_measurements = None
+
+    def _get_payload(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "infrahub_version": self.infrahub_version,
+            "start_time": self.start_time.isoformat(),
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+            "duration": (self.end_time - self.start_time).total_seconds() * 1000 if self.end_time else None,
+            "context": [item.model_dump() for item in self.context.values()],
+            "measurements": [item.model_dump() for item in self.measurements],
+            "metrics": self.metrics,
+            "host": self.host,
+            "env_vars": self.env_vars,
+            "test_info": self.test_info,
+        }
+
+    def send_results(self) -> None:
+        data = self._get_payload()
+
+        payload = {
+            "kind": PERFORMANCE_TEST_KIND,
+            "payload_format": PERFORMANCE_TEST_VERSION,
+            "data": data,
+            "checksum": hashlib.sha256(json.dumps(data).encode()).hexdigest(),
+        }
+
+        with httpx.Client() as client:
+            response = client.post(self.results_url, json=payload)
+            response.raise_for_status()

infrahub_testcontainers/plugin.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from .performance_test import InfrahubPerformanceTest
+
+if TYPE_CHECKING:
+    import _pytest
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    group = parser.getgroup("infrahub-performance-test")
+
+    group.addoption(
+        "--performance-result-address",
+        action="store",
+        dest="infrahub_performance_test_result_address",
+        default="https://webhook.site/25802839-4b34-433e-9dc4-59623cd73c80",
+        metavar="INFRAHUB_PERFORMANCE_TEST_RESULT_ADDRESS",
+        help="Address to send the results of the performance test (default: %(default)s)",
+    )
+    group.addoption(
+        "--performance-backup",
+        action="store",
+        dest="infrahub_performance_backup",
+        default="neo4j_database.db",
+        metavar="INFRAHUB_PERFORMANCE_BACKUP",
+        help="Name of the backup to use (default: %(default)s)",
+    )
+    group.addoption(
+        "--performance-backup-location",
+        action="store",
+        dest="infrahub_performance_backup_location",
+        metavar="INFRAHUB_PERFORMANCE_BACKUP_LOCATION",
+        help="Location of the backup to use (default: %(default)s)",
+    )
+    group.addoption(
+        "--performance-create-backup",
+        action="store_true",
+        default=False,
+        dest="infrahub_performance_create_backup",
+        help="Generate a backup of the database (default: %(default)s)",
+    )
+    group.addoption(
+        "--performance-use-backup",
+        action="store_true",
+        default=False,
+        dest="infrahub_performance_use_backup",
+        help="Use a backup of the database and skip all tests that creates data (default: %(default)s)",
+    )
+    group.addoption(
+        "--performance-report",
+        action="store_true",
+        dest="infrahub_performance_report",
+        default=False,
+        help="Display performance report at the end of the test run",
+    )
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    use_backup_skip = pytest.mark.skip(reason="load data from backup")
+    no_create_backup_skip = pytest.mark.skip(reason="no need to create a backup")
+    no_load_backup_skip = pytest.mark.skip(reason="no need to load a backup")
+
+    use_backup = config.getoption("infrahub_performance_use_backup")
+    create_backup = config.getoption("infrahub_performance_create_backup")
+
+    for item in items:
+        has_create_backup_marker = bool(
+            [True for marker in item.own_markers if marker.name == "performance_create_backup"]
+        )
+        has_load_backup_marker = bool([True for marker in item.own_markers if marker.name == "performance_load_backup"])
+        has_load_data_marker = bool([True for marker in item.own_markers if marker.name == "performance_load_data"])
+
+        if has_create_backup_marker and not create_backup:
+            item.add_marker(no_create_backup_skip)
+        if has_load_backup_marker and not use_backup:
+            item.add_marker(no_load_backup_skip)
+        if has_load_data_marker and use_backup:
+            item.add_marker(use_backup_skip)
+
+
+def pytest_sessionstart(session: pytest.Session) -> None:
+    session.infrahub_performance_test = InfrahubPerformanceTest(
+        results_url=session.config.getoption("infrahub_performance_test_result_address")
+    )
+
+
+def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:  # noqa: ARG001
+    """whole test run finishes."""
+    if not session.config.getoption("infrahub_performance_report"):
+        return
+
+    session.infrahub_performance_test.finalize(session=session)
+
+
+def pytest_runtest_teardown(item: pytest.Item, nextitem: pytest.Item | None) -> None:  # noqa: ARG001
+    """Fetch metrics at each test teardown because there's no better hook...
+    pytest_sessionfinish() is executed after fixtures has been finalized and pytest_fixture_post_finalizer() is too late"""
+    if not item.config.getoption("infrahub_performance_report"):
+        return
+
+    item.session.infrahub_performance_test.fetch_metrics()
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    if config.getoption("infrahub_performance_use_backup") and config.getoption("infrahub_performance_create_backup"):
+        raise pytest.UsageError("--performance-use-backup and --performance-create-backup are mutually exclusive")
+
+    config.addinivalue_line("markers", "performance_create_backup: Create a backup of the database")
+    config.addinivalue_line("markers", "performance_load_backup: Load the backup of the database")
+    config.addinivalue_line("markers", "performance_load_data: Load initial data into the database")
+
+
+def pytest_terminal_summary(
+    terminalreporter: _pytest.terminal.TerminalReporter,
+    exitstatus: int,  # noqa: ARG001
+    config: pytest.Config,
+) -> None:
+    if not config.getoption("infrahub_performance_report"):
+        return
+
+    performance_test = terminalreporter._session.infrahub_performance_test
+
+    report = [
+        f"{measurement.name}: {measurement.value} {measurement.unit.value}"
+        for measurement in performance_test.measurements
+    ]
+    terminalreporter.write("\n" + "\n".join(report) + "\n")
+
+
+@pytest.fixture(scope="session")
+def perf_test(request: pytest.FixtureRequest) -> InfrahubPerformanceTest:
+    return request.session.infrahub_performance_test

infrahub_testcontainers/prometheus.yml

@@ -0,0 +1,30 @@
+---
+global:
+  scrape_interval: 15s
+
+scrape_configs:
+  - job_name: 'infrahub-server'
+    dns_sd_configs:
+      - names:
+          - server
+        type: A
+        port: 8000
+
+  - job_name: 'infrahub-worker'
+    dns_sd_configs:
+      - names:
+          - task-worker
+        type: A
+        port: 8000
+
+  - job_name: 'cadvisor'
+    dns_sd_configs:
+      - names:
+          - cadvisor
+        type: A
+        port: 8080
+    metric_relabel_configs:
+      - source_labels: [container_label_namespace]
+        target_label: namespace
+      - source_labels: [container_label_pod]
+        target_label: pod

infrahub/message_bus/messages/event_branch_create.py

@@ -1,11 +0,0 @@
-from pydantic import Field
-
-from infrahub.message_bus import InfrahubMessage
-
-
-class EventBranchCreate(InfrahubMessage):
-    """Sent a new branch is created."""
-
-    branch: str = Field(..., description="The branch that was created")
-    branch_id: str = Field(..., description="The unique ID of the branch")
-    sync_with_git: bool = Field(..., description="Indicates if Infrahub should extend this branch to git.")

infrahub/message_bus/messages/event_branch_delete.py

@@ -1,11 +0,0 @@
-from pydantic import Field
-
-from infrahub.message_bus import InfrahubMessage
-
-
-class EventBranchDelete(InfrahubMessage):
-    """Sent when a branch has been deleted."""
-
-    branch: str = Field(..., description="The branch that was deleted")
-    branch_id: str = Field(..., description="The unique ID of the branch")
-    sync_with_git: bool = Field(..., description="Indicates if the branch was extended to Git")

infrahub/message_bus/messages/event_branch_rebased.py

@@ -1,9 +0,0 @@
-from pydantic import Field
-
-from infrahub.message_bus import InfrahubMessage
-
-
-class EventBranchRebased(InfrahubMessage):
-    """Sent when a branch has been rebased."""
-
-    branch: str = Field(..., description="The branch that was rebased")

infrahub/message_bus/messages/event_node_mutated.py

@@ -1,15 +0,0 @@
-from typing import Any
-
-from pydantic import Field
-
-from infrahub.message_bus import InfrahubMessage
-
-
-class EventNodeMutated(InfrahubMessage):
-    """Sent when a node has been mutated"""
-
-    branch: str = Field(..., description="The branch that was created")
-    kind: str = Field(..., description="The type of object modified")
-    node_id: str = Field(..., description="The ID of the mutated node")
-    action: str = Field(..., description="The action taken on the node")
-    data: dict[str, Any] = Field(..., description="Data on modified object")

infrahub/message_bus/messages/event_schema_update.py

@@ -1,9 +0,0 @@
-from pydantic import Field
-
-from infrahub.message_bus import InfrahubMessage
-
-
-class EventSchemaUpdate(InfrahubMessage):
-    """Sent when the schema on a branch has been updated."""
-
-    branch: str = Field(..., description="The branch where the update occurred")

infrahub/message_bus/operations/event/node.py

@@ -1,20 +0,0 @@
-from infrahub.message_bus import messages
-from infrahub.services import InfrahubServices
-
-
-async def mutated(
-    message: messages.EventNodeMutated,
-    service: InfrahubServices,
-) -> None:
-    """Event posted when a node is mutated"""
-    # Note for now this is only kept to facilitate publishing to other queues in the future
-    # This operation doesn't have a flow defined to avoid having the workers need to register
-    # the results in Prefect when they don't actually do anything aside from add noise.
-    service.log.debug(
-        "Mutation on node",
-        branch=message.branch,
-        node_id=message.node_id,
-        action=message.action,
-        kind=message.kind,
-        data=message.data,
-    )

infrahub/message_bus/operations/event/schema.py

@@ -1,17 +0,0 @@
-from prefect import flow
-
-from infrahub.log import get_logger
-from infrahub.message_bus import messages
-from infrahub.services import InfrahubServices
-
-log = get_logger()
-
-
-@flow(name="event-schema-update")
-async def update(message: messages.EventSchemaUpdate, service: InfrahubServices) -> None:
-    log.info("run_message", branch=message.branch)
-
-    msg = messages.RefreshRegistryBranches()
-
-    msg.assign_meta(parent=message)
-    await service.message_bus.send(message=msg)

infrahub/webhook/constants.py

@@ -1 +0,0 @@
-AUTOMATION_NAME_RUN = "Trigger-webhook-execution"