cleveragents-core/benchmarks/cloud_resource_bench.py

"""ASV benchmarks for cloud resource handler overhead.

Measures the performance of:
- Cloud credential resolution (aws, gcp, azure)
- Cloud credential validation
- CloudResourceHandler instantiation
- CloudSandboxStrategy validation
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

try:
    from cleveragents.resource.handlers.cloud import (
        CloudResourceHandler,
        CloudSandboxStrategy,
        resolve_credentials,
        validate_credentials,
    )
except ModuleNotFoundError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
    from cleveragents.resource.handlers.cloud import (
        CloudResourceHandler,
        CloudSandboxStrategy,
        resolve_credentials,
        validate_credentials,
    )


_CLOUD_ENV_VARS = [
    "AWS_ACCESS_KEY_ID",
    "AWS_SECRET_ACCESS_KEY",
    "AWS_SESSION_TOKEN",
    "AWS_REGION",
    "AWS_PROFILE",
    "GOOGLE_APPLICATION_CREDENTIALS",
    "GCLOUD_PROJECT",
    "GCP_REGION",
    "AZURE_SUBSCRIPTION_ID",
    "AZURE_TENANT_ID",
    "AZURE_CLIENT_ID",
    "AZURE_CLIENT_SECRET",
    "AZURE_REGION",
]


def _set_aws_env() -> dict[str, str]:
    """Set AWS env vars for benchmarking."""
    saved: dict[str, str] = {}
    for var in _CLOUD_ENV_VARS:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    os.environ["AWS_ACCESS_KEY_ID"] = "AKIAIOSFODNN7EXAMPLE"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "wJalrXUtnFEMI/EXAMPLEKEY"
    os.environ["AWS_REGION"] = "us-east-1"
    return saved


def _set_gcp_env() -> dict[str, str]:
    """Set GCP env vars for benchmarking."""
    saved: dict[str, str] = {}
    for var in _CLOUD_ENV_VARS:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/path/to/sa-key.json"
    os.environ["GCLOUD_PROJECT"] = "my-gcp-project"
    return saved


def _set_azure_env() -> dict[str, str]:
    """Set Azure env vars for benchmarking."""
    saved: dict[str, str] = {}
    for var in _CLOUD_ENV_VARS:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    os.environ["AZURE_SUBSCRIPTION_ID"] = "sub-12345"
    os.environ["AZURE_TENANT_ID"] = "tenant-67890"
    os.environ["AZURE_CLIENT_ID"] = "client-abcde"
    os.environ["AZURE_CLIENT_SECRET"] = "secret-fghij"
    return saved


def _restore_env(saved: dict[str, str]) -> None:
    """Restore env vars from saved dict."""
    for var in _CLOUD_ENV_VARS:
        if var in saved:
            os.environ[var] = saved[var]
        else:
            os.environ.pop(var, None)


class CloudCredentialResolutionSuite:
    """Benchmark credential resolution for all cloud providers."""

    def setup(self) -> None:
        self.saved: dict[str, str] = {}

    def teardown(self) -> None:
        _restore_env(self.saved)

    def time_resolve_aws_credentials(self) -> None:
        """Benchmark AWS credential resolution."""
        self.saved = _set_aws_env()
        resolve_credentials("aws", {})

    def time_resolve_gcp_credentials(self) -> None:
        """Benchmark GCP credential resolution."""
        self.saved = _set_gcp_env()
        resolve_credentials("gcp", {})

    def time_resolve_azure_credentials(self) -> None:
        """Benchmark Azure credential resolution."""
        self.saved = _set_azure_env()
        resolve_credentials("azure", {})


class CloudCredentialValidationSuite:
    """Benchmark credential validation for all cloud providers."""

    def setup(self) -> None:
        self.saved: dict[str, str] = {}

    def teardown(self) -> None:
        _restore_env(self.saved)

    def time_validate_aws_credentials(self) -> None:
        """Benchmark AWS credential validation."""
        self.saved = _set_aws_env()
        resolved = resolve_credentials("aws", {})
        validate_credentials("aws", resolved)

    def time_validate_gcp_credentials(self) -> None:
        """Benchmark GCP credential validation."""
        self.saved = _set_gcp_env()
        resolved = resolve_credentials("gcp", {})
        validate_credentials("gcp", resolved)

    def time_validate_azure_credentials(self) -> None:
        """Benchmark Azure credential validation."""
        self.saved = _set_azure_env()
        resolved = resolve_credentials("azure", {})
        validate_credentials("azure", resolved)


class CloudHandlerInstantiationSuite:
    """Benchmark CloudResourceHandler and CloudSandboxStrategy creation."""

    def time_create_handler(self) -> None:
        """Benchmark handler instantiation."""
        CloudResourceHandler()

    def time_create_sandbox_strategy_aws(self) -> None:
        """Benchmark sandbox strategy creation for AWS."""
        CloudSandboxStrategy("aws")

    def time_create_sandbox_strategy_gcp(self) -> None:
        """Benchmark sandbox strategy creation for GCP."""
        CloudSandboxStrategy("gcp")

    def time_create_sandbox_strategy_azure(self) -> None:
        """Benchmark sandbox strategy creation for Azure."""
        CloudSandboxStrategy("azure")


class CloudSandboxValidationSuite:
    """Benchmark sandbox strategy validation."""

    def setup(self) -> None:
        self.aws_strategy = CloudSandboxStrategy("aws")
        self.gcp_strategy = CloudSandboxStrategy("gcp")
        self.azure_strategy = CloudSandboxStrategy("azure")
        self.saved: dict[str, str] = {}

    def teardown(self) -> None:
        _restore_env(self.saved)

    def time_validate_aws_sandbox(self) -> None:
        """Benchmark AWS sandbox validation."""
        self.saved = _set_aws_env()
        self.aws_strategy.validate({})

    def time_validate_gcp_sandbox(self) -> None:
        """Benchmark GCP sandbox validation."""
        self.saved = _set_gcp_env()
        self.gcp_strategy.validate({})

    def time_validate_azure_sandbox(self) -> None:
        """Benchmark Azure sandbox validation."""
        self.saved = _set_azure_env()
        self.azure_strategy.validate({})