cleveragents-core/robot/helper_cloud_resources.py

"""Helper utilities for cloud resource Robot smoke tests.

Each command prints a single ``<command>-ok`` token on success so the
calling Robot test can assert on ``stdout``.
"""

from __future__ import annotations

import contextlib
import os
import sys
from typing import Any
from unittest.mock import MagicMock

from cleveragents.domain.models.core.resource import (
    PhysVirt,
    Resource,
    ResourceCapabilities,
)
from cleveragents.resource.handlers.cloud import (
    CloudResourceHandler,
    resolve_credentials,
    validate_credentials,
)
from cleveragents.resource.handlers.cloud_aws import CloudSandboxStrategy
from cleveragents.resource.handlers.protocol import ResourceHandler

_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 _clear_cloud_env() -> dict[str, str]:
    """Remove all cloud env vars and return saved values."""
    saved: dict[str, str] = {}
    for var in _CLOUD_ENV_VARS:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    return saved


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


def _make_resource(
    provider: str,
    properties: dict[str, Any] | None = None,
) -> Resource:
    """Create a test cloud resource."""
    return Resource(
        resource_id="01KJ5C5TPMP8GGX3QC83E2MAQS",
        resource_type_name=provider,
        classification=PhysVirt.PHYSICAL,
        location=None,
        properties=properties or {},
        capabilities=ResourceCapabilities(
            readable=True,
            writable=True,
            sandboxable=False,
            checkpointable=False,
        ),
    )


def _protocol_check() -> None:
    """Verify CloudResourceHandler satisfies the ResourceHandler protocol."""
    handler = CloudResourceHandler()
    assert isinstance(handler, ResourceHandler), (
        "CloudResourceHandler not a ResourceHandler"
    )
    print("protocol-check-ok")


def _aws_resolve() -> None:
    """Verify AWS credential resolution from env vars."""
    saved = _clear_cloud_env()
    try:
        os.environ["AWS_ACCESS_KEY_ID"] = "AKIAIOSFODNN7EXAMPLE"
        os.environ["AWS_SECRET_ACCESS_KEY"] = "wJalrXUtnFEMI/EXAMPLEKEY"
        os.environ["AWS_REGION"] = "us-east-1"
        resolved = resolve_credentials("aws", {})
        assert resolved["access-key-id"] is not None
        assert resolved["secret-access-key"] is not None
        assert resolved["region"] == "us-east-1"
        errors = validate_credentials("aws", resolved)
        assert len(errors) == 0, f"Unexpected errors: {errors}"
        print("aws-resolve-ok")
    finally:
        _restore_cloud_env(saved)


def _gcp_resolve() -> None:
    """Verify GCP credential resolution from env vars."""
    saved = _clear_cloud_env()
    try:
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/path/to/sa-key.json"
        os.environ["GCLOUD_PROJECT"] = "my-gcp-project"
        resolved = resolve_credentials("gcp", {})
        assert resolved["service-account-json-path"] is not None
        assert resolved["project-id"] == "my-gcp-project"
        errors = validate_credentials("gcp", resolved)
        assert len(errors) == 0, f"Unexpected errors: {errors}"
        print("gcp-resolve-ok")
    finally:
        _restore_cloud_env(saved)


def _azure_resolve() -> None:
    """Verify Azure credential resolution from env vars."""
    saved = _clear_cloud_env()
    try:
        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"
        resolved = resolve_credentials("azure", {})
        assert resolved["subscription-id"] is not None
        assert resolved["tenant-id"] is not None
        assert resolved["client-id"] is not None
        assert resolved["client-secret"] is not None
        errors = validate_credentials("azure", resolved)
        assert len(errors) == 0, f"Unexpected errors: {errors}"
        print("azure-resolve-ok")
    finally:
        _restore_cloud_env(saved)


def _stub_resolve() -> None:
    """Verify cloud handler raises ImportError or succeeds for AWS."""
    saved = _clear_cloud_env()
    try:
        os.environ["AWS_ACCESS_KEY_ID"] = "AKIAIOSFODNN7EXAMPLE"
        os.environ["AWS_SECRET_ACCESS_KEY"] = "wJalrXUtnFEMI/EXAMPLEKEY"
        handler = CloudResourceHandler()
        resource = _make_resource("aws")
        mock_manager = MagicMock()
        try:
            handler.resolve(
                resource=resource,
                plan_id="PLAN_ROBOT",
                slot_name="cloud-slot",
                sandbox_manager=mock_manager,
            )
            # boto3 is installed and resolve succeeded
            print("stub-resolve-ok")
        except (NotImplementedError, ImportError):
            # boto3 not installed or not yet implemented
            print("stub-resolve-ok")
    finally:
        _restore_cloud_env(saved)


def _missing_creds() -> None:
    """Verify cloud handler raises ValueError for missing credentials."""
    saved = _clear_cloud_env()
    try:
        handler = CloudResourceHandler()
        resource = _make_resource("aws")
        mock_manager = MagicMock()
        try:
            handler.resolve(
                resource=resource,
                plan_id="PLAN_ROBOT",
                slot_name="cloud-slot",
                sandbox_manager=mock_manager,
            )
            print("missing-creds-FAIL: no error raised")
        except ValueError:
            print("missing-creds-ok")
        except NotImplementedError:
            print("missing-creds-FAIL: NotImplementedError instead of ValueError")
    finally:
        _restore_cloud_env(saved)


def _sandbox_stubs() -> None:
    """Verify sandbox strategy methods raise NotImplementedError or ImportError."""
    for provider in ("gcp", "azure"):
        strategy = CloudSandboxStrategy(provider)
        for method_name in ("create", "commit", "rollback"):
            method = getattr(strategy, method_name)
            try:
                method("res-test", "plan-test")
                print(f"sandbox-stubs-FAIL: {provider}.{method_name} no error")
                return
            except NotImplementedError:
                pass
    # AWS raises ImportError (boto3 not installed) or succeeds (boto3 installed)
    aws_strategy = CloudSandboxStrategy("aws")
    for method_name in ("create", "commit", "rollback"):
        method = getattr(aws_strategy, method_name)
        with contextlib.suppress(NotImplementedError, ImportError):
            method("res-test", "plan-test")
    print("sandbox-stubs-ok")


def _aws_profile() -> None:
    """Verify AWS profile name satisfies credential requirements."""
    saved = _clear_cloud_env()
    try:
        os.environ["AWS_PROFILE"] = "my-profile"
        resolved = resolve_credentials("aws", {})
        errors = validate_credentials("aws", resolved)
        assert len(errors) == 0, f"Profile should satisfy creds: {errors}"
        print("aws-profile-ok")
    finally:
        _restore_cloud_env(saved)


def _builtin_types() -> None:
    """Verify cloud hierarchy types are in BUILTIN_TYPES."""
    from cleveragents.application.services._resource_registry_data import (
        BUILTIN_TYPES,
    )

    names = [t["name"] for t in BUILTIN_TYPES]
    # Core cloud types that must be present
    for expected in (
        "cloud-account",
        "cloud-region",
        "cloud-network",
        "aws-account",
        "aws-region",
        "aws-vpc",
        "gcp",
        "azure",
    ):
        assert expected in names, f"'{expected}' not in BUILTIN_TYPES"
    print("builtin-types-ok")


_COMMANDS: dict[str, Any] = {
    "protocol-check": _protocol_check,
    "aws-resolve": _aws_resolve,
    "gcp-resolve": _gcp_resolve,
    "azure-resolve": _azure_resolve,
    "stub-resolve": _stub_resolve,
    "missing-creds": _missing_creds,
    "sandbox-stubs": _sandbox_stubs,
    "aws-profile": _aws_profile,
    "builtin-types": _builtin_types,
}


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>", file=sys.stderr)
        sys.exit(1)
    _COMMANDS[sys.argv[1]]()