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

{infrahub_server-1.2.0rc0.dist-info → infrahub_server-1.2.1.dist-info}/entry_points.txt

@@ -9,4 +9,5 @@ infrahubasync=infrahub.workers.infrahub_async
 
 [pytest11]
 pytest-infrahub=infrahub_sdk.pytest_plugin.plugin
+pytest-infrahub-performance-test=infrahub_testcontainers.plugin
 

infrahub_testcontainers/constants.py

@@ -0,0 +1,2 @@
+PERFORMANCE_TEST_KIND: str = "end-2-end"
+PERFORMANCE_TEST_VERSION: str = "v20250216"

infrahub_testcontainers/container.py

@@ -1,10 +1,12 @@
 import os
+import shutil
 import uuid
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import cached_property
 from pathlib import Path
 
 from testcontainers.compose import DockerCompose
+from testcontainers.core.exceptions import ContainerIsNotRunning
 from typing_extensions import Self
 
 from infrahub_testcontainers import __version__ as infrahub_version
@@ -19,14 +21,21 @@ class ContainerService:
 INFRAHUB_SERVICES: dict[str, ContainerService] = {
     "server": ContainerService(container="infrahub-server-lb", port=8000),
     "task-manager": ContainerService(container="task-manager", port=4200),
+    "database": ContainerService(container="database", port=7687),
+    "scraper": ContainerService(container="scraper", port=8428),
+    "cadvisor": ContainerService(container="cadvisor", port=8080),
 }
 
 PROJECT_ENV_VARIABLES: dict[str, str] = {
+    "NEO4J_DOCKER_IMAGE": "neo4j:5.20.0-community",
+    "MESSAGE_QUEUE_DOCKER_IMAGE": "rabbitmq:3.13.7-management",
+    "CACHE_DOCKER_IMAGE": "redis:7.2.4",
     "INFRAHUB_TESTING_DOCKER_IMAGE": "registry.opsmill.io/opsmill/infrahub",
+    "INFRAHUB_TESTING_DOCKER_ENTRYPOINT": f"gunicorn --config backend/infrahub/serve/gunicorn_config.py -w {os.environ.get('INFRAHUB_TESTING_WEB_CONCURRENCY', 4)} --logger-class infrahub.serve.log.GunicornLogger infrahub.server:app",  # noqa: E501
     "INFRAHUB_TESTING_IMAGE_VERSION": infrahub_version,
     "INFRAHUB_TESTING_PRODUCTION": "false",
     "INFRAHUB_TESTING_DB_ADDRESS": "database",
-    "INFRAHUB_TESTING_LOG_LEVEL": "DEBUG",
+    "INFRAHUB_TESTING_LOG_LEVEL": "INFO",
     "INFRAHUB_TESTING_GIT_REPOSITORIES_DIRECTORY": "/opt/infrahub/git",
     "INFRAHUB_TESTING_API_TOKEN": "44af444d-3b26-410d-9546-b758657e026c",
     "INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN": "06438eb2-8019-4776-878c-0941b1f1d1ec",
@@ -37,6 +46,7 @@ PROJECT_ENV_VARIABLES: dict[str, str] = {
     "INFRAHUB_TESTING_BROKER_ADDRESS": "message-queue",
     "INFRAHUB_TESTING_CACHE_ADDRESS": "cache",
     "INFRAHUB_TESTING_WORKFLOW_ADDRESS": "task-manager",
+    "INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE": "infrahubasync",
     "INFRAHUB_TESTING_TIMEOUT": "60",
     "INFRAHUB_TESTING_PREFECT_API": "http://task-manager:4200/api",
     "INFRAHUB_TESTING_LOCAL_REMOTE_GIT_DIRECTORY": "repos",
@@ -52,6 +62,7 @@ PROJECT_ENV_VARIABLES: dict[str, str] = {
 @dataclass
 class InfrahubDockerCompose(DockerCompose):
     project_name: str | None = None
+    env_vars: dict[str, str] = field(default_factory=dict)
 
     @classmethod
     def init(cls, directory: Path | None = None, version: str | None = None) -> Self:
@@ -65,43 +76,111 @@ class InfrahubDockerCompose(DockerCompose):
         if version == "local" and infrahub_image_version:
             version = infrahub_image_version
 
-        cls.create_docker_file(directory=directory)
-        cls.create_env_file(directory=directory, version=version)
+        compose = cls(project_name=cls.generate_project_name(), context=directory)
+        compose.create_docker_file(directory=directory)
+        compose.create_env_file(directory=directory, version=version)
 
-        return cls(project_name=cls.generate_project_name(), context=directory)
+        return compose
+
+    def get_env_var(self, key: str) -> str:
+        if not self.env_vars:
+            raise ValueError("env_vars hasn't been initialized yet")
+        if key not in self.env_vars:
+            raise ValueError(f"{key} is not set in the environment variables")
+        return self.env_vars[key]
+
+    @property
+    def use_neo4j_enterprise(self) -> bool:
+        return "enterprise" in self.get_env_var("NEO4J_DOCKER_IMAGE")
+
+    @property
+    def internal_backup_dir(self) -> Path:
+        return Path(self.get_env_var("INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY"))
+
+    @property
+    def external_backup_dir(self) -> Path:
+        return Path(self.context) / Path(self.get_env_var("INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY"))
 
     @classmethod
     def generate_project_name(cls) -> str:
         project_id = str(uuid.uuid4())[:8]
         return f"infrahub-test-{project_id}"
 
-    @classmethod
-    def create_docker_file(cls, directory: Path) -> Path:
+    def create_docker_file(self, directory: Path) -> Path:
         current_directory = Path(__file__).resolve().parent
         compose_file = current_directory / "docker-compose.test.yml"
 
         test_compose_file = directory / "docker-compose.yml"
         test_compose_file.write_bytes(compose_file.read_bytes())
 
-        haproxy_config_file = current_directory / "haproxy.cfg"
+        for file in ["haproxy.cfg", "prometheus.yml"]:
+            config_file = current_directory / file
 
-        test_haproxy_config_file = directory / "haproxy.cfg"
-        test_haproxy_config_file.write_bytes(haproxy_config_file.read_bytes())
+            test_config_file = directory / file
+            test_config_file.write_bytes(config_file.read_bytes())
 
         return test_compose_file
 
-    @classmethod
-    def create_env_file(cls, directory: Path, version: str) -> Path:
+    def create_env_file(self, directory: Path, version: str) -> Path:
         env_file = directory / ".env"
 
         PROJECT_ENV_VARIABLES.update({"INFRAHUB_TESTING_IMAGE_VERSION": version})
+        if os.environ.get("INFRAHUB_TESTING_ENTERPRISE"):
+            PROJECT_ENV_VARIABLES.update(
+                {
+                    "INFRAHUB_TESTING_DOCKER_IMAGE": "registry.opsmill.io/opsmill/infrahub-enterprise",
+                    "INFRAHUB_TESTING_DOCKER_ENTRYPOINT": f"gunicorn --config community/backend/infrahub/serve/gunicorn_config.py -w {os.environ.get('INFRAHUB_TESTING_WEB_CONCURRENCY', 4)} --logger-class infrahub.serve.log.GunicornLogger infrahub_enterprise.server:app",  # noqa: E501
+                    "INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE": "infrahubentasync",
+                    "NEO4J_DOCKER_IMAGE": "neo4j:5.20.0-enterprise",
+                }
+            )
 
         with env_file.open(mode="w", encoding="utf-8") as file:
             for key, value in PROJECT_ENV_VARIABLES.items():
                 env_var_value = os.environ.get(key, value)
                 file.write(f"{key}={env_var_value}\n")
+                self.env_vars[key] = env_var_value
+
         return env_file.absolute()
 
+    def restart(self) -> None:
+        """
+        Restart the docker compose environment.
+
+        TODO Would be good to contribute this upstream
+        """
+        cmd = self.compose_command_property[:]
+        cmd += ["restart"]
+
+        if self.services:
+            cmd.extend(self.services)
+        self._run_command(cmd=cmd)
+
+    def start_container(self, service_name: str) -> None:
+        """
+        Starts a specific service of the docker compose environment.
+
+        TODO Would be good to contribute this upstream
+        """
+        base_cmd = self.compose_command_property or []
+
+        # pull means running a separate command before starting
+        if self.pull:
+            pull_cmd = [*base_cmd, "pull", service_name]
+            self._run_command(cmd=pull_cmd)
+
+        up_cmd = [*base_cmd, "up"]
+
+        # build means modifying the up command
+        if self.wait:
+            up_cmd.append("--wait")
+        else:
+            # we run in detached mode instead of blocking
+            up_cmd.append("--detach")
+
+        up_cmd.append(service_name)
+        self._run_command(cmd=up_cmd)
+
     # TODO would be good to the support for project_name upstream
     @cached_property
     def compose_command_property(self) -> list[str]:
@@ -120,3 +199,69 @@ class InfrahubDockerCompose(DockerCompose):
             service_name: int(self.get_service_port(service_name=service_data.container, port=service_data.port) or 0)
             for service_name, service_data in INFRAHUB_SERVICES.items()
         }
+
+    def database_create_backup(self, backup_name: str = "neo4j_database.backup", dest_dir: Path | None = None) -> None:
+        assert self.use_neo4j_enterprise
+
+        self.exec_in_container(
+            command=[
+                "neo4j-admin",
+                "database",
+                "backup",
+                "--compress=false",
+                "--to-path",
+                str(self.internal_backup_dir),
+            ],
+            service_name="database",
+        )
+
+        if dest_dir:
+            backup_files = list(self.external_backup_dir.glob("*.backup"))
+            if not backup_files:
+                raise FileNotFoundError(f"No .backup files found in {self.external_backup_dir}")
+
+            backup_file = backup_files[0]
+            shutil.copy(
+                backup_file,
+                dest_dir / backup_name,
+            )
+
+    def database_restore_backup(self, backup_file: Path) -> None:
+        assert self.use_neo4j_enterprise
+
+        shutil.copy(
+            str(backup_file),
+            str(self.external_backup_dir / backup_file.name),
+        )
+        service_name = "database"
+
+        # Ensure the database container is running otherwise start it
+        try:
+            self.get_container(service_name=service_name)
+        except ContainerIsNotRunning:
+            self.start_container(service_name=service_name)
+
+        self.exec_in_container(
+            command=["cypher-shell", "-u", "neo4j", "-p", "admin", "STOP DATABASE neo4j;"],
+            service_name=service_name,
+        )
+
+        self.exec_in_container(
+            command=[
+                "neo4j-admin",
+                "database",
+                "restore",
+                "--overwrite-destination",
+                "--from-path",
+                str(self.internal_backup_dir / backup_file.name),
+            ],
+            service_name=service_name,
+        )
+
+        self.exec_in_container(
+            command=["cypher-shell", "-d", "system", "-u", "neo4j", "-p", "admin", "START DATABASE neo4j;"],
+            service_name=service_name,
+        )
+
+        self.stop(down=False)
+        self.start()

infrahub_testcontainers/docker-compose.test.yml

@@ -102,11 +102,7 @@ services:
       mode: replicated
       replicas: ${INFRAHUB_TESTING_API_SERVER_COUNT}
     image: "${INFRAHUB_TESTING_DOCKER_IMAGE}:${INFRAHUB_TESTING_IMAGE_VERSION}"
-    command: >
-      gunicorn --config backend/infrahub/serve/gunicorn_config.py
-      -w ${INFRAHUB_TESTING_WEB_CONCURRENCY}
-      --logger-class infrahub.serve.log.GunicornLogger
-      infrahub.server:app
+    command: ${INFRAHUB_TESTING_DOCKER_ENTRYPOINT}
     environment:
       INFRAHUB_PRODUCTION: ${INFRAHUB_TESTING_PRODUCTION}
       INFRAHUB_LOG_LEVEL: ${INFRAHUB_TESTING_LOG_LEVEL:-INFO}
@@ -114,6 +110,7 @@ services:
       INFRAHUB_CACHE_ADDRESS: ${INFRAHUB_TESTING_CACHE_ADDRESS}
       INFRAHUB_DB_ADDRESS: ${INFRAHUB_TESTING_DB_ADDRESS}
       INFRAHUB_WORKFLOW_ADDRESS: ${INFRAHUB_TESTING_WORKFLOW_ADDRESS}
+      INFRAHUB_WORKFLOW_DEFAULT_WORKER_TYPE: ${INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE}
       INFRAHUB_INITIAL_ADMIN_TOKEN: ${INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN}
       INFRAHUB_INITIAL_AGENT_TOKEN: ${INFRAHUB_TESTING_INITIAL_AGENT_TOKEN}
       INFRAHUB_SECURITY_SECRET_KEY: ${INFRAHUB_TESTING_SECURITY_SECRET_KEY}
@@ -142,7 +139,7 @@ services:
       mode: replicated
       replicas: ${INFRAHUB_TESTING_TASK_WORKER_COUNT}
     image: "${INFRAHUB_TESTING_DOCKER_IMAGE}:${INFRAHUB_TESTING_IMAGE_VERSION}"
-    command: prefect worker start --type infrahubasync --pool infrahub-worker --with-healthcheck
+    command: prefect worker start --type ${INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE} --pool infrahub-worker --with-healthcheck
     environment:
       INFRAHUB_PRODUCTION: ${INFRAHUB_TESTING_PRODUCTION}
       INFRAHUB_LOG_LEVEL: ${INFRAHUB_TESTING_LOG_LEVEL}
@@ -163,8 +160,36 @@ services:
       - "./${INFRAHUB_TESTING_LOCAL_REMOTE_GIT_DIRECTORY}:${INFRAHUB_TESTING_INTERNAL_REMOTE_GIT_DIRECTORY}"
     tty: true
 
+  cadvisor:
+    image: "${CADVISOR_DOCKER_IMAGE:-gcr.io/cadvisor/cadvisor:v0.51.0}"
+    command:
+      - -disable_root_cgroup_stats=true
+      - -docker_only=true
+      - -store_container_labels=false
+      - -whitelisted_container_labels=com.docker.compose.project
+    privileged: true
+    volumes:
+      - /:/rootfs:ro
+      - /var/run:/var/run:ro
+      - /sys:/sys:ro
+      - /var/lib/docker:/var/lib/docker:ro
+      - /dev/disk/:/dev/disk:ro
+    ports:
+      - "${INFRAHUB_TESTING_CADVISOR_PORT:-0}:8080"
+
+  scraper:
+    image: "${SCRAPER_DOCKER_IMAGE:-victoriametrics/victoria-metrics:v1.110.0}"
+    volumes:
+      - vmdata:/victoria-metrics-data
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml:ro
+    command:
+      - "--promscrape.config=/etc/prometheus/prometheus.yml"
+    ports:
+      - ${INFRAHUB_TESTING_SCRAPER_PORT:-0}:8428
+
 volumes:
   database_data:
   database_logs:
   storage_data:
   workflow_db:
+  vmdata:

infrahub_testcontainers/helpers.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 import subprocess  # noqa: S404
 from pathlib import Path
 
@@ -15,7 +14,8 @@ class TestInfrahubDocker:
     def infrahub_version(self) -> str:
         return infrahub_version
 
-    def execute_ctl_run(self, address: str, script: str) -> str:
+    @staticmethod
+    def execute_ctl_run(address: str, script: str) -> str:
         env = os.environ.copy()
         env["INFRAHUB_ADDRESS"] = address
         env["INFRAHUB_API_TOKEN"] = PROJECT_ENV_VARIABLES["INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN"]
@@ -25,6 +25,17 @@ class TestInfrahubDocker:
         )
         return result.stdout
 
+    @staticmethod
+    def execute_command(command: str, address: str, concurrent_execution: int = 10) -> subprocess.CompletedProcess[str]:
+        env = os.environ.copy()
+        env["INFRAHUB_ADDRESS"] = address
+        env["INFRAHUB_API_TOKEN"] = PROJECT_ENV_VARIABLES["INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN"]
+        env["INFRAHUB_MAX_CONCURRENT_EXECUTION"] = f"{concurrent_execution}"
+        result = subprocess.run(  # noqa: S602
+            command, shell=True, capture_output=True, text=True, env=env, check=False
+        )
+        return result
+
     @pytest.fixture(scope="class")
     def tmp_directory(self, tmpdir_factory: pytest.TempdirFactory) -> Path:
         directory = Path(str(tmpdir_factory.getbasetemp().strpath))
@@ -65,7 +76,11 @@ class TestInfrahubDocker:
 
         request.addfinalizer(cleanup)
 
-        infrahub_compose.start()
+        try:
+            infrahub_compose.start()
+        except Exception as exc:
+            stdout, stderr = infrahub_compose.get_logs()
+            raise Exception(f"Failed to start docker compose:\nStdout:\n{stdout}\nStderr:\n{stderr}") from exc
 
         return infrahub_compose.get_services_port()
 
@@ -76,73 +91,3 @@ class TestInfrahubDocker:
     @pytest.fixture(scope="class")
     def task_manager_port(self, infrahub_app: dict[str, int]) -> int:
         return infrahub_app["task-manager"]
-
-    def backup_database(self, request: pytest.FixtureRequest, dest_dir: Path | None = None) -> None:
-        assert "enterprise" in os.environ.get("NEO4J_DOCKER_IMAGE", "")
-
-        backup_dir: Path = request.getfixturevalue("remote_backups_dir")
-        infrahub_compose: InfrahubDockerCompose = request.getfixturevalue("infrahub_compose")
-
-        infrahub_compose.exec_in_container(
-            command=[
-                "neo4j-admin",
-                "database",
-                "backup",
-                "--to-path",
-                os.environ.get(
-                    "INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY",
-                    PROJECT_ENV_VARIABLES["INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY"],
-                ),
-            ],
-            service_name="database",
-        )
-
-        if dest_dir:
-            shutil.copytree(
-                str(backup_dir),
-                str(dest_dir),
-            )
-
-    def restore_database(self, request: pytest.FixtureRequest, backup_file: Path) -> None:
-        assert "enterprise" in os.environ.get("NEO4J_DOCKER_IMAGE", "")
-
-        backup_dir: Path = request.getfixturevalue("remote_backups_dir")
-        infrahub_compose: InfrahubDockerCompose = request.getfixturevalue("infrahub_compose")
-
-        shutil.copy(
-            str(backup_file),
-            str(backup_dir / backup_file.name),
-        )
-
-        infrahub_compose.exec_in_container(
-            command=["cypher-shell", "-u", "neo4j", "-p", "admin", "STOP DATABASE neo4j;"],
-            service_name="database",
-        )
-
-        infrahub_compose.exec_in_container(
-            command=[
-                "neo4j-admin",
-                "database",
-                "restore",
-                "--overwrite-destination",
-                "--from-path",
-                str(
-                    Path(
-                        os.environ.get(
-                            "INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY",
-                            PROJECT_ENV_VARIABLES["INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY"],
-                        )
-                    )
-                    / backup_file.name
-                ),
-            ],
-            service_name="database",
-        )
-
-        infrahub_compose.exec_in_container(
-            command=["cypher-shell", "-d", "system", "-u", "neo4j", "-p", "admin", "START DATABASE neo4j;"],
-            service_name="database",
-        )
-
-        infrahub_compose.stop(down=False)
-        infrahub_compose.start()

infrahub_testcontainers/host.py

@@ -0,0 +1,41 @@
+import platform
+from typing import Any
+
+import psutil
+
+
+def get_system_stats() -> dict[str, Any]:
+    """
+    Gather and return key system statistics about CPU, memory and platform
+    Returns a dictionary with system information
+    """
+    # CPU information
+    cpu_freq = psutil.cpu_freq()
+    cpu_stats = {
+        "cpu_count_physical": psutil.cpu_count(logical=False),
+        "cpu_count_logical": psutil.cpu_count(logical=True),
+        "cpu_freq_current": float(f"{cpu_freq.current:.2f}") if cpu_freq else None,  # MHz
+        "cpu_freq_min": float(f"{cpu_freq.min:.2f}") if cpu_freq else None,  # MHz
+        "cpu_freq_max": float(f"{cpu_freq.max:.2f}") if cpu_freq else None,  # MHz
+        "cpu_percent": psutil.cpu_percent(interval=1, percpu=False),
+    }
+
+    # Memory information
+    memory = psutil.virtual_memory()
+    memory_stats = {
+        "total_memory": float(f"{memory.total / (1024**3):.2f}"),  # GB
+        "available_memory": float(f"{memory.available / (1024**3):.2f}"),  # GB
+        "memory_percent_used": memory.percent,
+        "memory_used": float(f"{memory.used / (1024**3):.2f}"),  # GB
+    }
+
+    # Platform information
+    platform_stats = {
+        "system": platform.system(),
+        "platform_release": platform.release(),
+        "platform_version": platform.version(),
+        "architecture": platform.machine(),
+        "processor": platform.processor(),
+    }
+
+    return {"cpu": cpu_stats, "memory": memory_stats, "platform": platform_stats}

infrahub_testcontainers/measurements.py

@@ -0,0 +1,93 @@
+from .models import ContextUnit, MeasurementDefinition
+
+NODE_QUERY_TIME = MeasurementDefinition(
+    name="node_query",
+    description="Query some nodes",
+    dimensions=["kind"],
+    unit=ContextUnit.TIME,
+)
+
+NODE_CREATE_TIME = MeasurementDefinition(
+    name="node_mutation_create",
+    description="Create a new node",
+    dimensions=["kind"],
+    unit=ContextUnit.TIME,
+)
+
+NODE_UPDATE_TIME = MeasurementDefinition(
+    name="node_mutation_update",
+    description="Update an existing node",
+    dimensions=["kind"],
+    unit=ContextUnit.TIME,
+)
+
+NODE_DELETE_TIME = MeasurementDefinition(
+    name="node_mutation_delete",
+    description="Delete a node",
+    dimensions=["kind"],
+    unit=ContextUnit.TIME,
+)
+
+SCHEMA_INITIAL_LOAD_TIME = MeasurementDefinition(
+    name="schema_initial_load",
+    description="Load the initial schema",
+    dimensions=["branch"],
+    unit=ContextUnit.TIME,
+)
+
+SCHEMA_UPDATE_TIME = MeasurementDefinition(
+    name="schema_update",
+    description="Update the schema",
+    dimensions=["branch"],
+    unit=ContextUnit.TIME,
+)
+
+DATABASE_SIZE = MeasurementDefinition(
+    name="database_size",
+    description="Size of the database",
+    unit=ContextUnit.DISK,
+)
+
+DIFF_CREATE_TIME = MeasurementDefinition(
+    name="diff_create",
+    description="Create a new diff",
+    unit=ContextUnit.TIME,
+)
+
+DIFF_APPLY_TIME = MeasurementDefinition(
+    name="diff_update",
+    description="Update an existing diff",
+    unit=ContextUnit.TIME,
+)
+
+BRANCH_CREATE_TIME = MeasurementDefinition(
+    name="branch_create",
+    description="Create a new branch",
+    unit=ContextUnit.TIME,
+)
+
+BRANCH_MERGE_TIME = MeasurementDefinition(
+    name="branch_merge",
+    description="Merge a branch",
+    unit=ContextUnit.TIME,
+)
+
+BRANCH_REBASE_TIME = MeasurementDefinition(
+    name="branch_rebase",
+    description="Rebase a branch",
+    unit=ContextUnit.TIME,
+)
+
+SCRIPT_EXECUTION_TIME = MeasurementDefinition(
+    name="script_execution",
+    description="Execute a script",
+    dimensions=["name"],
+    unit=ContextUnit.TIME,
+)
+
+GENERATOR_EXECUTION_TIME = MeasurementDefinition(
+    name="generator_execution",
+    description="Execute a generator",
+    dimensions=["name"],
+    unit=ContextUnit.TIME,
+)

infrahub_testcontainers/models.py

@@ -0,0 +1,38 @@
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class ContextUnit(str, Enum):
+    COUNT = "count"
+    TIME = "msec"  # time in milliseconds
+    MEMORY = "memory"  # memory in bytes
+    DISK = "disk"  # disk in bytes
+
+
+class MeasurementDefinition(BaseModel):
+    name: str
+    description: str
+    dimensions: list[str] = Field(default_factory=dict)
+    unit: ContextUnit
+
+
+class InfrahubResultContext(BaseModel):
+    name: str
+    value: int | float | str
+    unit: ContextUnit
+
+
+class InfrahubActiveMeasurementItem(BaseModel):
+    definition: MeasurementDefinition
+    start_time: datetime = datetime.now(UTC)
+    context: dict[str, Any] = Field(default_factory=dict)
+
+
+class InfrahubMeasurementItem(BaseModel):
+    name: str
+    value: int | float | str
+    unit: ContextUnit
+    context: dict[str, Any] = Field(default_factory=dict)