diff --git a/.github/workflows/shared-config.yml b/.github/workflows/shared-config.yml index 619afdb5b..c74610e8a 100644 --- a/.github/workflows/shared-config.yml +++ b/.github/workflows/shared-config.yml @@ -130,10 +130,13 @@ jobs: except ImportError: import tomli as tomllib # backport for Python < 3.11 import json - META = {'all', 'dev', 'ci', 'test-aws', 'test-k8s', 'all-providers', 'k8s-legacy', 'monitoring-aws'} + META = {'all', 'dev', 'ci', 'all-providers', 'k8s-legacy', 'monitoring-aws'} with open('pyproject.toml', 'rb') as f: data = tomllib.load(f) - extras = [k for k in data['project']['optional-dependencies'].keys() if k not in META] + extras = [ + k for k in data['project']['optional-dependencies'] + if k not in META and not k.startswith('test-') + ] print(json.dumps(extras if extras else ['cli', 'api', 'monitoring'])) ") diff --git a/dev-tools/setup/run_tool.sh b/dev-tools/setup/run_tool.sh index 826d66276..30b5d80b0 100755 --- a/dev-tools/setup/run_tool.sh +++ b/dev-tools/setup/run_tool.sh @@ -7,6 +7,7 @@ set -e TOOL_NAME="$1" shift # Remove tool name from arguments +read -r -a UV_RUN_OPTIONS <<< "${RUN_TOOL_UV_OPTIONS:-}" # Walk upward from $PWD looking for a UV/Python project root. # Prints the path on success; prints nothing and returns 1 if not found. @@ -105,12 +106,12 @@ run_tool() { adjusted_args=$(adjust_relative_args "$project_root" "$@") if [ -n "$adjusted_args" ]; then # shellcheck disable=SC2086 - (cd "$project_root" && uv run "${TOOL_NAME}" $adjusted_args) + (cd "$project_root" && uv run "${UV_RUN_OPTIONS[@]}" "${TOOL_NAME}" $adjusted_args) else - (cd "$project_root" && uv run "${TOOL_NAME}") + (cd "$project_root" && uv run "${UV_RUN_OPTIONS[@]}" "${TOOL_NAME}") fi else - uv run "${TOOL_NAME}" "$@" + uv run "${UV_RUN_OPTIONS[@]}" "${TOOL_NAME}" "$@" fi elif [ -f ".venv/bin/${TOOL_NAME}" ] && venv_usable ".venv"; then echo "Executing with venv..." @@ -125,12 +126,12 @@ run_tool() { adjusted_args=$(adjust_relative_args "$project_root" "$@") if [ -n "$adjusted_args" ]; then # shellcheck disable=SC2086 - (cd "$project_root" && uv run python -m "${TOOL_NAME}" $adjusted_args) + (cd "$project_root" && uv run "${UV_RUN_OPTIONS[@]}" python -m "${TOOL_NAME}" $adjusted_args) else - (cd "$project_root" && uv run python -m "${TOOL_NAME}") + (cd "$project_root" && uv run "${UV_RUN_OPTIONS[@]}" python -m "${TOOL_NAME}") fi else - uv run python -m "${TOOL_NAME}" "$@" + uv run "${UV_RUN_OPTIONS[@]}" python -m "${TOOL_NAME}" "$@" fi else python3 -m "${TOOL_NAME}" "$@" diff --git a/makefiles/ci.mk b/makefiles/ci.mk index 70e78b3c5..6472097ac 100644 --- a/makefiles/ci.mk +++ b/makefiles/ci.mk @@ -16,6 +16,7 @@ ci-quality-radon: ## Run radon complexity analysis $(call run-tool,radon,cc $(PACKAGE) --min B --show-complexity) $(call run-tool,radon,mi $(PACKAGE) --min B) +ci-quality-pyright: RUN_TOOL_UV_OPTIONS := --extra all-providers ci-quality-pyright: ## Run pyright type checking @echo "Running pyright type check..." $(call run-tool,pyright,) diff --git a/makefiles/common.mk b/makefiles/common.mk index 55b8c44d1..28af42851 100644 --- a/makefiles/common.mk +++ b/makefiles/common.mk @@ -71,8 +71,9 @@ DOCS_BUILD_DIR := $(DOCS_DIR)/site # Centralized tool execution function # Usage: $(call run-tool,tool-name,arguments[,working-dir]) +# Set RUN_TOOL_UV_OPTIONS for options that belong to `uv run`, not the tool. define run-tool - $(if $(3),cd $(3) && ../dev-tools/setup/run_tool.sh $(1) $(2),@dev-tools/setup/run_tool.sh $(1) $(2)) + $(if $(3),cd $(3) && RUN_TOOL_UV_OPTIONS="$(RUN_TOOL_UV_OPTIONS)" ../dev-tools/setup/run_tool.sh $(1) $(2),@RUN_TOOL_UV_OPTIONS="$(RUN_TOOL_UV_OPTIONS)" dev-tools/setup/run_tool.sh $(1) $(2)) endef # Virtual environment setup (common dependency for all makefiles) diff --git a/pyproject.toml b/pyproject.toml index f7fcf4e10..17de2d614 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,12 +89,17 @@ k8s = [ "kubernetes", ] -# Placeholder entries for future providers (add when provider is implemented) -# azure = [ -# "azure-mgmt-compute>=30.0.0", -# "azure-identity>=1.15.0", -# "azure-mgmt-network>=25.0.0", -# ] +# Azure provider runtime dependencies. Install with: +# pip install orb-py[azure] +azure = [ + "azure-core>=1.38.2", + "azure-identity>=1.25.2", + "azure-mgmt-compute>=37.2.0", + "azure-mgmt-network>=30.2.0", + "azure-mgmt-resource>=25.0.0", + "azure-mgmt-resource-subscriptions>=1.0.0b1", + "httpx>=0.27.0", +] # gcp = [ # "google-cloud-compute>=1.14.0", # "google-auth>=2.23.0", @@ -103,6 +108,7 @@ k8s = [ # Meta-extra: all currently implemented providers all-providers = [ "orb-py[aws]", + "orb-py[azure]", "orb-py[k8s]", ] @@ -194,6 +200,15 @@ test-aws = [ "requests-mock>=1.11.0,<2.0.0", ] +# Azure SDK packages imported directly by the Azure tests. These overlap with +# the runtime extra because the tests assert behavior using real SDK exception +# and model types, but the test group does not install unrelated runtime SDKs. +test-azure = [ + "azure-core>=1.38.2", + "azure-identity>=1.25.2", + "azure-mgmt-compute>=37.2.0", +] + # Kubernetes test dependencies — kmock is an in-process HTTP-level K8s API mock. # Tests point the real kubernetes SDK client at the kmock URL so the full # client stack is exercised without a live cluster. diff --git a/src/orb/domain/template/factory.py b/src/orb/domain/template/factory.py index 2837ee71b..d7f264d9d 100644 --- a/src/orb/domain/template/factory.py +++ b/src/orb/domain/template/factory.py @@ -113,9 +113,9 @@ def create_template( except Exception as e: if self._logger: self._logger.error("Failed to create %s template: %s", provider_type, e) - # Fall back to core template + raise - # Fall back to core template + # Use the core template only when no provider-specific class is registered. try: template = Template(**template_data) diff --git a/src/orb/domain/template/template_aggregate.py b/src/orb/domain/template/template_aggregate.py index b651af15d..cc5c213e2 100644 --- a/src/orb/domain/template/template_aggregate.py +++ b/src/orb/domain/template/template_aggregate.py @@ -74,6 +74,7 @@ class Template(BaseModel): # Tags and metadata tags: dict[str, Any] = Field(default_factory=dict) metadata: dict[str, Any] = Field(default_factory=dict) + provider_data: dict[str, Any] = Field(default_factory=dict) # Provider configuration (multi-provider support) provider_type: Optional[str] = None diff --git a/src/orb/providers/azure/__init__.py b/src/orb/providers/azure/__init__.py new file mode 100644 index 000000000..b29b5a233 --- /dev/null +++ b/src/orb/providers/azure/__init__.py @@ -0,0 +1,21 @@ +"""Azure Provider implementation.""" + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.configuration.template_extension import AzureTemplateExtensionConfig +from orb.providers.azure.registration import ( + get_azure_extension_defaults, + initialize_azure_provider, + is_azure_provider_registered, + register_azure_extensions, + register_azure_template_factory, +) + +__all__: list[str] = [ + "AzureProviderConfig", + "AzureTemplateExtensionConfig", + "get_azure_extension_defaults", + "initialize_azure_provider", + "is_azure_provider_registered", + "register_azure_extensions", + "register_azure_template_factory", +] diff --git a/src/orb/providers/azure/auth/__init__.py b/src/orb/providers/azure/auth/__init__.py new file mode 100644 index 000000000..8f40fc448 --- /dev/null +++ b/src/orb/providers/azure/auth/__init__.py @@ -0,0 +1 @@ +"""Azure authentication utilities.""" diff --git a/src/orb/providers/azure/auth/azure_auth_strategy.py b/src/orb/providers/azure/auth/azure_auth_strategy.py new file mode 100644 index 000000000..b6024070c --- /dev/null +++ b/src/orb/providers/azure/auth/azure_auth_strategy.py @@ -0,0 +1,143 @@ +"""Azure DefaultAzureCredential authentication strategy. + +Uses the shared Azure infrastructure credential factory so auth and +provider-runtime flows construct the same credential shape. +""" + +import asyncio +from typing import Any, Optional + +from orb.domain.base.ports import LoggingPort +from orb.infrastructure.adapters.ports.auth import ( + AuthContext, + AuthPort, + AuthResult, + AuthStatus, +) +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncAzureAccessTokenProviderProtocol, + AsyncDefaultAzureAccessTokenProvider, + AzureAccessTokenProviderProtocol, + DefaultAzureAccessTokenProvider, +) + + +@injectable +class AzureAuthStrategy(AuthPort): + """Authentication strategy using Azure DefaultAzureCredential.""" + + def __init__( + self, + logger: LoggingPort, + client_id: Optional[str] = None, + enabled: bool = True, + token_provider: Optional[AzureAccessTokenProviderProtocol] = None, + async_token_provider: Optional[AsyncAzureAccessTokenProviderProtocol] = None, + ) -> None: + """Initialise Azure auth with async-first token acquisition.""" + self._logger = logger + self.client_id = client_id + self.enabled = enabled + self._token_provider = token_provider + self._async_token_provider = async_token_provider + if self._token_provider is None and self._async_token_provider is None: + self._async_token_provider = AsyncDefaultAzureAccessTokenProvider( + client_id=client_id, + logger=logger, + ) + + @classmethod + def from_auth_config(cls, auth_config: Any) -> "AzureAuthStrategy": + """Build an Azure auth strategy from the shared AuthConfig object.""" + from orb.infrastructure.adapters.logging_adapter import LoggingAdapter + + provider_auth = getattr(auth_config, "provider_auth", None) + azure_cfg = getattr(provider_auth, "azure", None) if provider_auth is not None else None + client_id = getattr(azure_cfg, "client_id", None) if azure_cfg is not None else None + enabled = bool(getattr(auth_config, "enabled", True)) + return cls( + logger=LoggingAdapter(), + client_id=client_id, + enabled=enabled, + ) + + def _auth_error_types(self) -> tuple[type[Exception], ...]: + """Return the active provider's declared auth failure types.""" + if self._async_token_provider is not None: + return self._async_token_provider.get_auth_error_types() + if self._token_provider is None: + return DefaultAzureAccessTokenProvider( + client_id=self.client_id, + logger=self._logger, + ).get_auth_error_types() + return self._token_provider.get_auth_error_types() + + async def _get_access_token(self, scope: str) -> str: + """Resolve an ARM token without blocking the event loop.""" + if self._async_token_provider is not None: + return await self._async_token_provider.get_access_token(scope) + if self._token_provider is None: + raise RuntimeError("Azure auth strategy has no token provider configured") + return await asyncio.to_thread(self._token_provider.get_access_token, scope) + + async def authenticate(self, context: AuthContext) -> AuthResult: + if not self.enabled: + return AuthResult( + status=AuthStatus.FAILED, + error_message="Azure auth strategy disabled", + ) + try: + token = await self._get_access_token("https://management.azure.com/.default") + + return AuthResult( + status=AuthStatus.SUCCESS, + user_id=self.client_id or "azure-identity", + token=token, + user_roles=["provider"], + metadata={ + "strategy": "azure_default_credential", + }, + ) + except self._auth_error_types() as exc: + self._logger.error("Azure authentication failed: %s", exc) + return AuthResult( + status=AuthStatus.FAILED, + error_message=f"Azure authentication failed: {exc}", + ) + + async def validate_token(self, token: str) -> AuthResult: + """Token validation is handled by Azure SDK internally.""" + return AuthResult( + status=AuthStatus.SUCCESS, + token=token, + metadata={"strategy": "azure_default_credential"}, + ) + + async def refresh_token(self, refresh_token: str) -> AuthResult: + """DefaultAzureCredential handles token refresh automatically.""" + try: + token = await self._get_access_token("https://management.azure.com/.default") + return AuthResult( + status=AuthStatus.SUCCESS, + token=token, + metadata={"strategy": "azure_default_credential", "refreshed": True}, + ) + except self._auth_error_types() as exc: + return AuthResult( + status=AuthStatus.FAILED, + error_message=f"Token refresh failed: {exc}", + ) + + async def revoke_token(self, token: str) -> bool: + """Azure ARM tokens cannot be revoked directly.""" + self._logger.debug("Token revocation not supported for Azure ARM tokens") + return False + + def get_strategy_name(self) -> str: + """Return the identifier for the Azure default-credential strategy.""" + return "azure_default_credential" + + def is_enabled(self) -> bool: + """Return whether this auth strategy is currently enabled.""" + return self.enabled diff --git a/src/orb/providers/azure/capabilities.py b/src/orb/providers/azure/capabilities.py new file mode 100644 index 000000000..4d0973ea8 --- /dev/null +++ b/src/orb/providers/azure/capabilities.py @@ -0,0 +1,48 @@ +"""Shared Azure provider capability metadata.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +from orb.providers.azure.domain.template.value_objects import AzureProviderApi + +_AZURE_API_CAPABILITIES: dict[str, dict[str, Any]] = { + AzureProviderApi.VMSS.value: { + "supported_fleet_types": [], + "supports_spot": True, + "supports_on_demand": True, + "max_instances": 1000, + }, + AzureProviderApi.VMSS_UNIFORM.value: { + "supported_fleet_types": [], + "supports_spot": True, + "supports_on_demand": True, + "max_instances": 1000, + }, + AzureProviderApi.SINGLE_VM.value: { + "supported_fleet_types": [], + "supports_spot": True, + "supports_on_demand": True, + "max_instances": 1000, + }, + AzureProviderApi.CYCLECLOUD.value: { + "supported_fleet_types": [], + "supports_spot": False, + "supports_on_demand": True, + "requires_existing_cluster": True, + "required_create_fields": ["cluster_name", "node_array"], + "capacity_limit_source": "cluster_status.maxCount", + "supports_async_operations": True, + }, +} + + +def get_supported_api_capabilities() -> dict[str, dict[str, Any]]: + """Return Azure API capability metadata.""" + return deepcopy(_AZURE_API_CAPABILITIES) + + +def get_supported_apis() -> list[str]: + """Return the canonical list of Azure provider APIs.""" + return list(_AZURE_API_CAPABILITIES.keys()) diff --git a/src/orb/providers/azure/cli/__init__.py b/src/orb/providers/azure/cli/__init__.py new file mode 100644 index 000000000..d5de32869 --- /dev/null +++ b/src/orb/providers/azure/cli/__init__.py @@ -0,0 +1 @@ +"""Azure CLI integration helpers.""" diff --git a/src/orb/providers/azure/cli/azure_cli_spec.py b/src/orb/providers/azure/cli/azure_cli_spec.py new file mode 100644 index 000000000..db99cc441 --- /dev/null +++ b/src/orb/providers/azure/cli/azure_cli_spec.py @@ -0,0 +1,176 @@ +"""Azure-specific CLI argument specification.""" + +import argparse +import re +from typing import Any + + +class AzureCLISpec: + """CLI spec for the Azure provider.""" + + @staticmethod + def _arg(args: argparse.Namespace, name: str, default: Any = None) -> Any: + """Read a parsed CLI argument by name.""" + return vars(args).get(name, default) + + def add_arguments(self, parser: argparse.ArgumentParser) -> None: + """Add Azure provider arguments to the parser.""" + parser.add_argument( + "--azure-subscription-id", + dest="azure_subscription_id", + help="Azure subscription ID", + ) + parser.add_argument( + "--azure-resource-group", + dest="azure_resource_group", + help="Azure resource group", + ) + parser.add_argument( + "--azure-location", + dest="azure_location", + help="Azure location", + ) + parser.add_argument( + "--azure-client-id", + dest="azure_client_id", + help="Managed identity client ID", + ) + parser.add_argument( + "--azure-cyclecloud-url", + dest="azure_cyclecloud_url", + help="CycleCloud URL", + ) + parser.add_argument( + "--azure-cyclecloud-credential-path", + dest="azure_cyclecloud_credential_path", + help="Secret path or file path for CycleCloud credentials", + ) + parser.add_argument( + "--azure-cyclecloud-auth-mode", + dest="azure_cyclecloud_auth_mode", + help="CycleCloud auth mode override", + ) + parser.add_argument( + "--azure-cyclecloud-aad-scope", + dest="azure_cyclecloud_aad_scope", + help="CycleCloud AAD scope override", + ) + parser.add_argument( + "--azure-cyclecloud-verify-ssl", + dest="azure_cyclecloud_verify_ssl", + action="store_true", + help="Verify TLS certificates for CycleCloud", + ) + parser.add_argument( + "--azure-cyclecloud-no-verify-ssl", + dest="azure_cyclecloud_no_verify_ssl", + action="store_true", + help="Disable TLS certificate verification for CycleCloud", + ) + + def extract_config(self, args: argparse.Namespace) -> dict[str, Any]: + """Return full config dict from args.""" + result: dict[str, Any] = { + "subscription_id": self._arg(args, "azure_subscription_id"), + "resource_group": self._arg(args, "azure_resource_group"), + "region": self._arg(args, "azure_location") or "eastus2", + } + + client_id = self._arg(args, "azure_client_id") + if client_id: + result["client_id"] = client_id + + cyclecloud = self._extract_cyclecloud_config(args, partial=False) + if cyclecloud: + result["cyclecloud"] = cyclecloud + + return result + + def extract_partial_config(self, args: argparse.Namespace) -> dict[str, Any]: + """Return only explicitly provided Azure config fields.""" + result: dict[str, Any] = {} + + if self._arg(args, "azure_subscription_id") is not None: + result["subscription_id"] = args.azure_subscription_id + if self._arg(args, "azure_resource_group") is not None: + result["resource_group"] = args.azure_resource_group + if self._arg(args, "azure_location") is not None: + result["region"] = args.azure_location + if self._arg(args, "azure_client_id") is not None: + result["client_id"] = args.azure_client_id + + cyclecloud = self._extract_cyclecloud_config(args, partial=True) + if cyclecloud: + result["cyclecloud"] = cyclecloud + + return result + + def validate_add(self, args: argparse.Namespace) -> list[str]: + """Return error messages for missing required add fields.""" + errors: list[str] = [] + if not self._arg(args, "azure_subscription_id"): + errors.append("--azure-subscription-id is required") + if not self._arg(args, "azure_resource_group"): + errors.append("--azure-resource-group is required") + return errors + + def generate_name(self, args: argparse.Namespace) -> str: + """Generate a provider instance name from Azure subscription and location.""" + subscription_id = self._arg(args, "azure_subscription_id") or "default" + location = self._arg(args, "azure_location") or "eastus2" + sanitized_subscription = re.sub(r"[^a-zA-Z0-9\-_]", "-", subscription_id) + return f"azure_{sanitized_subscription}_{location}" + + def format_display(self, config: dict[str, Any]) -> list[tuple[str, str]]: + """Return (label, value) pairs for display.""" + cyclecloud = config.get("cyclecloud") or {} + return [ + ("Subscription", config.get("subscription_id", "-")), + ("Resource Group", config.get("resource_group", "-")), + ("Location", config.get("region", config.get("location", "-"))), + ("Client ID", config.get("client_id", "-")), + ("CycleCloud URL", cyclecloud.get("url", "-")), + ("CycleCloud Credential Path", cyclecloud.get("credential_path", "-")), + ] + + def _extract_cyclecloud_config( + self, args: argparse.Namespace, *, partial: bool + ) -> dict[str, Any]: + """Extract nested CycleCloud config from CLI args.""" + cyclecloud: dict[str, Any] = {} + + field_map = { + "url": "azure_cyclecloud_url", + "credential_path": "azure_cyclecloud_credential_path", + "auth_mode": "azure_cyclecloud_auth_mode", + "aad_scope": "azure_cyclecloud_aad_scope", + } + for config_key, arg_name in field_map.items(): + value = self._arg(args, arg_name) + if value is not None: + cyclecloud[config_key] = value + + verify_ssl = self._extract_verify_ssl(args) + if verify_ssl is not None: + cyclecloud["verify_ssl"] = verify_ssl + + if partial: + return cyclecloud + + return {k: v for k, v in cyclecloud.items() if v is not None} + + @staticmethod + def _extract_verify_ssl(args: argparse.Namespace) -> bool | None: + """Resolve CycleCloud TLS verification flags.""" + verify_ssl = vars(args).get("azure_cyclecloud_verify_ssl", False) + no_verify_ssl = vars(args).get("azure_cyclecloud_no_verify_ssl", False) + if verify_ssl and no_verify_ssl: + raise ValueError( + "Cannot specify both --azure-cyclecloud-verify-ssl and " + "--azure-cyclecloud-no-verify-ssl" + ) + if verify_ssl: + return True + if no_verify_ssl: + return False + return None diff --git a/src/orb/providers/azure/configuration/__init__.py b/src/orb/providers/azure/configuration/__init__.py new file mode 100644 index 000000000..8b9c821bb --- /dev/null +++ b/src/orb/providers/azure/configuration/__init__.py @@ -0,0 +1,15 @@ +"""Azure provider configuration.""" + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.configuration.template_extension import AzureTemplateExtensionConfig +from orb.providers.azure.configuration.validator import ( + validate_azure_config, + validate_azure_template, +) + +__all__: list[str] = [ + "AzureProviderConfig", + "AzureTemplateExtensionConfig", + "validate_azure_config", + "validate_azure_template", +] diff --git a/src/orb/providers/azure/configuration/config.py b/src/orb/providers/azure/configuration/config.py new file mode 100644 index 000000000..2bf3ffe0d --- /dev/null +++ b/src/orb/providers/azure/configuration/config.py @@ -0,0 +1,188 @@ +"""Azure configuration provider - single source of truth.""" + +import re +from typing import Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +from orb.infrastructure.interfaces.provider import BaseProviderConfig + + +class CycleCloudConfig(BaseModel): + """CycleCloud connection configuration.""" + + model_config = ConfigDict(extra="forbid") + + url: Optional[str] = Field(None, description="CycleCloud REST API base URL") + credential_path: Optional[str] = Field( + None, + description="Path to a JSON file containing CycleCloud credentials and optional auth overrides", + ) + verify_ssl: Optional[bool] = Field( + None, + description=( + "Whether to verify TLS certificates for CycleCloud API calls. " + "None means 'unset here' so request/template/credential-file/default " + "precedence can resolve the effective value later." + ), + ) + auth_mode: Optional[str] = Field( + None, + description="CycleCloud auth mode override, e.g. 'basic' or 'bearer'.", + ) + aad_scope: Optional[str] = Field( + None, + description="AAD scope used to resolve a bearer token for CycleCloud.", + ) + + @model_validator(mode="before") + @classmethod + def reject_inline_basic_auth(cls, data: object) -> object: + """Reject inline credentials that must be supplied via secret-path.""" + if not isinstance(data, dict): + return data + + forbidden_fields = { + "username", + "password", + "bearer_token", + "cyclecloud_username", + "cyclecloud_password", + "cyclecloud_bearer_token", + } + present = [field for field in forbidden_fields if data.get(field) not in (None, "")] + if present: + raise ValueError( + "CycleCloud inline username/password config is not supported; " + "use credential_path instead." + ) + return data + + +# --------------------------------------------------------------------------- +# Root provider config +# --------------------------------------------------------------------------- + + +class AzureProviderConfig(BaseSettings, BaseProviderConfig): # type: ignore[misc] # Pydantic model mixin MRO is valid at runtime. + """Configuration for the Azure provider (VMSS / Compute Fleet). + + The shared provider interface uses ``region`` across providers. Azure's + platform term is ``location``. This model keeps the shared ``region`` + field to stay aligned with ``BaseProviderConfig`` while accepting + Azure-native ``location`` input at the boundary. + """ + + # Pydantic's inherited model_config annotations disagree despite compatible runtime values. + model_config = SettingsConfigDict( # type: ignore[assignment] + env_prefix="ORB_AZURE_", + case_sensitive=False, + populate_by_name=True, + env_nested_delimiter="__", + extra="forbid", + ) + + # ------------------------------------------------------------------ + # Provider identity + # ------------------------------------------------------------------ + provider_type: str = Field("azure", description="Provider type identifier") + region: str = Field( + default="eastus2", + description="Azure location slug", + ) + + # ------------------------------------------------------------------ + # Azure subscription & resource targeting + # ------------------------------------------------------------------ + subscription_id: Optional[str] = Field(None, description="Azure subscription ID (UUID)") + resource_group: Optional[str] = Field( + None, + description="Default resource group for created resources (1-90 chars)", + ) + + # ------------------------------------------------------------------ + # Authentication + # ------------------------------------------------------------------ + client_id: Optional[str] = Field( + None, + description="Managed-identity client ID for user-assigned identity selection", + ) + + # ------------------------------------------------------------------ + # Retry / timeout + # ------------------------------------------------------------------ + max_retries: int = Field(3, ge=0, description="Maximum SDK retry attempts for transient errors") + connect_timeout: int = Field( + 30, ge=1, description="Connection timeout for ARM API calls in seconds" + ) + read_timeout: int = Field(60, ge=1, description="Read timeout for ARM API calls in seconds") + + # ------------------------------------------------------------------ + # CycleCloud + # ------------------------------------------------------------------ + cyclecloud: Optional[CycleCloudConfig] = Field( + default=None, + description="CycleCloud integration configuration (URL, credentials, TLS verification)", + ) + + # ------------------------------------------------------------------ + # Validators + # ------------------------------------------------------------------ + @model_validator(mode="before") + @classmethod + def normalise_location_input(cls, data: object) -> object: + """Translate Azure-native ``location`` input onto the shared ``region`` field.""" + if not isinstance(data, dict): + return data + + data = dict(data) + location = data.get("location") + region = data.get("region") + + if location not in (None, "") and region not in (None, "") and location != region: + raise ValueError( + "Azure provider config received conflicting 'location' and 'region' values" + ) + + if region in (None, "") and location not in (None, ""): + data["region"] = location + if location not in (None, ""): + data.pop("location", None) + + return data + + @field_validator("resource_group") + @classmethod + def validate_resource_group(cls, v: Optional[str]) -> Optional[str]: + """Enforce ARM resource-group naming rules.""" + if v is None: + return v + if not (1 <= len(v) <= 90): + raise ValueError("resource_group must be 1-90 characters") + if not re.match(r"^[a-zA-Z0-9_\-.()\[\]]+$", v): + raise ValueError( + "resource_group contains invalid characters " + "(allowed: alphanumeric, _, -, ., (, ), [, ])" + ) + return v + + @field_validator("subscription_id") + @classmethod + def validate_subscription_id(cls, v: Optional[str]) -> Optional[str]: + """Validate UUID-like subscription ID format.""" + if v is None: + return v + if not re.match( + r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$", + v, + ): + raise ValueError( + "subscription_id must be a valid UUID (xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx)" + ) + return v + + @property + def location(self) -> str: + """Azure-native accessor for callers operating on Azure concepts.""" + return self.region diff --git a/src/orb/providers/azure/configuration/template_extension.py b/src/orb/providers/azure/configuration/template_extension.py new file mode 100644 index 000000000..c1e3d85e7 --- /dev/null +++ b/src/orb/providers/azure/configuration/template_extension.py @@ -0,0 +1,224 @@ +"""Azure template extension configuration.""" + +from typing import Any, Optional + +from pydantic import BaseModel, ConfigDict, Field + +from orb.providers.azure.domain.template.value_objects import ( + AzureAllocationStrategy, + AzureCapacityReservationGroupId, + AzureDataDisk, + AzureDiskEncryptionSetId, + AzureEvictionPolicy, + AzureImageReference, + AzureNetworkConfig, + AzureOSDiskConfig, + AzurePriority, + AzureProviderApi, + AzureProximityPlacementGroupId, + AzureResourceGroupName, + AzureSecurityType, + AzureUpgradePolicyMode, + AzureVmSizePreference, + AzureVMSSOrchestrationMode, +) +from orb.providers.azure.services.spot_placement_planner import PlacementSplitStrategy + + +class AzureTemplateExtensionConfig(BaseModel): + """Azure-specific template extension defaults. + + Registered with ``TemplateExtensionRegistry`` so the template factory + can apply Azure defaults when ``provider_type == "azure"``. + """ + + model_config = ConfigDict(populate_by_name=True) + + # Azure DTO-only fields preserved by TemplateDTO.provider_config round-trips. + resource_group: Optional[AzureResourceGroupName] = Field( + default=None, description="Azure resource group for the template" + ) + location: Optional[str] = Field(default=None, description="Azure location for the template") + subscription_id: Optional[str] = Field( + default=None, description="Azure subscription override for the template" + ) + vmss_name: Optional[str] = Field(default=None, description="Explicit VMSS name") + orchestration_mode: Optional[AzureVMSSOrchestrationMode] = Field( + default=None, description="VMSS orchestration mode" + ) + platform_fault_domain_count: Optional[int] = Field( + default=None, description="Fault domain count for Flexible orchestration" + ) + single_placement_group: Optional[bool] = Field( + default=None, description="Restrict VMSS to a single placement group" + ) + image: Optional[AzureImageReference] = Field( + default=None, description="Azure VM image reference" + ) + eviction_policy: Optional[AzureEvictionPolicy] = Field( + default=None, description="Spot eviction policy" + ) + billing_profile_max_price: Optional[float] = Field( + default=None, description="Maximum Spot VM price" + ) + spot_percentage: Optional[int] = Field( + default=None, description="Desired percentage of Spot VMs" + ) + base_regular_priority_count: Optional[int] = Field( + default=None, description="Base regular-priority VM count for priority mix" + ) + spot_restore_enabled: Optional[bool] = Field( + default=None, description="Enable Spot Try-Restore" + ) + spot_restore_timeout: Optional[str] = Field( + default=None, description="ISO 8601 Spot Try-Restore timeout" + ) + spot_placement_score_enabled: Optional[bool] = Field( + default=None, description="Enable Azure Spot Placement Score planning before launch" + ) + placement_split_strategy: Optional[PlacementSplitStrategy] = Field( + default=None, description="How Spot Placement Score launches split capacity" + ) + placement_primary_share_percent: Optional[int] = Field( + default=None, description="Capacity percentage assigned to the top placement candidate" + ) + placement_regions: Optional[list[str]] = Field( + default=None, description="Azure regions considered for Spot Placement Score planning" + ) + placement_zones: Optional[list[str]] = Field( + default=None, description="Azure zones considered for Spot Placement Score planning" + ) + zones: Optional[list[str]] = Field(default=None, description="Availability zones") + zone_balance: Optional[bool] = Field(default=None, description="Enable zone balancing") + proximity_placement_group_id: Optional[AzureProximityPlacementGroupId] = Field( + default=None, description="Proximity placement group ARM resource ID" + ) + capacity_reservation_group_id: Optional[AzureCapacityReservationGroupId] = Field( + default=None, description="Capacity reservation group ARM resource ID" + ) + os_disk: Optional[AzureOSDiskConfig] = Field(default=None, description="OS disk config") + data_disks: Optional[list[AzureDataDisk]] = Field(default=None, description="Data disks") + network_config: Optional[AzureNetworkConfig] = Field( + default=None, description="Azure networking config" + ) + security_type: Optional[AzureSecurityType] = Field(default=None, description="VM security type") + secure_boot_enabled: Optional[bool] = Field(default=None, description="Enable UEFI Secure Boot") + vtpm_enabled: Optional[bool] = Field(default=None, description="Enable vTPM") + encryption_at_host: Optional[bool] = Field( + default=None, description="Enable host-based encryption" + ) + disk_encryption_set_id: Optional[AzureDiskEncryptionSetId] = Field( + default=None, description="Disk encryption set ARM resource ID" + ) + ssh_key_name: Optional[str] = Field( + default=None, description="Azure SSH Public Key resource name" + ) + ssh_public_keys: Optional[list[str]] = Field(default=None, description="Inline SSH public keys") + user_assigned_identity_ids: Optional[list[str]] = Field( + default=None, description="User-assigned managed identity ARM resource IDs" + ) + system_assigned_identity: Optional[bool] = Field( + default=None, description="Enable system-assigned managed identity" + ) + custom_data: Optional[str] = Field(default=None, description="Base64 custom-data payload") + extension_profile: Optional[list[dict[str, Any]]] = Field( + default=None, description="VMSS extension definitions" + ) + overprovision: Optional[bool] = Field(default=None, description="Enable VMSS overprovisioning") + upgrade_policy_mode: Optional[AzureUpgradePolicyMode] = Field( + default=None, description="VMSS upgrade policy mode" + ) + provider_api_spec: Optional[dict[str, Any]] = Field( + default=None, description="Raw Azure provider request payload override" + ) + provider_api_spec_file: Optional[str] = Field( + default=None, description="Path to a native Azure provider spec file" + ) + cluster_name: Optional[str] = Field(default=None, description="CycleCloud cluster name") + node_array: Optional[str] = Field(default=None, description="CycleCloud node array") + cyclecloud_url: Optional[str] = Field(default=None, description="CycleCloud API URL") + cyclecloud_credential_path: Optional[str] = Field( + default=None, description="CycleCloud credential reference path" + ) + cyclecloud_verify_ssl: Optional[bool] = Field( + default=None, description="CycleCloud SSL verification setting" + ) + cyclecloud_auth_mode: Optional[str] = Field( + default=None, description="CycleCloud auth mode override" + ) + cyclecloud_aad_scope: Optional[str] = Field(default=None, description="CycleCloud AAD scope") + provider_api: Optional[AzureProviderApi] = Field(default=None, description="Azure provider API") + + # VM configuration + vm_size: Optional[str] = Field( + default=None, + description="Explicit Azure VM size default", + ) + vm_sizes: Optional[list[str]] = Field( + default=None, + description="Additional Azure VM size candidates for generic instance mix", + ) + vm_size_preferences: Optional[list[AzureVmSizePreference]] = Field( + default=None, + description="Ranked Azure VM size candidates for Prioritized VMSS instance mix", + ) + vmss_allocation_strategy: Optional[AzureAllocationStrategy] = Field( + default=None, + description="Azure VMSS instance-mix allocation strategy", + ) + + # Pricing + priority: AzurePriority = Field( + default=AzurePriority.REGULAR, description="Default VM priority" + ) + + # OS disk + os_disk_type: Optional[str] = Field(default=None, description="Default OS disk type") + os_disk_size_gb: Optional[int] = Field( + default=None, description="OS disk size in GiB (None = image default)" + ) + + # Identity + admin_username: str = Field(default="azureuser", description="Default admin username") + + # Freeform attributes + node_attributes: dict[str, Any] = Field( + default_factory=dict, + description="Freeform attributes merged into the template", + ) + + def to_template_defaults(self) -> dict[str, Any]: + """Convert extension to a dict of default values for template creation.""" + defaults: dict[str, Any] = { + "priority": self.priority, + "admin_username": self.admin_username, + "node_attributes": self.node_attributes, + } + if self.vm_size: + defaults["vm_size"] = self.vm_size + if self.vm_sizes: + defaults["vm_sizes"] = self.vm_sizes + if self.vm_size_preferences: + defaults["vm_size_preferences"] = self.vm_size_preferences + if self.vmss_allocation_strategy: + defaults["vmss_allocation_strategy"] = self.vmss_allocation_strategy + if self.spot_placement_score_enabled is not None: + defaults["spot_placement_score_enabled"] = self.spot_placement_score_enabled + if self.placement_split_strategy: + defaults["placement_split_strategy"] = self.placement_split_strategy + if self.placement_primary_share_percent is not None: + defaults["placement_primary_share_percent"] = self.placement_primary_share_percent + if self.placement_regions: + defaults["placement_regions"] = self.placement_regions + if self.placement_zones: + defaults["placement_zones"] = self.placement_zones + if self.os_disk is not None: + defaults["os_disk"] = self.os_disk.model_dump(mode="json", exclude_none=True) + elif self.os_disk_type is not None or self.os_disk_size_gb is not None: + legacy_os_disk: dict[str, Any] = { + "storage_account_type": self.os_disk_type or "Premium_LRS", + } + if self.os_disk_size_gb is not None: + legacy_os_disk["disk_size_gb"] = self.os_disk_size_gb + defaults["os_disk"] = legacy_os_disk + return defaults diff --git a/src/orb/providers/azure/configuration/validator.py b/src/orb/providers/azure/configuration/validator.py new file mode 100644 index 000000000..57796ab79 --- /dev/null +++ b/src/orb/providers/azure/configuration/validator.py @@ -0,0 +1,68 @@ +"""Azure configuration and template validation utilities.""" + +from typing import Any + +from pydantic import ValidationError + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate + + +def validate_azure_config(config: AzureProviderConfig) -> dict[str, Any]: + """Validate an AzureProviderConfig and return a structured result. + + Returns: + dict with keys: valid (bool), errors (list[str]), warnings (list[str]) + """ + errors: list[str] = [] + warnings: list[str] = [] + + if not config.subscription_id: + errors.append("subscription_id is required for Azure operations") + if not config.resource_group: + warnings.append("resource_group not set; templates must specify it explicitly") + if not config.region: + warnings.append("region not set; defaulting to 'eastus2'") + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": warnings, + } + + +def validate_azure_template(template_config: dict[str, Any]) -> dict[str, Any]: + """Validate an Azure template configuration dict via ``AzureTemplate``. + + ``AzureTemplate`` is the authoritative validation path for Azure template + configuration. This helper exists only to normalize model-validation + failures into the validation result shape used by the strategy and + validation adapter. + """ + warnings: list[str] = [] + + vm_size = template_config.get("vm_size", "") + if vm_size and not vm_size.startswith("Standard_"): + warnings.append( + f"Uncommon VM size format: '{vm_size}'. " + "Azure VM sizes typically start with 'Standard_'." + ) + + try: + AzureTemplate.model_validate(template_config) + errors: list[str] = [] + except ValidationError as exc: + errors = [] + for issue in exc.errors(): + message = issue.get("msg", "Validation error") + location = ".".join(str(part) for part in issue.get("loc", [])) + errors.append(f"{location}: {message}" if location else message) + except Exception as exc: + errors = [str(exc)] + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": warnings, + "validated_fields": list(template_config.keys()), + } diff --git a/src/orb/providers/azure/defaults_loader.py b/src/orb/providers/azure/defaults_loader.py new file mode 100644 index 000000000..0e7c33433 --- /dev/null +++ b/src/orb/providers/azure/defaults_loader.py @@ -0,0 +1,38 @@ +"""Azure provider defaults loader.""" + +from __future__ import annotations + +from orb.domain.base.ports.provider_defaults_loader_port import ProviderDefaultsLoaderPort +from orb.providers.azure.capabilities import get_supported_api_capabilities +from orb.providers.azure.registration import get_azure_extension_defaults + +_HANDLER_CLASSES = { + "VMSS": "VMSSHandler", + "VMSSUniform": "VMSSHandler", + "SingleVM": "SingleVMHandler", + "CycleCloud": "CycleCloudHandler", +} + + +class AzureDefaultsLoader: + """Loads Azure provider defaults from the provider-owned template extension.""" + + def load_defaults(self) -> dict: + """Return Azure provider defaults contributed by the Azure provider.""" + handlers = { + api: {"handler_class": _HANDLER_CLASSES[api], **capabilities} + for api, capabilities in get_supported_api_capabilities().items() + } + return { + "provider": { + "provider_defaults": { + "azure": { + "handlers": handlers, + "template_defaults": get_azure_extension_defaults(), + } + } + } + } + + +assert isinstance(AzureDefaultsLoader(), ProviderDefaultsLoaderPort) diff --git a/src/orb/providers/azure/domain/__init__.py b/src/orb/providers/azure/domain/__init__.py new file mode 100644 index 000000000..998e06afa --- /dev/null +++ b/src/orb/providers/azure/domain/__init__.py @@ -0,0 +1,37 @@ +"""Azure domain layer.""" + +from orb.providers.azure.domain.template import ( + AzureAllocationStrategy, + AzureCapacityReservationGroupId, + AzureDiskEncryptionSetId, + AzureEvictionPolicy, + AzureLocationName, + AzureOSDiskType, + AzurePriority, + AzureProviderApi, + AzureProximityPlacementGroupId, + AzureResourceGroupName, + AzureSecurityType, + AzureTemplate, + AzureUpgradePolicyMode, + AzureVmSizePreference, + AzureVMSSOrchestrationMode, +) + +__all__: list[str] = [ + "AzureAllocationStrategy", + "AzureCapacityReservationGroupId", + "AzureDiskEncryptionSetId", + "AzureEvictionPolicy", + "AzureLocationName", + "AzureOSDiskType", + "AzurePriority", + "AzureProximityPlacementGroupId", + "AzureProviderApi", + "AzureResourceGroupName", + "AzureSecurityType", + "AzureTemplate", + "AzureUpgradePolicyMode", + "AzureVmSizePreference", + "AzureVMSSOrchestrationMode", +] diff --git a/src/orb/providers/azure/domain/template/__init__.py b/src/orb/providers/azure/domain/template/__init__.py new file mode 100644 index 000000000..0be751656 --- /dev/null +++ b/src/orb/providers/azure/domain/template/__init__.py @@ -0,0 +1,37 @@ +"""Azure template domain - VMSS-based template configuration.""" + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import ( + AzureAllocationStrategy, + AzureCapacityReservationGroupId, + AzureDiskEncryptionSetId, + AzureEvictionPolicy, + AzureLocationName, + AzureOSDiskType, + AzurePriority, + AzureProviderApi, + AzureProximityPlacementGroupId, + AzureResourceGroupName, + AzureSecurityType, + AzureUpgradePolicyMode, + AzureVmSizePreference, + AzureVMSSOrchestrationMode, +) + +__all__: list[str] = [ + "AzureAllocationStrategy", + "AzureCapacityReservationGroupId", + "AzureDiskEncryptionSetId", + "AzureEvictionPolicy", + "AzureLocationName", + "AzureOSDiskType", + "AzurePriority", + "AzureProximityPlacementGroupId", + "AzureProviderApi", + "AzureResourceGroupName", + "AzureSecurityType", + "AzureTemplate", + "AzureUpgradePolicyMode", + "AzureVmSizePreference", + "AzureVMSSOrchestrationMode", +] diff --git a/src/orb/providers/azure/domain/template/azure_template_aggregate.py b/src/orb/providers/azure/domain/template/azure_template_aggregate.py new file mode 100644 index 000000000..c8d6112f6 --- /dev/null +++ b/src/orb/providers/azure/domain/template/azure_template_aggregate.py @@ -0,0 +1,759 @@ +"""Azure VMSS-based template aggregate. + +Provides an Azure-specific extension of the core Template that targets +Virtual Machine Scale Sets (VMSS) as the primary provisioning mechanism. + +Key Azure VMSS concepts modelled here: +- VM size selection (single or multiple for Spot/Flex) +- Spot / low-priority pricing with eviction policies +- VMSS orchestration mode (Flexible vs Uniform) +- Availability zone spreading +- OS disk & data disk configuration +- Networking (subnet, NSG, accelerated networking) +- Image references (Marketplace, custom, gallery) +- Proximity placement groups & capacity reservations +- User data / custom-data / cloud-init +- Extension profiles (custom script, etc.) +- Overprovisioning and upgrade policies + +See: https://learn.microsoft.com/en-us/azure/virtual-machine-scale-sets/overview + https://learn.microsoft.com/en-us/rest/api/compute/virtual-machine-scale-sets +""" + +from typing import Any, Optional + +from pydantic import AliasChoices, ConfigDict, Field, field_validator, model_validator + +from orb.domain.template.template_aggregate import Template +from orb.providers.azure.domain.template.value_objects import ( + AzureAllocationStrategy, + AzureCapacityReservationGroupId, + AzureDataDisk, + AzureDiskEncryptionSetId, + AzureEvictionPolicy, + AzureImageReference, + AzureLocationName, + AzureNetworkConfig, + AzureOSDiskConfig, + AzurePriority, + AzureProviderApi, + AzureProximityPlacementGroupId, + AzureResourceGroupName, + AzureSecurityType, + AzureUpgradePolicyMode, + AzureVmSizePreference, + AzureVMSSOrchestrationMode, +) +from orb.providers.azure.services.spot_placement_planner import PlacementSplitStrategy + + +class AzureTemplate(Template): + """Azure VMSS-specific template with full Azure compute extensions. + + Extends the provider-agnostic ``Template`` with every configuration + knob exposed by the Azure VMSS (Flexible orchestration) API so that + callers can declaratively describe the desired VM fleet. + + Note that defaults set to None will typically use the provider default in Config.py + """ + + model_config = ConfigDict( + arbitrary_types_allowed=True, + populate_by_name=True, + extra="forbid", + frozen=True, + validate_assignment=True, + ) + + # ------------------------------------------------------------------ + # Provider identification + # ------------------------------------------------------------------ + provider_api: AzureProviderApi = AzureProviderApi.VMSS + + # ------------------------------------------------------------------ + # Resource group & location + # ------------------------------------------------------------------ + resource_group: AzureResourceGroupName = Field( + ..., + description="Azure resource group for the VMSS.", + validation_alias=AliasChoices("resource_group", "resourceGroup"), + ) + location: AzureLocationName = Field( + ..., + description=( + "Azure location, e.g. 'eastus2'. Azure templates use the Azure-native " + "term ``location`` even though shared provider config uses ``region``." + ), + validation_alias=AliasChoices("location", "region"), + ) + subscription_id: Optional[str] = Field( + default=None, + description="Azure subscription ID (overrides default).", + validation_alias=AliasChoices("subscription_id", "subscriptionId"), + ) + + # ------------------------------------------------------------------ + # VMSS configuration + # ------------------------------------------------------------------ + vmss_name: Optional[str] = Field( + default=None, + description="Explicit VMSS name. Auto-generated if omitted.", + validation_alias=AliasChoices("vmss_name", "vmssName"), + ) + orchestration_mode: AzureVMSSOrchestrationMode = Field( + default=AzureVMSSOrchestrationMode.FLEXIBLE, + description="VMSS orchestration mode.", + validation_alias=AliasChoices("orchestration_mode", "orchestrationMode"), + ) + platform_fault_domain_count: Optional[int] = Field( + default=None, + ge=1, + le=5, + description="Fault domain count for Flexible orchestration.", + validation_alias=AliasChoices("platform_fault_domain_count", "platformFaultDomainCount"), + ) + single_placement_group: bool = Field( + default=False, + description="Restrict VMSS to a single placement group (max ~100 VMs).", + validation_alias=AliasChoices("single_placement_group", "singlePlacementGroup"), + ) + + # ------------------------------------------------------------------ + # VM size selection + # ------------------------------------------------------------------ + vm_size: str = Field( + default=..., + description="Primary Azure VM size, e.g. 'Standard_D4s_v5'.", + validation_alias=AliasChoices("vm_size", "vmSize"), + ) + vm_sizes: list[str] = Field( + default_factory=list, + description="Additional unranked Azure VM size candidates.", + validation_alias=AliasChoices("vm_sizes", "vmSizes"), + ) + vm_size_preferences: list[AzureVmSizePreference] = Field( + default_factory=list, + description=( + "Additional Azure VM size candidates for VMSS instance mix. " + "The primary vm_size remains the first choice." + ), + validation_alias=AliasChoices("vm_size_preferences", "vmSizePreferences"), + ) + + # ------------------------------------------------------------------ + # Image + # ------------------------------------------------------------------ + image: Optional[AzureImageReference] = Field( + default=None, + description="VM image reference (Marketplace or custom).", + ) + + # ------------------------------------------------------------------ + # Pricing / Spot + # ------------------------------------------------------------------ + priority: AzurePriority = Field( + default=AzurePriority.REGULAR, + description="VM priority (Regular, Spot, Low).", + ) + eviction_policy: Optional[AzureEvictionPolicy] = Field( + default=None, + description="Eviction policy for Spot VMs.", + validation_alias=AliasChoices("eviction_policy", "evictionPolicy"), + ) + billing_profile_max_price: Optional[float] = Field( + default=None, + description="Max price ($/hr) for Spot VMs. -1 = market price.", + validation_alias=AliasChoices("billing_profile_max_price", "billingProfileMaxPrice"), + ) + spot_percentage: Optional[int] = Field( + default=None, + ge=0, + le=100, + description=( + "Desired percentage of Spot VMs above the regular-priority base count. " + "Mapped to Azure Flexible VMSS priorityMixPolicy." + ), + validation_alias=AliasChoices("spot_percentage", "spotPercentage"), + ) + base_regular_priority_count: int = Field( + default=0, + ge=0, + description=( + "Minimum number of regular-priority VMs to keep when using Spot Priority Mix." + ), + validation_alias=AliasChoices("base_regular_priority_count", "baseRegularPriorityCount"), + ) + vmss_allocation_strategy: Optional[AzureAllocationStrategy] = Field( + default=None, + description=( + "Azure VMSS instance-mix allocation strategy " + "(LowestPrice, CapacityOptimized, Prioritized)." + ), + validation_alias=AliasChoices("vmss_allocation_strategy", "vmssAllocationStrategy"), + ) + spot_restore_enabled: bool = Field( + default=False, + description="Enable Spot Try-Restore to automatically re-create evicted instances.", + validation_alias=AliasChoices("spot_restore_enabled", "spotRestoreEnabled"), + ) + spot_restore_timeout: Optional[str] = Field( + default=None, + description="ISO 8601 duration for spot restore timeout, e.g. 'PT1H'.", + validation_alias=AliasChoices("spot_restore_timeout", "spotRestoreTimeout"), + ) + spot_placement_score_enabled: bool = Field( + default=False, + description="Enable Azure Spot Placement Score planning before launch.", + validation_alias=AliasChoices("spot_placement_score_enabled", "spotPlacementScoreEnabled"), + ) + placement_split_strategy: PlacementSplitStrategy = Field( + default=PlacementSplitStrategy.HYBRID, + description="How spot placement score launches split capacity across candidates.", + validation_alias=AliasChoices("placement_split_strategy", "placementSplitStrategy"), + ) + placement_primary_share_percent: int = Field( + default=80, + ge=0, + le=100, + description="Capacity percentage assigned to the top spot placement candidate.", + validation_alias=AliasChoices( + "placement_primary_share_percent", "placementPrimarySharePercent" + ), + ) + placement_regions: list[str] = Field( + default_factory=list, + description="Azure regions considered for spot placement score planning.", + validation_alias=AliasChoices("placement_regions", "placementRegions"), + ) + placement_zones: list[str] = Field( + default_factory=list, + description="Azure zones considered for spot placement score planning.", + validation_alias=AliasChoices("placement_zones", "placementZones"), + ) + + # ------------------------------------------------------------------ + # Availability + # ------------------------------------------------------------------ + zones: Optional[list[str]] = Field( + default=None, + description="Availability zones, e.g. ['1', '2', '3'].", + ) + zone_balance: bool = Field( + default=False, + description="Strictly balance instances across zones.", + validation_alias=AliasChoices("zone_balance", "zoneBalance"), + ) + proximity_placement_group_id: Optional[AzureProximityPlacementGroupId] = Field( + default=None, + description="ARM resource ID of a proximity placement group.", + validation_alias=AliasChoices("proximity_placement_group_id", "proximityPlacementGroupId"), + ) + capacity_reservation_group_id: Optional[AzureCapacityReservationGroupId] = Field( + default=None, + description="ARM resource ID of a capacity reservation group.", + validation_alias=AliasChoices( + "capacity_reservation_group_id", "capacityReservationGroupId" + ), + ) + + # ------------------------------------------------------------------ + # OS & storage + # ------------------------------------------------------------------ + os_disk: Optional[AzureOSDiskConfig] = Field( + default=None, + description="OS disk configuration.", + validation_alias=AliasChoices("os_disk", "osDisk"), + ) + data_disks: list[AzureDataDisk] = Field( + default_factory=list, + description="Additional data disks.", + validation_alias=AliasChoices("data_disks", "dataDisks"), + ) + + # ------------------------------------------------------------------ + # Networking + # ------------------------------------------------------------------ + network_config: Optional[AzureNetworkConfig] = Field( + default=None, + description="Primary NIC / subnet / NSG configuration.", + validation_alias=AliasChoices("network_config", "networkConfig"), + ) + + # ------------------------------------------------------------------ + # Security + # ------------------------------------------------------------------ + security_type: Optional[AzureSecurityType] = Field( + default=None, + description="VM security type (TrustedLaunch, ConfidentialVM).", + validation_alias=AliasChoices("security_type", "securityType"), + ) + secure_boot_enabled: Optional[bool] = Field( + default=None, + description="Enable UEFI Secure Boot (requires TrustedLaunch).", + validation_alias=AliasChoices("secure_boot_enabled", "secureBootEnabled"), + ) + vtpm_enabled: Optional[bool] = Field( + default=None, + description="Enable vTPM (requires TrustedLaunch).", + validation_alias=AliasChoices("vtpm_enabled", "vtpmEnabled"), + ) + encryption_at_host: Optional[bool] = Field( + default=None, + description="Enable host-based encryption for all disks.", + validation_alias=AliasChoices("encryption_at_host", "encryptionAtHost"), + ) + disk_encryption_set_id: Optional[AzureDiskEncryptionSetId] = Field( + default=None, + description="ARM resource ID of a disk encryption set (CMK).", + validation_alias=AliasChoices("disk_encryption_set_id", "diskEncryptionSetId"), + ) + + # ------------------------------------------------------------------ + # Identity & access + # ------------------------------------------------------------------ + admin_username: str = Field( + default="azureuser", + description="VM admin username.", + validation_alias=AliasChoices("admin_username", "adminUsername"), + ) + ssh_key_name: Optional[str] = Field( + default=None, + description=( + "Name of an Azure SSH Public Key resource " + "(Microsoft.Compute/sshPublicKeys) in the same resource group. " + "The handler resolves the actual key data at provisioning time. " + ), + validation_alias=AliasChoices("ssh_key_name", "sshKeyName"), + ) + ssh_public_keys: list[str] = Field( + default_factory=list, + description=( + "Inline SSH public key strings for Linux VMs. " + "Prefer ssh_key_name to reference an Azure-managed key instead." + ), + validation_alias=AliasChoices("ssh_public_keys", "sshPublicKeys"), + ) + user_assigned_identity_ids: list[str] = Field( + default_factory=list, + description="ARM resource IDs of user-assigned managed identities.", + validation_alias=AliasChoices("user_assigned_identity_ids", "userAssignedIdentityIds"), + ) + system_assigned_identity: bool = Field( + default=False, + description="Enable system-assigned managed identity.", + validation_alias=AliasChoices("system_assigned_identity", "systemAssignedIdentity"), + ) + + # ------------------------------------------------------------------ + # Bootstrapping + # ------------------------------------------------------------------ + custom_data: Optional[str] = Field( + default=None, + description="Base64-encoded custom data / cloud-init payload.", + validation_alias=AliasChoices("custom_data", "customData"), + ) + extension_profile: Optional[list[dict[str, Any]]] = Field( + default=None, + description="VMSS VM extension definitions (custom script, monitoring, etc.).", + validation_alias=AliasChoices("extension_profile", "extensionProfile"), + ) + + # ------------------------------------------------------------------ + # Scale & upgrade behaviour + # ------------------------------------------------------------------ + overprovision: bool = Field( + default=False, + description="Enable overprovisioning (VMSS creates extra VMs then removes surplus).", + ) + upgrade_policy_mode: AzureUpgradePolicyMode = Field( + default=AzureUpgradePolicyMode.MANUAL, + description="Upgrade policy: Manual, Rolling, or Automatic.", + validation_alias=AliasChoices("upgrade_policy_mode", "upgradePolicyMode"), + ) + + # ------------------------------------------------------------------ + # Freeform pass-through + # ------------------------------------------------------------------ + provider_api_spec: Optional[dict[str, Any]] = Field( + default=None, + description="Optional raw Azure provider request payload override/overlay.", + validation_alias=AliasChoices("provider_api_spec", "providerApiSpec"), + ) + provider_api_spec_file: Optional[str] = Field( + default=None, + description="Path to a JSON native spec file for Azure provider request payloads.", + validation_alias=AliasChoices("provider_api_spec_file", "providerApiSpecFile"), + ) + node_attributes: dict[str, Any] = Field( + default_factory=dict, + description="Freeform pass-through properties appended to the VMSS ARM payload.", + validation_alias=AliasChoices("node_attributes", "nodeAttributes"), + ) + + # ------------------------------------------------------------------ + # CycleCloud configuration + # ------------------------------------------------------------------ + cluster_name: Optional[str] = Field( + default=None, + description="CycleCloud cluster name to add nodes to.", + validation_alias=AliasChoices("cluster_name", "clusterName"), + ) + node_array: str = Field( + default="execute", + description="CycleCloud node array (partition) to target, e.g. 'execute', 'hpc', 'htc'.", + validation_alias=AliasChoices("node_array", "nodeArray"), + ) + cyclecloud_url: Optional[str] = Field( + default=None, + description="CycleCloud REST API base URL, e.g. 'https://cyclecloud.example.com'.", + validation_alias=AliasChoices("cyclecloud_url", "cyclecloudUrl"), + ) + cyclecloud_credential_path: Optional[str] = Field( + default=None, + description="Path to a JSON file containing CycleCloud credentials and optional auth overrides.", + validation_alias=AliasChoices("cyclecloud_credential_path", "cyclecloudCredentialPath"), + ) + cyclecloud_verify_ssl: Optional[bool] = Field( + default=None, + description="Whether to verify SSL certificates for CycleCloud API calls.", + validation_alias=AliasChoices("cyclecloud_verify_ssl", "cyclecloudVerifySsl"), + ) + # TODO: Bearer functionality not tested in live environment + cyclecloud_auth_mode: Optional[str] = Field( + default=None, + description="CycleCloud auth mode override, e.g. 'basic' or 'bearer'.", + validation_alias=AliasChoices("cyclecloud_auth_mode", "cyclecloudAuthMode"), + ) + cyclecloud_aad_scope: Optional[str] = Field( + default=None, + description="AAD scope used to resolve a bearer token for CycleCloud.", + validation_alias=AliasChoices("cyclecloud_aad_scope", "cyclecloudAadScope"), + ) + version: Optional[str] = Field( + default=None, + description="Template catalog or DTO version metadata.", + ) + + # ------------------------------------------------------------------ + # Initialisation + # ------------------------------------------------------------------ + + def __init__(self, **data: Any) -> None: + """Initialise the Azure template and set provider_type.""" + data["provider_type"] = "azure" + super().__init__(**data) + + # ------------------------------------------------------------------ + # Validators + # ------------------------------------------------------------------ + + @field_validator("vm_sizes", mode="before") + @classmethod + def normalize_vm_sizes(cls, value: Any) -> Any: + """Accept HostFactory ``vmTypes`` dictionaries as Azure VM size lists.""" + if isinstance(value, dict): + return [str(vm_size) for vm_size in value.keys()] + return value + + @model_validator(mode="before") + @classmethod + def apply_implied_defaults(cls, data: Any) -> Any: + """Normalise input data before construction. + + This is input normalisation, not mutation of a constructed model. + The ``mode="after"`` validator below is purely rejecting — it never + modifies state. Conditional defaults that depend on other field + values belong here so that every construction path (strategy, + factory, test helpers) gets consistent behaviour. + + Azure templates use ``location`` as the canonical field because this + is the Azure platform term. We still accept ``region`` at this input + boundary because provider-level config and some shared call paths use + the cross-provider ``region`` name. + """ + if not isinstance(data, dict): + return data + + data = dict(data) + raw_provider_config = data.pop("provider_config", None) + if raw_provider_config is not None: + if hasattr(raw_provider_config, "model_dump"): + provider_config = raw_provider_config.model_dump(mode="json") + elif isinstance(raw_provider_config, dict): + provider_config = dict(raw_provider_config) + else: + provider_config = {} + for key, value in provider_config.items(): + if value is not None and data.get(key) in (None, ""): + data[key] = value + + location = data.get("location") + region = data.get("region") + if location not in (None, "") and region not in (None, "") and location != region: + raise ValueError("Azure templates received conflicting 'location' and 'region' values") + if location in (None, "") and region not in (None, ""): + data["location"] = region + data.pop("region", None) + + if "max_number" in data: + if "max_instances" not in data: + data["max_instances"] = data["max_number"] + data.pop("max_number", None) + + if data.get("allocation_strategy") not in (None, ""): + raise ValueError( + "Azure templates do not support allocation_strategy; use " + "vmss_allocation_strategy for Azure VMSS instance mix or " + "spot_placement_score_enabled for Spot Placement Score planning" + ) + # The shared Template base still owns an AWS-shaped allocation_strategy + # field and mutates it when absent. Azure does not use that contract, so + # set an inert value here after rejecting explicit non-empty input. + data["allocation_strategy"] = "" + + if data.get("spot_percentage") is not None and data.get("priority") in ( + None, + "", + "Regular", + ): + data["priority"] = "Spot" + + # Spot VMs need an eviction policy. + if data.get("priority") == "Spot": + data.setdefault("eviction_policy", "Deallocate") + + # Trusted Launch implies secure boot and vTPM. + if data.get("security_type") == "TrustedLaunch": + data.setdefault("secure_boot_enabled", True) + data.setdefault("vtpm_enabled", True) + + return data + + @model_validator(mode="after") + def validate_azure_template(self) -> "AzureTemplate": + """Azure-specific template validation.""" + if self.provider_api_spec and self.provider_api_spec_file: + raise ValueError("Cannot specify both provider_api_spec and provider_api_spec_file") + + if self.provider_api == AzureProviderApi.VMSS_UNIFORM: + if self.orchestration_mode != AzureVMSSOrchestrationMode.UNIFORM: + raise ValueError("provider_api 'VMSSUniform' requires orchestration_mode 'Uniform'") + + if self.spot_placement_score_enabled: + candidate_sizes = self.candidate_vm_sizes + if len(candidate_sizes) < 2: + raise ValueError( + "spot_placement_score_enabled requires at least two candidate vm sizes" + ) + + if self.vm_sizes and self.vm_size_preferences: + raise ValueError( + "Specify either vm_sizes or vm_size_preferences for Azure instance mix, not both" + ) + + if self.vm_sizes: + if self.vm_size in self.vm_sizes: + raise ValueError("vm_sizes must not repeat the primary vm_size") + if len(self.vm_sizes) != len(set(self.vm_sizes)): + raise ValueError("vm_sizes must contain unique VM sizes") + + preference_names = [preference.name for preference in self.vm_size_preferences] + explicit_ranks = [ + preference.rank + for preference in self.vm_size_preferences + if preference.rank is not None + ] + + if self.vm_size in preference_names: + raise ValueError("vm_size_preferences must not repeat the primary vm_size") + + if len(preference_names) != len(set(preference_names)): + raise ValueError("vm_size_preferences must contain unique VM sizes") + + if ( + self.vm_size_preferences + and self.vmss_allocation_strategy != AzureAllocationStrategy.PRIORITIZED + ): + raise ValueError( + "vm_size_preferences is only valid with vmss_allocation_strategy='Prioritized'; " + "use vm_sizes for unranked Azure instance mix" + ) + + if ( + self.vmss_allocation_strategy == AzureAllocationStrategy.PRIORITIZED + and not self.vm_size_preferences + ): + raise ValueError( + "vmss_allocation_strategy='Prioritized' requires at least one vm_size_preference" + ) + + if self.vmss_allocation_strategy != AzureAllocationStrategy.PRIORITIZED: + if explicit_ranks: + raise ValueError( + "vm_size_preferences ranks are only valid with vmss_allocation_strategy='Prioritized'" + ) + + if self.vmss_allocation_strategy == AzureAllocationStrategy.PRIORITIZED: + if any(rank == 0 for rank in explicit_ranks): + raise ValueError( + "vm_size_preferences ranks must start at 1; primary vm_size is rank 0" + ) + if len(explicit_ranks) != len(set(explicit_ranks)): + raise ValueError("vm_size_preferences ranks must be unique") + + if self.spot_percentage is not None: + if self.provider_api not in ( + AzureProviderApi.VMSS, + AzureProviderApi.VMSS_UNIFORM, + ): + raise ValueError("spot_percentage is only supported for VMSS-based Azure templates") + if self.orchestration_mode != AzureVMSSOrchestrationMode.FLEXIBLE: + raise ValueError("spot_percentage requires Flexible orchestration mode") + if self.single_placement_group: + raise ValueError( + "spot_percentage is not supported when single_placement_group is enabled" + ) + # Azure requires Spot priority on the scale set when using priorityMixPolicy. + if self.priority == AzurePriority.LOW: + raise ValueError( + "spot_percentage is not compatible with Low priority VMs; use Spot" + ) + if self.priority != AzurePriority.SPOT: + raise ValueError("spot_percentage requires priority='Spot'") + + # Spot VMs require an eviction policy. + if self.priority == AzurePriority.SPOT: + if self.eviction_policy is None: + raise ValueError("eviction_policy is required for Spot priority VMs") + + # Non-spot VMs should not have spot-specific settings + if self.priority == AzurePriority.REGULAR and self.eviction_policy is not None: + raise ValueError("eviction_policy is only valid for Spot or Low priority VMs") + if self.priority != AzurePriority.SPOT and self.billing_profile_max_price is not None: + raise ValueError("billing_profile_max_price is only valid for Spot priority VMs") + + # Zone balance requires zones + if self.zone_balance and not self.zones: + raise ValueError("zone_balance requires at least one availability zone") + + # overprovision is only valid for Uniform orchestration + if self.overprovision and self.orchestration_mode != AzureVMSSOrchestrationMode.UNIFORM: + raise ValueError("overprovision is only valid for Uniform orchestration mode") + + # Trusted Launch validation + if self.security_type == AzureSecurityType.TRUSTED_LAUNCH: + if self.secure_boot_enabled is None: + raise ValueError( + "secure_boot_enabled is required when security_type is TrustedLaunch" + ) + if self.vtpm_enabled is None: + raise ValueError("vtpm_enabled is required when security_type is TrustedLaunch") + + # Capacity reservation and Spot are mutually exclusive + if self.capacity_reservation_group_id and self.priority == AzurePriority.SPOT: + raise ValueError("Capacity reservations cannot be used with Spot priority VMs") + + # SSH access is required for Linux VMs — never fall back to + # generating a random admin password. Callers must provide either + # ssh_key_name (a reference to an Azure SSH Public Key resource) or inline ssh_public_keys. + # CycleCloud manages SSH access internally so this is not required. + if self.provider_api != AzureProviderApi.CYCLECLOUD: + if self.image is None and self.image_id is None: + raise ValueError( + "An Azure image source is required. Provide either 'image' or 'image_id'." + ) + if not self.ssh_key_name and not self.ssh_public_keys: + raise ValueError( + "SSH access is required for Azure Linux VMs. Provide either " + "'ssh_key_name' (name of an Azure SSH Public Key resource) " + "or 'ssh_public_keys' (inline key data). " + "Password-based authentication is not supported." + ) + + # CycleCloud-specific validation + if self.provider_api == AzureProviderApi.CYCLECLOUD: + if not self.cluster_name: + raise ValueError( + "cluster_name is required for CycleCloud templates. " + "Specify the name of an existing CycleCloud cluster." + ) + + return self + + @property + def candidate_vm_sizes(self) -> list[str]: + """Return the ordered VM sizes considered for provisioning.""" + if self.vm_sizes: + return [self.vm_size, *self.vm_sizes] + return [self.vm_size, *[preference.name for preference in self.vm_size_preferences]] + + @property + def uses_vm_size_mix(self) -> bool: + """Return whether the template allows more than one VM size.""" + return bool(self.vm_sizes or self.vm_size_preferences) + + def build_vm_size_profile(self) -> list[dict[str, Any]]: + """Build Azure skuProfile.vmSizes entries for instance mix.""" + if not self.uses_vm_size_mix: + return [] + + is_prioritized = self.vmss_allocation_strategy == AzureAllocationStrategy.PRIORITIZED + + if self.vm_sizes: + return [{"name": vm_size} for vm_size in [self.vm_size, *self.vm_sizes]] + + entries: list[dict[str, Any]] = [] + primary_entry: dict[str, Any] = {"name": self.vm_size} + if is_prioritized: + primary_entry["rank"] = 0 + entries.append(primary_entry) + + used_ranks = {0} if is_prioritized else set() + next_rank = 1 + for preference in self.vm_size_preferences: + entry: dict[str, Any] = {"name": preference.name} + if is_prioritized: + if preference.rank is not None: + entry["rank"] = preference.rank + used_ranks.add(preference.rank) + entries.append(entry) + + if is_prioritized: + used_ranks = {entry["rank"] for entry in entries if "rank" in entry} + for entry in entries[1:]: + if "rank" in entry: + continue + while next_rank in used_ranks: + next_rank += 1 + entry["rank"] = next_rank + used_ranks.add(next_rank) + next_rank += 1 + + return entries + + @classmethod + def from_azure_format(cls, data: dict[str, Any]) -> "AzureTemplate": + """Create an AzureTemplate from a flat configuration dict. + + Accepts both snake_case and camelCase keys (via AliasChoices). + """ + return cls.model_validate(data) + + def __str__(self) -> str: + """Return string representation.""" + return ( + f"AzureTemplate(id={self.template_id}, vm_size={self.vm_size}, " + f"location={self.location}, priority={self.priority.value}, " + f"instances={self.max_instances})" + ) + + def __repr__(self) -> str: + """Detailed string representation.""" + return ( + f"AzureTemplate(template_id='{self.template_id}', " + f"vm_size='{self.vm_size}', location='{self.location}', " + f"resource_group='{self.resource_group}', " + f"priority='{self.priority.value}', " + f"orchestration_mode='{self.orchestration_mode.value}', " + f"max_instances={self.max_instances})" + ) diff --git a/src/orb/providers/azure/domain/template/value_objects.py b/src/orb/providers/azure/domain/template/value_objects.py new file mode 100644 index 000000000..3df745a26 --- /dev/null +++ b/src/orb/providers/azure/domain/template/value_objects.py @@ -0,0 +1,460 @@ +"""Azure-specific value objects and domain primitives for VMSS-based provisioning.""" + +import re +from abc import ABC +from enum import Enum +from typing import Any, ClassVar, Optional + +from pydantic import ( + ConfigDict, + Field, + ValidationInfo, + field_validator, + model_serializer, + model_validator, +) + +from orb.domain.base.value_objects import PriceType, ValueObject + +# --------------------------------------------------------------------------- +# Enumerations +# --------------------------------------------------------------------------- + + +class AzureProviderApi(str, Enum): + """Azure compute provisioning API types. + + Maps to the Azure SDK / ARM API surface used to create VMs. + """ + + VMSS = "VMSS" # Virtual Machine Scale Set (Flexible orchestration) + VMSS_UNIFORM = "VMSSUniform" # VMSS Uniform orchestration (legacy) + SINGLE_VM = "SingleVM" # Individual VM via compute API + CYCLECLOUD = "CycleCloud" # Azure CycleCloud cluster node management + + +class AzureVMSSOrchestrationMode(str, Enum): + """VMSS orchestration mode. + + See: https://learn.microsoft.com/en-us/azure/virtual-machine-scale-sets/ + virtual-machine-scale-sets-orchestration-modes + """ + + FLEXIBLE = "Flexible" + UNIFORM = "Uniform" + + +class AzurePriority(str, Enum): + """VM priority (maps to core PriceType). + + See: https://learn.microsoft.com/en-us/azure/virtual-machines/spot-vms + """ + + REGULAR = "Regular" + SPOT = "Spot" + LOW = "Low" # Legacy low-priority (batch workloads) + + @classmethod + def from_price_type(cls, price_type: PriceType) -> "AzurePriority": + """Convert core PriceType to Azure priority.""" + mapping = { + PriceType.ONDEMAND: cls.REGULAR, + PriceType.SPOT: cls.SPOT, + PriceType.RESERVED: cls.REGULAR, # Reserved still uses Regular VMs + PriceType.HETEROGENEOUS: cls.REGULAR, + } + return mapping.get(price_type, cls.REGULAR) + + +class AzureEvictionPolicy(str, Enum): + """Eviction policy for Spot / low-priority VMs. + + See: https://learn.microsoft.com/en-us/azure/virtual-machines/spot-vms#eviction-policy + """ + + DEALLOCATE = "Deallocate" + DELETE = "Delete" + + +class AzureAllocationStrategy(str, Enum): + """VMSS allocation / placement strategy for instance mix. + + See: https://learn.microsoft.com/en-us/azure/virtual-machine-scale-sets/ + instance-mix-overview + """ + + LOWEST_PRICE = "LowestPrice" + CAPACITY_OPTIMIZED = "CapacityOptimized" + PRIORITIZED = "Prioritized" + + def to_arm_value(self) -> str: + """Return the value expected by the ARM / REST API.""" + return self.value + + +class AzureVmSizePreference(ValueObject): + """Additional VM size candidate for Azure instance mix.""" + + name: str = Field(..., description="Azure VM size name, e.g. Standard_D8s_v5.") + rank: Optional[int] = Field( + default=None, + ge=1, + description=( + "Preferred order for Azure Prioritized instance mix. " + "Lower values are preferred after the primary vm_size." + ), + ) + + @field_validator("name") + @classmethod + def validate_name(cls, value: str) -> str: + """Reject empty VM size names.""" + normalised = value.strip() + if not normalised: + raise ValueError("AzureVmSizePreference.name cannot be empty") + return normalised + + +class AzureOSDiskType(str, Enum): + """Managed disk storage account type. + + See: https://learn.microsoft.com/en-us/azure/virtual-machines/disks-types + """ + + STANDARD_LRS = "Standard_LRS" # Standard HDD + STANDARD_SSD_LRS = "StandardSSD_LRS" # Standard SSD + PREMIUM_LRS = "Premium_LRS" # Premium SSD + PREMIUM_SSD_V2_LRS = "PremiumV2_LRS" # Premium SSD v2 + ULTRA_SSD_LRS = "UltraSSD_LRS" # Ultra Disk + + +class AzureSecurityType(str, Enum): + """VM security type. + + See: https://learn.microsoft.com/en-us/azure/virtual-machines/trusted-launch + """ + + STANDARD = "Standard" + TRUSTED_LAUNCH = "TrustedLaunch" + CONFIDENTIAL_VM = "ConfidentialVM" + + +class AzureCachingType(str, Enum): + """OS / data disk caching type.""" + + NONE = "None" + READ_ONLY = "ReadOnly" + READ_WRITE = "ReadWrite" + + +class AzureUpgradePolicyMode(str, Enum): + """VMSS upgrade policy mode.""" + + MANUAL = "Manual" + ROLLING = "Rolling" + AUTOMATIC = "Automatic" + + +class AzureStringValue(ValueObject, ABC): + """Base class for Azure scalar string value objects.""" + + value: str + + @model_validator(mode="before") + @classmethod + def coerce_from_string(cls, data: object) -> object: + """Accept raw strings at the model boundary.""" + if isinstance(data, str): + return {"value": data} + return data + + @field_validator("value") + @classmethod + def validate_value(cls, value: str) -> str: + """Validate and normalise scalar string values.""" + if not isinstance(value, str): + raise TypeError(f"{cls.__name__} must be a string") + + normalised = value.strip() + if not normalised: + raise ValueError(f"{cls.__name__} cannot be empty") + + cls._validate_normalised(normalised) + return normalised + + @classmethod + def _validate_normalised(cls, value: str) -> None: + """Subclass hook for value-specific validation.""" + + @model_serializer + def serialise_model(self) -> str: + """Serialise as the raw string value.""" + return self.value + + def __str__(self) -> str: + return self.value + + +class AzureResourceGroupName(AzureStringValue): + """Azure resource-group name.""" + + _RESOURCE_GROUP_RE: ClassVar[re.Pattern[str]] = re.compile(r"^[a-zA-Z0-9_\-.()\[\]]+$") + + @classmethod + def _validate_normalised(cls, value: str) -> None: + if not (1 <= len(value) <= 90): + raise ValueError("resource_group must be 1-90 characters") + if not cls._RESOURCE_GROUP_RE.match(value): + raise ValueError( + "resource_group contains invalid characters " + "(allowed: alphanumeric, _, -, ., (, ), [, ])" + ) + + +class AzureLocationName(AzureStringValue): + """Azure location slug used by the compute APIs.""" + + _LOCATION_RE: ClassVar[re.Pattern[str]] = re.compile(r"^[A-Za-z0-9][A-Za-z0-9-]*$") + + @classmethod + def _validate_normalised(cls, value: str) -> None: + if not cls._LOCATION_RE.match(value): + raise ValueError( + "location must be an Azure location slug like 'eastus2' or 'westeurope'" + ) + + +class AzureArmResourceId(AzureStringValue): + """Base class for Azure ARM resource identifiers.""" + + expected_segment: ClassVar[Optional[str]] = None + + @classmethod + def _validate_normalised(cls, value: str) -> None: + """Validate Azure ARM resource ID format.""" + casefolded_value = value.casefold() + if not casefolded_value.startswith("/subscriptions/"): + raise ValueError(f"{cls.__name__} must be a full ARM resource ID") + if "/providers/" not in casefolded_value: + raise ValueError(f"{cls.__name__} must include an Azure provider segment") + + expected_segment = cls.expected_segment.casefold() if cls.expected_segment else None + if expected_segment and expected_segment not in casefolded_value: + raise ValueError(f"{cls.__name__} must reference '{cls.expected_segment}'") + + +class AzureProximityPlacementGroupId(AzureArmResourceId): + """ARM resource ID for an Azure proximity placement group.""" + + expected_segment = "/providers/Microsoft.Compute/proximityPlacementGroups/" + + +class AzureCapacityReservationGroupId(AzureArmResourceId): + """ARM resource ID for an Azure capacity reservation group.""" + + expected_segment = "/providers/Microsoft.Compute/capacityReservationGroups/" + + +class AzureDiskEncryptionSetId(AzureArmResourceId): + """ARM resource ID for an Azure disk encryption set.""" + + expected_segment = "/providers/Microsoft.Compute/diskEncryptionSets/" + + +# --------------------------------------------------------------------------- +# Value Objects +# --------------------------------------------------------------------------- + + +class AzureImageReference(ValueObject): + """Azure VM image reference - either a Marketplace image or a custom/gallery image. + + See: https://learn.microsoft.com/en-us/azure/virtual-machines/linux/cli-ps-findimage + """ + + model_config = ConfigDict(populate_by_name=True) + + publisher: Optional[str] = None # e.g. "Canonical" + offer: Optional[str] = None # e.g. "0001-com-ubuntu-server-jammy" + sku: Optional[str] = None # e.g. "22_04-lts-gen2" + version: str = "latest" + + # For custom / Shared Image Gallery / Community Gallery images + image_id: Optional[str] = Field( + default=None, + description="Full resource ID of a custom image, shared image gallery image, " + "or community gallery image.", + ) + + @model_validator(mode="after") + def _validate_image_source(self) -> "AzureImageReference": + """Ensure either marketplace fields or image_id is provided.""" + has_marketplace = self.publisher and self.offer and self.sku + if not has_marketplace and not self.image_id: + raise ValueError( + "Provide either (publisher, offer, sku) for a Marketplace image " + "or image_id for a custom / gallery image." + ) + if has_marketplace and self.image_id: + raise ValueError( + "Cannot specify both Marketplace image fields (publisher/offer/sku) " + "and image_id at the same time." + ) + return self + + def to_arm_dict(self) -> dict[str, Any]: + """Serialise to the ARM imageReference format.""" + if self.image_id: + return {"id": self.image_id} + return { + "publisher": self.publisher, + "offer": self.offer, + "sku": self.sku, + "version": self.version, + } + + +class AzureNetworkConfig(ValueObject): + """Networking configuration for VMSS instances.""" + + model_config = ConfigDict(populate_by_name=True) + + subnet_id: str = Field( + ..., + description="Full ARM resource ID of the subnet.", + ) + network_security_group_id: Optional[str] = Field( + default=None, + description="Full ARM resource ID of the NSG to attach.", + ) + accelerated_networking: Optional[bool] = Field( + default=None, + description="Enable accelerated networking (SR-IOV).", + ) + public_ip_enabled: bool = False + load_balancer_backend_pool_ids: list[str] = Field( + default_factory=list, + description="Existing Azure Load Balancer backend address pool ARM IDs to attach.", + ) + load_balancer_inbound_nat_pool_ids: list[str] = Field( + default_factory=list, + description="Existing Azure Load Balancer inbound NAT pool ARM IDs to attach.", + ) + application_gateway_backend_pool_ids: list[str] = Field( + default_factory=list, + description="Existing Application Gateway backend pool ARM IDs to attach.", + ) + + def to_arm_dict(self) -> dict[str, Any]: + """Serialise to the ARM networkInterfaceConfiguration format.""" + ip_config_properties: dict[str, Any] = {"subnet": {"id": self.subnet_id}} + if self.load_balancer_backend_pool_ids: + ip_config_properties["loadBalancerBackendAddressPools"] = [ + {"id": pool_id} for pool_id in self.load_balancer_backend_pool_ids + ] + if self.load_balancer_inbound_nat_pool_ids: + ip_config_properties["loadBalancerInboundNatPools"] = [ + {"id": pool_id} for pool_id in self.load_balancer_inbound_nat_pool_ids + ] + if self.application_gateway_backend_pool_ids: + ip_config_properties["applicationGatewayBackendAddressPools"] = [ + {"id": pool_id} for pool_id in self.application_gateway_backend_pool_ids + ] + + ip_config: dict[str, Any] = { + "name": "ipconfig1", + "properties": ip_config_properties, + } + if self.public_ip_enabled: + ip_config["properties"]["publicIPAddressConfiguration"] = { + "name": "publicip", + "properties": {"deleteOption": "Delete"}, + } + + nic: dict[str, Any] = { + "name": "nic-config", + "properties": { + "deleteOption": "Delete", + "primary": True, + "ipConfigurations": [ip_config], + }, + } + if self.accelerated_networking is not None: + nic["properties"]["enableAcceleratedNetworking"] = self.accelerated_networking + if self.network_security_group_id: + nic["properties"]["networkSecurityGroup"] = { + "id": self.network_security_group_id, + } + return nic + + +class AzureOSDiskConfig(ValueObject): + """OS disk configuration for VMSS instances.""" + + model_config = ConfigDict(populate_by_name=True) + + disk_size_gb: Optional[int] = None + storage_account_type: AzureOSDiskType = AzureOSDiskType.PREMIUM_LRS + caching: AzureCachingType = AzureCachingType.READ_WRITE + ephemeral_os_disk: bool = False + ephemeral_placement: Optional[str] = Field( + default=None, + validate_default=True, + description="Placement for ephemeral OS disk: 'CacheDisk' or 'ResourceDisk'.", + ) + + @field_validator("ephemeral_placement", mode="before") + @classmethod + def _apply_ephemeral_defaults(cls, value: Optional[str], info: ValidationInfo) -> Optional[str]: + if value is None and info.data.get("ephemeral_os_disk"): + return "CacheDisk" + return value + + @model_validator(mode="after") + def _validate_ephemeral(self) -> "AzureOSDiskConfig": + if not self.ephemeral_os_disk and self.ephemeral_placement is not None: + raise ValueError("ephemeral_placement requires ephemeral_os_disk=True") + return self + + def to_arm_dict(self) -> dict[str, Any]: + """Serialise to the ARM osDisk format.""" + disk: dict[str, Any] = { + "createOption": "FromImage", + "deleteOption": "Delete", + "caching": self.caching.value, + "managedDisk": { + "storageAccountType": self.storage_account_type.value, + }, + } + if self.disk_size_gb: + disk["diskSizeGB"] = self.disk_size_gb + if self.ephemeral_os_disk: + disk["diffDiskSettings"] = { + "option": "Local", + "placement": self.ephemeral_placement, + } + return disk + + +class AzureDataDisk(ValueObject): + """A single data disk specification.""" + + model_config = ConfigDict(populate_by_name=True) + + lun: int = Field(..., ge=0, le=63) + disk_size_gb: int + storage_account_type: AzureOSDiskType = AzureOSDiskType.PREMIUM_LRS + caching: AzureCachingType = AzureCachingType.NONE + + def to_arm_dict(self) -> dict[str, Any]: + """Serialise to the ARM dataDisk format.""" + return { + "lun": self.lun, + "createOption": "Empty", + "deleteOption": "Delete", + "diskSizeGB": self.disk_size_gb, + "caching": self.caching.value, + "managedDisk": { + "storageAccountType": self.storage_account_type.value, + }, + } diff --git a/src/orb/providers/azure/exceptions/__init__.py b/src/orb/providers/azure/exceptions/__init__.py new file mode 100644 index 000000000..d80859ef8 --- /dev/null +++ b/src/orb/providers/azure/exceptions/__init__.py @@ -0,0 +1,59 @@ +"""Azure exceptions package.""" + +from orb.providers.azure.exceptions.azure_exceptions import ( + AuthenticationError, + AuthorizationError, + AzureConfigurationError, + AzureEntityNotFoundError, + AzureError, + AzureInfrastructureError, + AzureValidationError, + CycleCloudClusterNotFoundError, + CycleCloudConnectionError, + CycleCloudError, + CycleCloudNodeError, + ImageValidationError, + LaunchError, + NetworkError, + QuotaExceededError, + RateLimitError, + ResourceCleanupError, + ResourceInUseError, + ResourceStateError, + SubnetValidationError, + TaggingError, + TerminationError, + VMNotFoundError, + VMSizeValidationError, + VMSSCreationError, + VMSSNotFoundError, +) + +__all__: list[str] = [ + "AuthenticationError", + "AuthorizationError", + "AzureConfigurationError", + "AzureEntityNotFoundError", + "AzureError", + "AzureInfrastructureError", + "AzureValidationError", + "CycleCloudClusterNotFoundError", + "CycleCloudConnectionError", + "CycleCloudError", + "CycleCloudNodeError", + "ImageValidationError", + "LaunchError", + "NetworkError", + "QuotaExceededError", + "RateLimitError", + "ResourceCleanupError", + "ResourceInUseError", + "ResourceStateError", + "SubnetValidationError", + "TaggingError", + "TerminationError", + "VMNotFoundError", + "VMSizeValidationError", + "VMSSCreationError", + "VMSSNotFoundError", +] diff --git a/src/orb/providers/azure/exceptions/azure_exceptions.py b/src/orb/providers/azure/exceptions/azure_exceptions.py new file mode 100644 index 000000000..f51737885 --- /dev/null +++ b/src/orb/providers/azure/exceptions/azure_exceptions.py @@ -0,0 +1,409 @@ +"""Azure-specific exceptions. + +Provides a full exception hierarchy for Azure provider operations +""" + +from typing import Any, Optional + +from orb.providers.base.exceptions import ( + ProviderAuthError, + ProviderConfigError, + ProviderError, + ProviderPermanentError, + ProviderQuotaError, + ProviderTransientError, +) + + +class AzureError(ProviderError): + """Base class for Azure-related errors.""" + + def __init__( + self, + message: str, + details: Optional[dict[str, Any]] = None, + error_code: Optional[str] = None, + *, + provider_name: Optional[str] = None, + underlying_exception: Optional[BaseException] = None, + is_retryable: Optional[bool] = None, + ) -> None: + super().__init__( + message, + provider_type="azure", + provider_name=provider_name, + underlying_exception=underlying_exception, + details=details, + is_retryable=is_retryable, + ) + self.error_code = error_code or self.__class__.__name__ + + def to_dict(self) -> dict[str, Any]: + """Serialize the exception to a dict, including error_code if non-default.""" + result: dict[str, Any] = super().to_dict() + if self.error_code and self.error_code != self.__class__.__name__: + result["error_code"] = self.error_code + return result + + def safe_to_dict(self) -> dict[str, Any]: + """Serialize the exception without exposing its underlying exception.""" + result = super().safe_to_dict() + if self.error_code and self.error_code != self.__class__.__name__: + result["error_code"] = self.error_code + return result + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +class AzureValidationError(AzureError, ProviderPermanentError): + """Raised when Azure resource validation fails.""" + + +class VMSizeValidationError(AzureValidationError): + """Invalid or unsupported VM size.""" + + def __init__(self, message: str, vm_size: str, details: Optional[dict[str, Any]] = None): + super().__init__(message, details={"vm_size": vm_size, **(details or {})}) + self.vm_size = vm_size + + +class SubnetValidationError(AzureValidationError): + """Invalid subnet configuration.""" + + def __init__(self, message: str, subnet_id: str, details: Optional[dict[str, Any]] = None): + super().__init__(message, details={"subnet_id": subnet_id, **(details or {})}) + self.subnet_id = subnet_id + + +class ImageValidationError(AzureValidationError): + """Invalid image reference.""" + + def __init__(self, message: str, image_ref: str, details: Optional[dict[str, Any]] = None): + super().__init__(message, details={"image_ref": image_ref, **(details or {})}) + self.image_ref = image_ref + + +# --------------------------------------------------------------------------- +# Entity not found +# --------------------------------------------------------------------------- + + +class AzureEntityNotFoundError(AzureError, ProviderPermanentError): + """Raised when an Azure resource is not found.""" + + +class VMNotFoundError(AzureEntityNotFoundError): + """A specific VM instance was not found.""" + + def __init__(self, message: str, instance_id: str, details: Optional[dict[str, Any]] = None): + super().__init__(message, details={"instance_id": instance_id, **(details or {})}) + self.instance_id = instance_id + + +class VMSSNotFoundError(AzureEntityNotFoundError): + """A VMSS resource was not found.""" + + def __init__(self, message: str, vmss_name: str, details: Optional[dict[str, Any]] = None): + super().__init__(message, details={"vmss_name": vmss_name, **(details or {})}) + self.vmss_name = vmss_name + + +# --------------------------------------------------------------------------- +# Quotas & capacity +# --------------------------------------------------------------------------- + + +class QuotaExceededError(AzureError, ProviderQuotaError): + """Raised when Azure service quotas would be exceeded.""" + + +class ServiceQuotaError(QuotaExceededError): + """Detailed quota error with specific service information.""" + + def __init__( + self, + message: str, + service: str, + quota_name: str, + current_value: int, + quota_value: int, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={ + "service": service, + "quota_name": quota_name, + "current_value": current_value, + "quota_value": quota_value, + **(details or {}), + }, + ) + self.service = service + self.quota_name = quota_name + self.current_value = current_value + self.quota_value = quota_value + + +# --------------------------------------------------------------------------- +# Resource state +# --------------------------------------------------------------------------- + + +class ResourceInUseError(AzureError, ProviderTransientError): + """Raised when an Azure resource is already in use.""" + + +class ResourceStateError(AzureError, ProviderTransientError): + """Raised when a resource is in an invalid state for the operation.""" + + def __init__( + self, + message: str, + resource_id: str, + current_state: str, + expected_states: list[str], + details: Optional[dict[str, Any]] = None, + ) -> None: + super().__init__( + message, + details={ + "resource_id": resource_id, + "current_state": current_state, + "expected_states": expected_states, + **(details or {}), + }, + ) + self.resource_id = resource_id + self.current_state = current_state + self.expected_states = expected_states + + +# --------------------------------------------------------------------------- +# Authentication & authorization +# --------------------------------------------------------------------------- + + +class AuthenticationError(AzureError, ProviderAuthError): + """Raised when Azure authentication fails.""" + + +class AuthorizationError(AzureError, ProviderAuthError): + """Raised when there are insufficient permissions.""" + + +class RBACError(AuthorizationError): + """Insufficient RBAC permissions for the operation.""" + + def __init__( + self, + message: str, + role: Optional[str] = None, + scope: Optional[str] = None, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={"role": role, "scope": scope, **(details or {})}, + ) + self.role = role + self.scope = scope + + +# --------------------------------------------------------------------------- +# Network & rate limiting +# --------------------------------------------------------------------------- + + +class RateLimitError(AzureError, ProviderQuotaError): + """Raised when Azure API rate limits are exceeded.""" + + +class NetworkError(AzureError, ProviderTransientError): + """Raised when there are network-related issues.""" + + +# --------------------------------------------------------------------------- +# Infrastructure & configuration +# --------------------------------------------------------------------------- + + +class AzureInfrastructureError(AzureError, ProviderTransientError): + """Raised for general Azure infrastructure errors.""" + + +class AzureConfigurationError(AzureError, ProviderConfigError): + """Raised when Azure configuration is invalid.""" + + +# --------------------------------------------------------------------------- +# Provisioning +# --------------------------------------------------------------------------- + + +class LaunchError(AzureError, ProviderTransientError): + """Instance / VMSS launch failed.""" + + def __init__( + self, + message: str, + template_id: str, + launch_params: Optional[dict[str, Any]] = None, + details: Optional[dict[str, Any]] = None, + error_code: Optional[str] = None, + ): + super().__init__( + message, + details={ + "template_id": template_id, + "launch_params": launch_params or {}, + **(details or {}), + }, + error_code=error_code, + ) + self.template_id = template_id + self.launch_params = launch_params or {} + + +class VMSSCreationError(LaunchError): + """VMSS creation specifically failed.""" + + def __init__( + self, + message: str, + template_id: str, + vmss_name: Optional[str] = None, + launch_params: Optional[dict[str, Any]] = None, + details: Optional[dict[str, Any]] = None, + error_code: Optional[str] = None, + ): + super().__init__( + message, + template_id=template_id, + launch_params=launch_params, + details={"vmss_name": vmss_name, **(details or {})}, + error_code=error_code, + ) + self.vmss_name = vmss_name + + +class TerminationError(AzureError, ProviderTransientError): + """Instance termination failed.""" + + def __init__( + self, + message: str, + resource_ids: list[str], + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={"resource_ids": resource_ids, **(details or {})}, + ) + self.resource_ids = resource_ids + + +class ResourceCleanupError(AzureError, ProviderTransientError): + """Failed to clean up Azure resources.""" + + def __init__( + self, + message: str, + resource_id: str, + resource_type: str, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={ + "resource_id": resource_id, + "resource_type": resource_type, + **(details or {}), + }, + ) + self.resource_id = resource_id + self.resource_type = resource_type + + +class TaggingError(AzureError, ProviderTransientError): + """Failed to tag Azure resources.""" + + def __init__( + self, + message: str, + resource_id: str, + tags: dict[str, str], + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={"resource_id": resource_id, "tags": tags, **(details or {})}, + ) + self.resource_id = resource_id + self.tags = tags + + +# --------------------------------------------------------------------------- +# CycleCloud +# --------------------------------------------------------------------------- + + +class CycleCloudError(AzureError): + """Base class for CycleCloud-related errors.""" + + +class CycleCloudConnectionError(CycleCloudError, ProviderTransientError): + """Failed to connect to the CycleCloud REST API.""" + + def __init__( + self, + message: str, + url: Optional[str] = None, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={"url": url, **(details or {})}, + ) + self.url = url + + +class CycleCloudClusterNotFoundError(CycleCloudError, ProviderPermanentError): + """The specified CycleCloud cluster was not found.""" + + def __init__( + self, + message: str, + cluster_name: str, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={"cluster_name": cluster_name, **(details or {})}, + ) + self.cluster_name = cluster_name + + +class CycleCloudNodeError(CycleCloudError, ProviderTransientError): + """Failed to add or manage CycleCloud nodes.""" + + def __init__( + self, + message: str, + cluster_name: str, + node_array: Optional[str] = None, + details: Optional[dict[str, Any]] = None, + ): + super().__init__( + message, + details={ + "cluster_name": cluster_name, + "node_array": node_array, + **(details or {}), + }, + ) + self.cluster_name = cluster_name + self.node_array = node_array diff --git a/src/orb/providers/azure/infrastructure/__init__.py b/src/orb/providers/azure/infrastructure/__init__.py new file mode 100644 index 000000000..6c1ddd856 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/__init__.py @@ -0,0 +1 @@ +"""Azure infrastructure layer.""" diff --git a/src/orb/providers/azure/infrastructure/adapters/__init__.py b/src/orb/providers/azure/infrastructure/adapters/__init__.py new file mode 100644 index 000000000..bde8b6311 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/adapters/__init__.py @@ -0,0 +1,7 @@ +"""Azure infrastructure adapters.""" + +from orb.providers.azure.infrastructure.adapters.template_adapter import AzureTemplateAdapter + +__all__: list[str] = [ + "AzureTemplateAdapter", +] diff --git a/src/orb/providers/azure/infrastructure/adapters/azure_validation_adapter.py b/src/orb/providers/azure/infrastructure/adapters/azure_validation_adapter.py new file mode 100644 index 000000000..9a3dc7e93 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/adapters/azure_validation_adapter.py @@ -0,0 +1,114 @@ +"""Adapter bridging Azure template validation into the ORB validation port.""" + +from typing import Any + +from orb.domain.base.ports.logging_port import LoggingPort +from orb.domain.base.ports.provider_validation_port import BaseProviderValidationAdapter +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.capabilities import get_supported_api_capabilities, get_supported_apis +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.configuration.validator import validate_azure_template + + +@injectable +class AzureValidationAdapter(BaseProviderValidationAdapter): + """ + azure implementation of the ProviderValidationAdapter interface. + + Encapsulates the Azure-specific validation logic that + requires access to Azure configuration. + + Features: + - Validates Azure provider API endpoints. + - Retrieves supported Azure provider APIs. + - Identifies the provider type as Azure. + + This adapter is wired via ``validator_factory`` in ``azure/registration.py`` + so validation can use Azure-native capability metadata without + constructing the full provider strategy. + """ + + def __init__(self, config: AzureProviderConfig, logger: LoggingPort) -> None: + """ + Initialize Azure validation adapter. + + Args: + config: Azure provider configuration + logger: Logger for validation operations + """ + self._config = config + self._logger = logger + + def get_provider_type(self) -> str: + """Get the provider type this adapter supports.""" + return "azure" + + def validate_provider_api(self, api: str) -> bool: + """ + Validate if a provider API is supported by azure. + + Args: + api: The provider API identifier to validate + + Returns: + True if the API is supported by azure configuration + """ + supported_apis = set(get_supported_apis()) + is_valid = api in supported_apis + if not is_valid: + self._logger.debug("azure API validation failed: %s not in %s", api, supported_apis) + return is_valid + + def get_supported_provider_apis(self) -> list[str]: + """ + Get list of all supported azure provider APIs. + + Returns: + List of supported azure provider API identifiers + """ + return sorted(get_supported_apis()) + + @staticmethod + def get_api_capabilities(api: str) -> dict[str, Any]: + """Get capability metadata for a specific Azure provider API.""" + capabilities = get_supported_api_capabilities().get(api) + if capabilities is None: + raise ValueError(f"Unsupported Azure provider API: {api}") + return capabilities + + def validate_template_configuration(self, template_config: dict[str, Any]) -> dict[str, Any]: + """Validate a complete Azure template configuration.""" + base_result = super().validate_template_configuration(template_config) + azure_result = validate_azure_template(template_config) + + errors: list[str] = [] + warnings: list[str] = [] + validated_fields: list[str] = [] + + for item in base_result.get("errors", []): + if item not in errors: + errors.append(item) + for item in azure_result.get("errors", []): + if item not in errors: + errors.append(item) + + for item in base_result.get("warnings", []): + if item not in warnings: + warnings.append(item) + for item in azure_result.get("warnings", []): + if item not in warnings: + warnings.append(item) + + for item in base_result.get("validated_fields", []): + if item not in validated_fields: + validated_fields.append(item) + for item in azure_result.get("validated_fields", []): + if item not in validated_fields: + validated_fields.append(item) + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": warnings, + "validated_fields": validated_fields, + } diff --git a/src/orb/providers/azure/infrastructure/adapters/template_adapter.py b/src/orb/providers/azure/infrastructure/adapters/template_adapter.py new file mode 100644 index 000000000..94fac09a0 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/adapters/template_adapter.py @@ -0,0 +1,54 @@ +"""Azure Template Adapter. + +Converts between configuration-layer template dicts and the domain +``AzureTemplate`` aggregate, following the adapter/port pattern. +""" + +from typing import Any + +from orb.domain.base.ports import LoggingPort +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate + + +@injectable +class AzureTemplateAdapter: + """Adapter for Azure template operations.""" + + def __init__(self, logger: LoggingPort) -> None: + self._logger = logger + + def create_template(self, template_data: dict[str, Any]) -> AzureTemplate: + """Create an ``AzureTemplate`` from a raw config dict. + + The dict may use either snake_case or camelCase keys thanks + to the ``AliasChoices`` defined on the aggregate. + """ + self._logger.debug("Creating AzureTemplate from data: %s", template_data.get("template_id")) + return AzureTemplate.from_azure_format(template_data) + + @staticmethod + def template_to_dict(template: AzureTemplate) -> dict[str, Any]: + """Serialise an ``AzureTemplate`` to a plain dict.""" + return template.model_dump(mode="json", by_alias=False, exclude_none=True) + + def validate_template_data(self, template_data: dict[str, Any]) -> dict[str, Any]: + """Validate template data by attempting to construct the aggregate. + + Returns: + dict with keys: valid (bool), errors (list[str]), warnings (list[str]) + """ + errors: list[str] = [] + warnings: list[str] = [] + + try: + self.create_template(template_data) + except Exception as exc: + errors.append(str(exc)) + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": warnings, + "validated_fields": list(template_data.keys()), + } diff --git a/src/orb/providers/azure/infrastructure/adapters/template_example_generator_adapter.py b/src/orb/providers/azure/infrastructure/adapters/template_example_generator_adapter.py new file mode 100644 index 000000000..c58c63af6 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/adapters/template_example_generator_adapter.py @@ -0,0 +1,22 @@ +"""Azure adapter for provider-owned template example generation.""" + +from typing import Any, Optional + +from orb.providers.azure.infrastructure.azure_handler_factory import ( + generate_azure_example_templates, +) + + +class AzureTemplateExampleGeneratorAdapter: + """Expose Azure handler examples through the shared generator port.""" + + def generate_example_templates( + self, + provider_name: str, + provider_api: Optional[str] = None, + ) -> list[Any]: + """Return Azure examples, optionally restricted to one provider API.""" + examples = generate_azure_example_templates() + if provider_api is None: + return examples + return [example for example in examples if example.get("provider_api") == provider_api] diff --git a/src/orb/providers/azure/infrastructure/azure_client.py b/src/orb/providers/azure/infrastructure/azure_client.py new file mode 100644 index 000000000..c9f219aca --- /dev/null +++ b/src/orb/providers/azure/infrastructure/azure_client.py @@ -0,0 +1,700 @@ +"""Azure client wrapper with lazy service client initialization. + +This module provides an integrated wrapper for Azure SDK interactions with: +- Lazy initialization of Azure service clients +- Explicit lifecycle cleanup for owned Azure SDK resources +- Explicit runtime config assembly before client construction + +Azure SDK clients wrapped: +- ComputeManagementClient (VMs, VMSS, Disks, Images, Galleries) +- NetworkManagementClient (VNets, Subnets, NICs, NSGs, Public IPs, LBs) +- ResourceManagementClient (Resource groups, deployments) +- SubscriptionClient (Subscriptions, locations) + +Note: + Actual Azure SDK packages (``azure-identity``, ``azure-mgmt-compute``, etc.) + are **not** imported at module level so that the rest of the codebase can be + loaded and tested without them installed. They are imported lazily inside + property accessors the first time a client is requested. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from threading import RLock +from typing import TYPE_CHECKING, Any, Callable, Optional + +from orb.config import PerformanceConfig +from orb.domain.base.ports import LoggingPort +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.exceptions.azure_exceptions import ( + AuthenticationError, + AzureConfigurationError, +) +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncAzureCredentialProtocol, + create_default_azure_credential_async, + format_import_error, +) +from orb.providers.azure.infrastructure.services.arm_resource_id_parser import ( + ArmResourceIdParser, + ParsedArmResourceId, +) +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentity, + AzureNetworkIdentityResolver, +) + +if TYPE_CHECKING: + from azure.mgmt.compute.aio import ComputeManagementClient as AsyncComputeManagementClient + from azure.mgmt.network.aio import NetworkManagementClient as AsyncNetworkManagementClient + from azure.mgmt.resource.resources.aio import ( + ResourceManagementClient as AsyncResourceManagementClient, + ) + from azure.mgmt.resource.subscriptions.aio import SubscriptionClient as AsyncSubscriptionClient + + +@dataclass(frozen=True) +class AzureClientRuntimeConfig: + """Infrastructure-owned config resolved before AzureClient construction.""" + + azure_config: AzureProviderConfig + performance_config: PerformanceConfig = field( + default_factory=lambda: PerformanceConfig.model_validate({}) + ) + + +@injectable +class AzureClient: + """Wrapper for Azure API interactions. + + * Provider and shared runtime config are resolved before construction + and passed in as a typed infrastructure object. + * Azure SDK management clients are created **lazily** on first access + through ``@property`` accessors, keeping startup cost near zero. + """ + + def __init__( + self, + runtime_config: AzureClientRuntimeConfig, + logger: LoggingPort, + ) -> None: + """Initialise the Azure client wrapper. + + Args: + runtime_config: Fully resolved Azure client runtime settings. + logger: Logger for diagnostic and operational messages. + """ + self._runtime_config = runtime_config + self._logger = logger + self._azure_config = runtime_config.azure_config + + self.region_name: str = ( + self._azure_config.region if self._azure_config.region else "eastus2" + ) + self.subscription_id: Optional[str] = self._azure_config.subscription_id + self.resource_group: Optional[str] = self._azure_config.resource_group + + self._logger.debug("Azure client region determined: %s", self.region_name) + + max_retries = self._resolve_int_config_value( + field_name="max_retries", + raw_value=self._azure_config.max_retries, + default=3, + minimum=0, + ) + connect_timeout = self._resolve_int_config_value( + field_name="connect_timeout", + raw_value=self._azure_config.connect_timeout, + default=30, + minimum=1, + ) + read_timeout = self._resolve_int_config_value( + field_name="read_timeout", + raw_value=self._azure_config.read_timeout, + default=60, + minimum=1, + ) + + self._max_retries = max_retries + self._connect_timeout = connect_timeout + self._read_timeout = read_timeout + + self.perf_config = self._map_performance_config(runtime_config.performance_config) + + # Lazy Azure SDK client slots + self._async_credential: Optional[AsyncAzureCredentialProtocol] = None + self._async_compute_client: Optional[AsyncComputeManagementClient] = None + self._async_network_client: Optional[AsyncNetworkManagementClient] = None + self._async_resource_client: Optional[AsyncResourceManagementClient] = None + self._async_subscription_client: Optional[AsyncSubscriptionClient] = None + self._credentials_validated = False + self._closed = False + self._pending_async_close_task: asyncio.Task[None] | None = None + self._lazy_init_lock = RLock() + self._arm_resource_id_parser = ArmResourceIdParser() + self._network_identity_resolver = AzureNetworkIdentityResolver( + async_network_client_getter=self.get_async_network_client, + logger=logger, + arm_resource_id_parser=self._arm_resource_id_parser, + network_lookup_error_types=self._network_lookup_error_types, + ) + + self._logger.info( + "Azure client initialised: region=%s, subscription=%s, resource_group=%s, " + "retries=%d, connect_timeout=%ds, read_timeout=%ds", + self.region_name, + self.subscription_id or "(not set)", + self.resource_group or "(not set)", + max_retries, + connect_timeout, + read_timeout, + ) + + def get_provider_config(self) -> Optional[AzureProviderConfig]: + """Return the active Azure provider configuration, if available.""" + return self._azure_config + + # ------------------------------------------------------------------ + # Credential management + # ------------------------------------------------------------------ + + def _management_client_kwargs(self) -> dict[str, Any]: + """Common azure-core kwargs for management clients.""" + return { + "retry_total": self._max_retries, + "connection_timeout": self._connect_timeout, + "read_timeout": self._read_timeout, + } + + @staticmethod + def _resolve_int_config_value( + *, + field_name: str, + raw_value: Optional[int], + default: int, + minimum: int, + ) -> int: + """Resolve an integer config value with explicit validation.""" + if raw_value is None: + return default + + try: + value = int(raw_value) + except (TypeError, ValueError) as exc: + raise AzureConfigurationError( + f"Azure provider config '{field_name}' must be an integer, got {raw_value!r}" + ) from exc + + if value < minimum: + raise AzureConfigurationError( + f"Azure provider config '{field_name}' must be >= {minimum}, got {value}" + ) + + return value + + def _ensure_open(self) -> None: + """Raise if the Azure client has already been closed.""" + if self._closed: + raise RuntimeError("AzureClient has been closed") + + async def _async_management_client_credential(self) -> Any: + """Return the async credential object to hand to async Azure SDK client constructors.""" + return await self.get_async_credential() + + async def _async_management_client_init_kwargs( + self, + *, + requires_subscription_id: bool, + ) -> dict[str, Any]: + """Build constructor kwargs shared by async Azure management clients.""" + client_kwargs = { + "credential": await self._async_management_client_credential(), + **self._management_client_kwargs(), + } + if requires_subscription_id: + self._ensure_subscription_id() + client_kwargs["subscription_id"] = self.subscription_id + return client_kwargs + + async def _build_management_client_async( + self, + *, + loader: Callable[[], Any], + client_name: str, + missing_package_message: str, + requires_subscription_id: bool, + ) -> Any: + """Construct a lazily imported async Azure SDK management client.""" + self._ensure_open() + self._logger.debug("Initialising async %s on first use", client_name) + try: + client_class = loader() + except ImportError as exc: + raise AzureConfigurationError(missing_package_message) from exc + return client_class( + **await self._async_management_client_init_kwargs( + requires_subscription_id=requires_subscription_id + ) + ) + + async def get_async_credential(self) -> AsyncAzureCredentialProtocol: + """Return an async Azure credential, creating it on first use.""" + with self._lazy_init_lock: + self._ensure_open() + existing = self._async_credential + if existing is not None: + return existing + self._logger.debug("Creating async Azure credential on first use") + try: + created = create_default_azure_credential_async( + client_id=self._azure_config.client_id if self._azure_config else None, + logger=self._logger, + ) + except ImportError as exc: + raise AuthenticationError(format_import_error("azure-identity", exc)) from exc + with self._lazy_init_lock: + self._ensure_open() + if self._async_credential is None: + self._async_credential = created + return created + existing = self._async_credential + await self._close_async_resource("async_credential", created) + return existing + + def close(self) -> None: + """Close owned Azure SDK resources and prevent further use. + + The Azure client owns its lazily-created credentials and management + clients, so cleanup belongs here rather than in unrelated orchestration + layers. + """ + close_errors: list[Exception] = [] + + with self._lazy_init_lock: + if self._closed: + return + + def close_resource(close_fn: Any, *close_args: Any) -> None: + """Invoke a close function, logging and collecting any errors.""" + resource_name = str(close_args[0]) + resource = close_args[-1] + if resource is None: + return + try: + close_fn(*close_args) + except Exception as exc: # pragma: no cover - exercised via public close tests + close_errors.append(exc) + self._logger.warning( + "Failed closing Azure resource %s: %s", + resource_name, + exc, + ) + + async_resources_present = any( + resource is not None + for resource in ( + self._async_subscription_client, + self._async_resource_client, + self._async_network_client, + self._async_compute_client, + self._async_credential, + ) + ) + + self._credentials_validated = False + self._closed = True + + if async_resources_present: + try: + loop = asyncio.get_running_loop() + except RuntimeError: + asyncio.run(self.aclose()) + else: + close_task = loop.create_task(self.aclose()) + with self._lazy_init_lock: + self._pending_async_close_task = close_task + + def _log_async_close_completion(task: asyncio.Task[None]) -> None: + with self._lazy_init_lock: + if self._pending_async_close_task is task: + self._pending_async_close_task = None + try: + task.result() + except Exception as exc: + self._logger.warning( + "Failed closing async Azure resources after AzureClient.close(): %s", + exc, + ) + + close_task.add_done_callback(_log_async_close_completion) + + if close_errors: + raise close_errors[0] + + async def aclose(self) -> None: + """Close owned async Azure SDK resources.""" + close_errors: list[Exception] = [] + + with self._lazy_init_lock: + self._credentials_validated = False + self._closed = True + async_subscription_client = self._async_subscription_client + self._async_subscription_client = None + async_resource_client = self._async_resource_client + self._async_resource_client = None + async_network_client = self._async_network_client + self._async_network_client = None + async_compute_client = self._async_compute_client + self._async_compute_client = None + async_credential = self._async_credential + self._async_credential = None + + for resource_name, resource in ( + ("async_subscription_client", async_subscription_client), + ("async_resource_client", async_resource_client), + ("async_network_client", async_network_client), + ("async_compute_client", async_compute_client), + ("async_credential", async_credential), + ): + if resource is None: + continue + try: + await self._close_async_resource(resource_name, resource) + except Exception as exc: + close_errors.append(exc) + self._logger.warning( + "Failed closing Azure resource %s: %s", + resource_name, + exc, + ) + + if close_errors: + raise close_errors[0] + + def __enter__(self) -> AzureClient: + """Enter a managed Azure client scope.""" + self._ensure_open() + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + """Exit a managed Azure client scope and close owned resources.""" + self.close() + + # ------------------------------------------------------------------ + # Lazy management-client properties + # ------------------------------------------------------------------ + + async def get_async_compute_client(self) -> AsyncComputeManagementClient: + """Lazy initialisation of the async Azure Compute management client.""" + with self._lazy_init_lock: + self._ensure_open() + existing = self._async_compute_client + if existing is not None: + return existing + + created = await self._build_management_client_async( + loader=self._load_async_compute_management_client, + client_name="ComputeManagementClient", + missing_package_message="azure-mgmt-compute package is not installed", + requires_subscription_id=True, + ) + with self._lazy_init_lock: + self._ensure_open() + if self._async_compute_client is None: + self._async_compute_client = created + return created + surviving_client = self._async_compute_client + await self._close_async_resource("async_compute_client", created) + if surviving_client is None: + raise RuntimeError("ComputeManagementClient disappeared during async initialization") + return surviving_client + + async def get_async_network_client(self) -> AsyncNetworkManagementClient: + """Lazy initialisation of the async Azure Network management client.""" + with self._lazy_init_lock: + self._ensure_open() + existing = self._async_network_client + if existing is not None: + return existing + + created = await self._build_management_client_async( + loader=self._load_async_network_management_client, + client_name="NetworkManagementClient", + missing_package_message="azure-mgmt-network package is not installed", + requires_subscription_id=True, + ) + with self._lazy_init_lock: + self._ensure_open() + if self._async_network_client is None: + self._async_network_client = created + return created + surviving_client = self._async_network_client + await self._close_async_resource("async_network_client", created) + if surviving_client is None: + raise RuntimeError("NetworkManagementClient disappeared during async initialization") + return surviving_client + + async def get_async_resource_client(self) -> AsyncResourceManagementClient: + """Lazy initialisation of the async Azure Resource management client.""" + with self._lazy_init_lock: + self._ensure_open() + existing = self._async_resource_client + if existing is not None: + return existing + + created = await self._build_management_client_async( + loader=self._load_async_resource_management_client, + client_name="ResourceManagementClient", + missing_package_message="azure-mgmt-resource package is not installed", + requires_subscription_id=True, + ) + with self._lazy_init_lock: + self._ensure_open() + if self._async_resource_client is None: + self._async_resource_client = created + return created + surviving_client = self._async_resource_client + await self._close_async_resource("async_resource_client", created) + if surviving_client is None: + raise RuntimeError("ResourceManagementClient disappeared during async initialization") + return surviving_client + + async def get_async_subscription_client(self) -> AsyncSubscriptionClient: + """Lazy initialisation of the async Azure Subscription client.""" + with self._lazy_init_lock: + self._ensure_open() + existing = self._async_subscription_client + if existing is not None: + return existing + + created = await self._build_management_client_async( + loader=self._load_async_subscription_client_class, + client_name="SubscriptionClient", + missing_package_message="azure-mgmt-resource package is not installed", + requires_subscription_id=False, + ) + with self._lazy_init_lock: + self._ensure_open() + if self._async_subscription_client is None: + self._async_subscription_client = created + return created + surviving_client = self._async_subscription_client + await self._close_async_resource("async_subscription_client", created) + if surviving_client is None: + raise RuntimeError("SubscriptionClient disappeared during async initialization") + return surviving_client + + # ------------------------------------------------------------------ + # Validation helpers + # ------------------------------------------------------------------ + + async def validate_credentials_async(self) -> bool: + """Async variant of credential validation using the async Azure credential.""" + if self._credentials_validated: + return True + + try: + credential = await self.get_async_credential() + await credential.get_token("https://management.azure.com/.default") + self._credentials_validated = True + self._logger.info("Azure credentials validated successfully") + return True + except self._credential_validation_error_types() as exc: + self._logger.error("Azure credential validation failed: %s", exc) + return False + + async def validate_subscription_async(self) -> bool: + """Async variant of subscription validation using the async Azure SDK client.""" + if not self.subscription_id: + self._logger.error("No subscription_id configured") + return False + + try: + subscription_client = await self.get_async_subscription_client() + sub = await subscription_client.subscriptions.get(self.subscription_id) + self._logger.info( + "Azure subscription validated: %s (%s)", + sub.display_name, + sub.state, + ) + return True + except self._subscription_validation_error_types() as exc: + self._logger.error( + "Azure subscription validation failed for %s: %s", + self.subscription_id, + exc, + ) + return False + + # ------------------------------------------------------------------ + # Performance / caching configuration + # ------------------------------------------------------------------ + + @staticmethod + def _default_performance_config() -> dict[str, Any]: + """Return the default Azure client performance settings.""" + return { + "enable_batching": True, + "batch_sizes": { + "deallocate_vms": 25, + "create_tags": 20, + "describe_vms": 25, + }, + "enable_parallel": True, + "max_workers": 10, + "enable_caching": True, + "cache_ttl": 300, + } + + @classmethod + def _map_performance_config(cls, perf_config: PerformanceConfig) -> dict[str, Any]: + """Project the shared performance config onto Azure client settings.""" + performance_settings = cls._default_performance_config() + performance_settings.update( + { + "enable_batching": perf_config.enable_batching, + "enable_parallel": perf_config.enable_parallel, + "max_workers": perf_config.max_workers, + "enable_caching": perf_config.caching.request_status.enabled, + "cache_ttl": perf_config.caching.request_status.ttl_seconds, + } + ) + return performance_settings + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _ensure_subscription_id(self) -> None: + """Raise if ``subscription_id`` is not configured.""" + if not self.subscription_id: + raise AzureConfigurationError( + "subscription_id is required but not configured. " + "Set it in the Azure provider configuration." + ) + + @staticmethod + def _collect_error_types( + *base_error_types: type[BaseException], + optional_error_loaders: tuple[Callable[[], Any], ...], + ) -> tuple[type[BaseException], ...]: + """Build a de-duplicated exception tuple with optional lazy imports.""" + error_types: list[type[BaseException]] = list(base_error_types) + for loader in optional_error_loaders: + try: + error_type = loader() + except ImportError: + continue + error_types.append(error_type) + return tuple(dict.fromkeys(error_types)) + + @classmethod + def _credential_validation_error_types(cls) -> tuple[type[BaseException], ...]: + """Return errors that should result in a soft credential validation failure.""" + return cls._collect_error_types( + AuthenticationError, + optional_error_loaders=( + cls._load_azure_error_type, + cls._load_client_authentication_error_type, + ), + ) + + @classmethod + def _subscription_validation_error_types(cls) -> tuple[type[BaseException], ...]: + """Return errors that should result in a soft subscription validation failure.""" + return cls._collect_error_types( + AzureConfigurationError, + *cls._credential_validation_error_types(), + optional_error_loaders=(), + ) + + @classmethod + def _network_lookup_error_types(cls) -> tuple[type[BaseException], ...]: + """Return Azure/network lookup errors that should not abort status enrichment.""" + return cls._collect_error_types( + AzureConfigurationError, + optional_error_loaders=(cls._load_azure_error_type,), + ) + + def _get_network_identity_resolver(self) -> AzureNetworkIdentityResolver: + """Return the initialized network identity resolver.""" + return self._network_identity_resolver + + async def _close_async_resource(self, resource_name: str, resource: Any) -> None: + """Close one async Azure SDK resource owned by this wrapper.""" + try: + await resource.close() + except Exception: + raise + self._logger.debug("Closed Azure resource %s", resource_name) + + @staticmethod + def _load_async_compute_management_client() -> Any: + from azure.mgmt.compute.aio import ComputeManagementClient + + return ComputeManagementClient + + @staticmethod + def _load_async_network_management_client() -> Any: + from azure.mgmt.network.aio import NetworkManagementClient + + return NetworkManagementClient + + @staticmethod + def _load_async_resource_management_client() -> Any: + from azure.mgmt.resource.resources.aio import ResourceManagementClient + + return ResourceManagementClient + + @staticmethod + def _load_async_subscription_client_class() -> Any: + from azure.mgmt.resource.subscriptions.aio import SubscriptionClient + + return SubscriptionClient + + @staticmethod + def _load_azure_error_type() -> Any: + from azure.core.exceptions import AzureError + + return AzureError + + @staticmethod + def _load_client_authentication_error_type() -> Any: + from azure.core.exceptions import ClientAuthenticationError + + return ClientAuthenticationError + + @classmethod + def _parse_arm_resource_id( + cls, + arm_id: str, + ) -> Optional[ParsedArmResourceId]: + """Parse an ARM resource ID only when it matches the canonical shape.""" + return ArmResourceIdParser.parse(arm_id) + + @classmethod + def extract_resource_group_and_name_from_arm_id( + cls, + arm_id: str, + ) -> Optional[tuple[str, str]]: + """Extract ``(resource_group, resource_name)`` from an ARM resource ID.""" + return ArmResourceIdParser().extract_resource_group_and_name(arm_id) + + @classmethod + def subnet_id_to_vnet_id(cls, subnet_id: Optional[str]) -> Optional[str]: + """Return the parent VNet ARM ID from a subnet ARM ID.""" + return ArmResourceIdParser().subnet_id_to_vnet_id(subnet_id) + + async def resolve_network_identity_from_vm_async(self, vm: Any) -> AzureNetworkIdentity: + """Async variant of VM network identity resolution.""" + return await self._get_network_identity_resolver().resolve_from_vm_async(vm) + + async def resolve_network_identity_from_nic_refs_async( + self, + nic_refs: list[Any], + ) -> AzureNetworkIdentity: + """Async variant of NIC-ref network identity resolution.""" + return await self._get_network_identity_resolver().resolve_from_nic_refs_async(nic_refs) diff --git a/src/orb/providers/azure/infrastructure/azure_handler_factory.py b/src/orb/providers/azure/infrastructure/azure_handler_factory.py new file mode 100644 index 000000000..494a393af --- /dev/null +++ b/src/orb/providers/azure/infrastructure/azure_handler_factory.py @@ -0,0 +1,158 @@ +"""Azure Handler Factory. + +Creates and caches Azure handlers based on ``provider_api`` values +""" + +from threading import RLock +from typing import TYPE_CHECKING + +from orb.domain.base.ports import LoggingPort +from orb.domain.template.template_aggregate import Template +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.infrastructure.handlers.azure_handler import AzureHandler + +if TYPE_CHECKING: + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + +@injectable +class AzureHandlerFactory: + """Factory for creating Azure handlers keyed by ``AzureProviderApi``.""" + + def __init__( + self, + azure_client: AzureClient, + logger: LoggingPort, + azure_native_spec_service: "AzureNativeSpecService | None" = None, + azure_resource_manager: "AzureResourceManager | None" = None, + ) -> None: + """Initialize the factory with an Azure client and register handler classes.""" + self._azure_client = azure_client + self._logger = logger + self._azure_native_spec_service = azure_native_spec_service + self._azure_resource_manager = azure_resource_manager + self._lock = RLock() + self._handlers: dict[AzureProviderApi, AzureHandler] = {} + self._handler_classes: dict[AzureProviderApi, type[AzureHandler]] = {} + self._register_handler_classes() + + @property + def azure_client(self) -> AzureClient: + """Return the underlying Azure client instance.""" + return self._azure_client + + @staticmethod + def _normalize_handler_type(handler_type: AzureProviderApi | str) -> AzureProviderApi: + if isinstance(handler_type, AzureProviderApi): + return handler_type + return AzureProviderApi(handler_type) + + def create_handler(self, handler_type: AzureProviderApi | str) -> AzureHandler: + """Create (or return cached) handler for *handler_type*. + + Raises: + AzureValidationError: If *handler_type* is unknown. + """ + try: + handler_type_key = self._normalize_handler_type(handler_type) + except ValueError: + raise AzureValidationError(f"Invalid Azure handler type: {handler_type}") + with self._lock: + if handler_type_key in self._handlers: + return self._handlers[handler_type_key] + + if handler_type_key not in self._handler_classes: + raise AzureValidationError( + f"No handler class registered for type: {handler_type_key.value}" + ) + + if handler_type_key is AzureProviderApi.SINGLE_VM: + from orb.providers.azure.infrastructure.handlers.single_vm_handler import ( + SingleVMHandler, + ) + + handler = SingleVMHandler( + azure_client=self._azure_client, + logger=self._logger, + azure_native_spec_service=self._azure_native_spec_service, + ) + elif handler_type_key in (AzureProviderApi.VMSS, AzureProviderApi.VMSS_UNIFORM): + from orb.providers.azure.infrastructure.handlers.vmss_handler import VMSSHandler + + handler = VMSSHandler( + azure_client=self._azure_client, + logger=self._logger, + azure_native_spec_service=self._azure_native_spec_service, + azure_resource_manager=self._azure_resource_manager, + ) + else: + from orb.providers.azure.infrastructure.handlers.cyclecloud_handler import ( + CycleCloudHandler, + ) + + handler = CycleCloudHandler( + azure_client=self._azure_client, + logger=self._logger, + ) + self._handlers[handler_type_key] = handler + self._logger.debug("Created Azure handler for type: %s", handler_type_key.value) + return handler + + def create_handler_for_template(self, template: Template) -> AzureHandler: + """Create handler appropriate for *template.provider_api*.""" + handler_type = template.provider_api or AzureProviderApi.VMSS + return self.create_handler(handler_type) + + def _register_handler_classes(self) -> None: + self._handler_classes = self.get_handler_classes() + self._logger.debug( + "Registered Azure handler classes: %s", + [handler_type.value for handler_type in self._handler_classes], + ) + + @staticmethod + def get_handler_classes() -> dict[AzureProviderApi, type[AzureHandler]]: + """Return the handler classes that define Azure's supported APIs.""" + from orb.providers.azure.infrastructure.handlers.cyclecloud_handler import ( + CycleCloudHandler, + ) + from orb.providers.azure.infrastructure.handlers.single_vm_handler import ( + SingleVMHandler, + ) + from orb.providers.azure.infrastructure.handlers.vmss_handler import VMSSHandler + + return { + AzureProviderApi.VMSS: VMSSHandler, + AzureProviderApi.VMSS_UNIFORM: VMSSHandler, + AzureProviderApi.SINGLE_VM: SingleVMHandler, + AzureProviderApi.CYCLECLOUD: CycleCloudHandler, + } + + def registered_handler_types(self) -> tuple[AzureProviderApi, ...]: + """Return the registered handler enums in factory-owned canonical order.""" + return tuple(self._handler_classes.keys()) + + def get_all_handlers(self) -> dict[str, AzureHandler]: + """Materialize and return handlers keyed by serialized Azure provider API values.""" + return { + handler_type.value: self.create_handler(handler_type) + for handler_type in self.registered_handler_types() + } + + def generate_example_templates(self) -> list[dict]: + """Collect example templates from all registered handlers.""" + return generate_azure_example_templates() + + +def generate_azure_example_templates() -> list[dict[str, object]]: + """Generate examples from the same handler set used for live dispatch.""" + examples: list[dict[str, object]] = [] + for handler_class in AzureHandlerFactory.get_handler_classes().values(): + examples.extend(handler_class.get_example_templates()) + return examples diff --git a/src/orb/providers/azure/infrastructure/credential_factory.py b/src/orb/providers/azure/infrastructure/credential_factory.py new file mode 100644 index 000000000..b954f1ad4 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/credential_factory.py @@ -0,0 +1,218 @@ +"""Azure credential construction owned by Azure infrastructure.""" + +from __future__ import annotations + +from typing import Any, Optional, Protocol + +from orb.domain.base.ports import LoggingPort + + +class AzureCredentialProtocol(Protocol): + """Credential surface required by Azure provider code.""" + + def get_token(self, *scopes: str, **kwargs: Any) -> Any: + """Request an access token for the given scopes.""" + ... + + def close(self) -> None: + """Release any resources held by this credential.""" + ... + + +class AsyncAzureCredentialProtocol(Protocol): + """Async credential surface required by native async Azure SDK flows.""" + + async def get_token(self, *scopes: str, **kwargs: Any) -> Any: + """Asynchronously request an access token for the given scopes.""" + ... + + async def close(self) -> None: + """Release any resources held by this async credential.""" + ... + + +class AzureAccessTokenProviderProtocol(Protocol): + """Short-lived Azure token provider used by auth flows.""" + + def get_access_token(self, scope: str) -> str: + """Return a raw access-token string for the given scope.""" + ... + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return exception types that signal authentication failures.""" + ... + + +class AsyncAzureAccessTokenProviderProtocol(Protocol): + """Async short-lived Azure token provider used by native async auth flows.""" + + async def get_access_token(self, scope: str) -> str: + """Return a raw access-token string for the given scope.""" + ... + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return exception types that signal authentication failures.""" + ... + + +class AzureCredentialAccessTokenProvider(AzureAccessTokenProviderProtocol): + """Adapt an existing Azure credential to the token-provider protocol.""" + + def __init__(self, credential: AzureCredentialProtocol) -> None: + self._credential = credential + + def get_access_token(self, scope: str) -> str: + """Delegate token retrieval to the wrapped credential.""" + token = self._credential.get_token(scope) + return token.token + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return Azure SDK credential error types.""" + return get_default_azure_credential_error_types() + + +class AsyncAzureCredentialAccessTokenProvider(AsyncAzureAccessTokenProviderProtocol): + """Adapt an async Azure credential to the async token-provider protocol.""" + + def __init__(self, credential: AsyncAzureCredentialProtocol) -> None: + self._credential = credential + + async def get_access_token(self, scope: str) -> str: + """Delegate token retrieval to the wrapped async credential.""" + token = await self._credential.get_token(scope) + return token.token + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return Azure SDK credential error types.""" + return get_default_azure_credential_error_types() + + +def format_import_error(package_name: str, exc: ImportError) -> str: + """Return a precise dependency error message without masking nested imports.""" + message = str(exc).strip() + if not message: + return f"{package_name} package is not installed" + return f"{package_name} dependency error: {message}" + + +def get_default_azure_credential_error_types() -> tuple[type[Exception], ...]: + """Return the expected Azure SDK exception types for credential operations.""" + try: + from azure.core.exceptions import ClientAuthenticationError + from azure.identity import CredentialUnavailableError + except ImportError: + return () + + return CredentialUnavailableError, ClientAuthenticationError + + +def _default_credential_kwargs(*, client_id: Optional[str]) -> dict[str, Any]: + """Build the shared DefaultAzureCredential kwargs for sync and async factories.""" + credential_kwargs: dict[str, Any] = {} + if client_id: + credential_kwargs["managed_identity_client_id"] = client_id + return credential_kwargs + + +def create_default_azure_credential( + *, + client_id: Optional[str], + logger: LoggingPort, +) -> AzureCredentialProtocol: + """Create the canonical Azure DefaultAzureCredential for ORB Azure flows.""" + try: + from azure.identity import DefaultAzureCredential + except ImportError as exc: + logger.error(format_import_error("azure-identity", exc)) + raise + + credential_kwargs = _default_credential_kwargs(client_id=client_id) + + try: + credential = DefaultAzureCredential(**credential_kwargs) + except ImportError as exc: + logger.error(format_import_error("azure-identity", exc)) + raise + logger.info("Azure DefaultAzureCredential initialised") + return credential + + +def create_default_azure_credential_async( + *, + client_id: Optional[str], + logger: LoggingPort, +) -> AsyncAzureCredentialProtocol: + """Create the canonical async Azure DefaultAzureCredential for ORB Azure flows.""" + try: + from azure.identity.aio import DefaultAzureCredential + except ImportError as exc: + logger.error(format_import_error("azure-identity", exc)) + raise + + credential_kwargs = _default_credential_kwargs(client_id=client_id) + + try: + credential = DefaultAzureCredential(**credential_kwargs) + except ImportError as exc: + logger.error(format_import_error("azure-identity", exc)) + raise + logger.info("Async Azure DefaultAzureCredential initialised") + return credential + + +class DefaultAzureAccessTokenProvider(AzureAccessTokenProviderProtocol): + """Resolve Azure access tokens with short-lived credentials.""" + + def __init__( + self, + *, + client_id: Optional[str], + logger: LoggingPort, + ) -> None: + self._client_id = client_id + self._logger = logger + + def get_access_token(self, scope: str) -> str: + """Create a short-lived credential, fetch a token, then close it.""" + credential = create_default_azure_credential( + client_id=self._client_id, + logger=self._logger, + ) + try: + token = credential.get_token(scope) + return token.token + finally: + credential.close() + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return Azure SDK and ImportError types for credential failures.""" + return ImportError, *get_default_azure_credential_error_types() + + +class AsyncDefaultAzureAccessTokenProvider(AsyncAzureAccessTokenProviderProtocol): + """Resolve Azure access tokens with short-lived async credentials.""" + + def __init__( + self, + *, + client_id: Optional[str], + logger: LoggingPort, + ) -> None: + self._client_id = client_id + self._logger = logger + + async def get_access_token(self, scope: str) -> str: + """Create a short-lived async credential, fetch a token, then close it.""" + credential = create_default_azure_credential_async( + client_id=self._client_id, + logger=self._logger, + ) + try: + token = await credential.get_token(scope) + return token.token + finally: + await credential.close() + + def get_auth_error_types(self) -> tuple[type[Exception], ...]: + """Return Azure SDK and ImportError types for credential failures.""" + return ImportError, *get_default_azure_credential_error_types() diff --git a/src/orb/providers/azure/infrastructure/cyclecloud_session.py b/src/orb/providers/azure/infrastructure/cyclecloud_session.py new file mode 100644 index 000000000..70c2ea098 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/cyclecloud_session.py @@ -0,0 +1,164 @@ +"""CycleCloud infrastructure session context.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional + +import httpx + + +def _coerce_optional_bool(value: Any) -> Optional[bool]: + if value in (None, ""): + return None + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "y", "on"}: + return True + if normalized in {"false", "0", "no", "n", "off"}: + return False + return bool(value) + + +def _mapping_value(data: dict[str, Any], *keys: str) -> Any: + for key in keys: + value = data.get(key) + if value not in (None, ""): + return value + return None + + +@dataclass(frozen=True) +class CycleCloudCredentialData: + """CycleCloud credential material resolved from a credential file.""" + + url: Optional[str] = None + verify_ssl: Optional[bool] = None + auth_mode: Optional[str] = None + username: Optional[str] = field(default=None, repr=False) + password: Optional[str] = field(default=None, repr=False) + bearer_token: Optional[str] = field(default=None, repr=False) + aad_scope: Optional[str] = None + + @classmethod + def from_mapping(cls, data: dict[str, Any]) -> CycleCloudCredentialData: + """Construct credential data from a flat config mapping.""" + return cls( + url=_mapping_value(data, "cyclecloud_url", "url"), + verify_ssl=_coerce_optional_bool( + _mapping_value(data, "cyclecloud_verify_ssl", "verify_ssl") + ), + auth_mode=_mapping_value(data, "cyclecloud_auth_mode", "auth_mode"), + username=_mapping_value(data, "cyclecloud_username", "username"), + password=_mapping_value(data, "cyclecloud_password", "password"), + bearer_token=_mapping_value(data, "cyclecloud_bearer_token", "bearer_token"), + aad_scope=_mapping_value(data, "cyclecloud_aad_scope", "aad_scope"), + ) + + +@dataclass(frozen=True) +class CycleCloudSessionSettings: + """Resolved CycleCloud transport and auth settings before session creation.""" + + base_url: str + verify_ssl: bool + auth_mode: Optional[str] + credential_path: Optional[str] + + +@dataclass(frozen=True) +class CycleCloudRequestContext: + """Typed CycleCloud request/follow-up context carried through handler flows.""" + + cluster_name: Optional[str] = None + node_array: Optional[str] = None + node_ids: tuple[str, ...] = () + operation_id: Optional[str] = None + operation_location: Optional[str] = None + added_count: Optional[int] = None + cyclecloud_url: Optional[str] = None + cyclecloud_credential_path: Optional[str] = None + cyclecloud_verify_ssl: Optional[bool] = None + cyclecloud_auth_mode: Optional[str] = None + cyclecloud_aad_scope: Optional[str] = None + + @classmethod + def from_mapping(cls, data: Optional[dict[str, Any]]) -> CycleCloudRequestContext: + """Construct a request context from an optional metadata mapping.""" + if not data: + return cls() + + raw_node_ids = data.get("node_ids") or () + if isinstance(raw_node_ids, (list, tuple)): + node_ids = tuple(str(node_id) for node_id in raw_node_ids if node_id not in (None, "")) + else: + node_ids = () + + raw_added_count = data.get("added_count") + added_count = int(raw_added_count) if raw_added_count not in (None, "") else None + + return cls( + cluster_name=data.get("cluster_name"), + node_array=data.get("node_array"), + node_ids=node_ids, + operation_id=data.get("operation_id"), + operation_location=data.get("operation_location"), + added_count=added_count, + cyclecloud_url=data.get("cyclecloud_url"), + cyclecloud_credential_path=data.get("cyclecloud_credential_path"), + cyclecloud_verify_ssl=_coerce_optional_bool(data.get("cyclecloud_verify_ssl")), + cyclecloud_auth_mode=data.get("cyclecloud_auth_mode"), + cyclecloud_aad_scope=data.get("cyclecloud_aad_scope"), + ) + + def to_metadata(self) -> dict[str, Any]: + """Serialize non-empty fields to a metadata dict for transport.""" + metadata: dict[str, Any] = {} + if self.cluster_name not in (None, ""): + metadata["cluster_name"] = self.cluster_name + if self.node_array not in (None, ""): + metadata["node_array"] = self.node_array + if self.node_ids: + metadata["node_ids"] = list(self.node_ids) + if self.operation_id not in (None, ""): + metadata["operation_id"] = self.operation_id + if self.operation_location not in (None, ""): + metadata["operation_location"] = self.operation_location + if self.added_count is not None: + metadata["added_count"] = self.added_count + if self.cyclecloud_url not in (None, ""): + metadata["cyclecloud_url"] = self.cyclecloud_url + if self.cyclecloud_credential_path not in (None, ""): + metadata["cyclecloud_credential_path"] = self.cyclecloud_credential_path + if self.cyclecloud_verify_ssl not in (None, ""): + metadata["cyclecloud_verify_ssl"] = self.cyclecloud_verify_ssl + if self.cyclecloud_auth_mode not in (None, ""): + metadata["cyclecloud_auth_mode"] = self.cyclecloud_auth_mode + if self.cyclecloud_aad_scope not in (None, ""): + metadata["cyclecloud_aad_scope"] = self.cyclecloud_aad_scope + return metadata + + +@dataclass(frozen=True) +class AsyncCycleCloudSessionContext: + """Resolved async CycleCloud HTTP session plus ORB-specific connection metadata.""" + + client: httpx.AsyncClient = field(repr=False) + base_url: str + auth_mode: Optional[str] + credential_path: Optional[str] + verify_ssl: bool + + def __repr__(self) -> str: + """Return a safe repr that avoids leaking client internals or auth material.""" + return ( + "AsyncCycleCloudSessionContext(" + f"base_url={self.base_url!r}, " + f"auth_mode={self.auth_mode!r}, " + f"credential_path={self.credential_path!r}, " + f"verify_ssl={self.verify_ssl!r})" + ) diff --git a/src/orb/providers/azure/infrastructure/cyclecloud_session_builder.py b/src/orb/providers/azure/infrastructure/cyclecloud_session_builder.py new file mode 100644 index 000000000..676cd041f --- /dev/null +++ b/src/orb/providers/azure/infrastructure/cyclecloud_session_builder.py @@ -0,0 +1,235 @@ +"""CycleCloud session settings resolution for Azure infrastructure.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Optional +from urllib.parse import urlparse + +import httpx + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.exceptions.azure_exceptions import CycleCloudConnectionError +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncAzureAccessTokenProviderProtocol, +) +from orb.providers.azure.infrastructure.cyclecloud_session import ( + CycleCloudCredentialData, + CycleCloudRequestContext, + CycleCloudSessionSettings, +) + + +class CycleCloudSessionBuilder: + """Resolve CycleCloud credential and transport settings before session creation.""" + + def __init__( + self, + *, + cc_url: Optional[str], + verify_ssl: Optional[bool], + template: Optional[AzureTemplate], + request_context: Optional[CycleCloudRequestContext], + provider_cfg: Optional[AzureProviderConfig], + async_token_provider: Optional[AsyncAzureAccessTokenProviderProtocol] = None, + ): + self._cc_url = cc_url + self._verify_ssl = verify_ssl + self._template = template + self._request_context = request_context or CycleCloudRequestContext() + self._provider_cfg = provider_cfg + self._async_token_provider = async_token_provider + + @classmethod + def _load_credential_file(cls, credential_path: str) -> CycleCloudCredentialData: + path = Path(credential_path).expanduser() + try: + with path.open(encoding="utf-8") as handle: + data = json.load(handle) + except FileNotFoundError as exc: + raise CycleCloudConnectionError( + f"CycleCloud credential file not found: {path}", + url=None, + ) from exc + except json.JSONDecodeError as exc: + raise CycleCloudConnectionError( + f"CycleCloud credential file is not valid JSON: {path}", + url=None, + ) from exc + except OSError as exc: + raise CycleCloudConnectionError( + f"Failed to read CycleCloud credential file {path}: {exc}", + url=None, + ) from exc + + if not isinstance(data, dict): + raise CycleCloudConnectionError( + f"CycleCloud credential file must contain a JSON object: {path}", + url=None, + ) + + return CycleCloudCredentialData.from_mapping(data) + + def _provider_cyclecloud(self): + if self._provider_cfg is None: + return None + return self._provider_cfg.cyclecloud + + @staticmethod + def _resolve_cascaded_value( + *sources: object, + default: object = None, + ) -> object: + """Return the first configured value from the resolution cascade.""" + for value in sources: + if value not in (None, ""): + return value + return default + + async def _get_azure_bearer_token_async(self, scopes: list[str]) -> Optional[str]: + if self._async_token_provider is None: + return None + for scope in scopes: + if not scope: + continue + try: + token = await self._async_token_provider.get_access_token(scope) + if token: + return token + except self._async_token_provider.get_auth_error_types(): + continue + return None + + def _resolve_credential_path(self) -> Optional[str]: + provider_cyclecloud = self._provider_cyclecloud() + credential_path = self._resolve_cascaded_value( + None if self._template is None else self._template.cyclecloud_credential_path, + self._request_context.cyclecloud_credential_path, + None if provider_cyclecloud is None else provider_cyclecloud.credential_path, + ) + if credential_path in (None, ""): + return None + return str(credential_path) + + def _resolve_transport_settings( + self, + credential_data: CycleCloudCredentialData, + ) -> tuple[str, bool]: + provider_cyclecloud = self._provider_cyclecloud() + resolved_url = self._resolve_cascaded_value( + self._cc_url, + None if self._template is None else self._template.cyclecloud_url, + self._request_context.cyclecloud_url, + None if provider_cyclecloud is None else provider_cyclecloud.url, + credential_data.url, + ) + + verify_resolved = self._resolve_cascaded_value( + self._verify_ssl, + None if self._template is None else self._template.cyclecloud_verify_ssl, + self._request_context.cyclecloud_verify_ssl, + None if provider_cyclecloud is None else provider_cyclecloud.verify_ssl, + credential_data.verify_ssl, + default=True, + ) + + if not resolved_url: + raise CycleCloudConnectionError( + "cyclecloud_url is required in the template, request context, or provider configuration.", + url=None, + ) + + return str(resolved_url).rstrip("/"), bool(verify_resolved) + + def _resolve_auth_mode( + self, + credential_data: CycleCloudCredentialData, + ) -> Optional[str]: + provider_cyclecloud = self._provider_cyclecloud() + auth_mode = self._resolve_cascaded_value( + None if self._template is None else self._template.cyclecloud_auth_mode, + self._request_context.cyclecloud_auth_mode, + None if provider_cyclecloud is None else provider_cyclecloud.auth_mode, + credential_data.auth_mode, + ) + return str(auth_mode).strip().lower() if auth_mode else None + + async def _resolve_bearer_token_async( + self, + *, + base_url: str, + credential_data: CycleCloudCredentialData, + ) -> Optional[str]: + if credential_data.bearer_token: + return str(credential_data.bearer_token) + + provider_cyclecloud = self._provider_cyclecloud() + aad_scope = self._resolve_cascaded_value( + None if self._template is None else self._template.cyclecloud_aad_scope, + self._request_context.cyclecloud_aad_scope, + None if provider_cyclecloud is None else provider_cyclecloud.aad_scope, + credential_data.aad_scope, + ) + + parsed = urlparse(base_url) + host_scope = ( + f"{parsed.scheme}://{parsed.netloc}/.default" if parsed.scheme and parsed.netloc else "" + ) + scopes = [str(aad_scope)] if aad_scope else [] + scopes.extend([host_scope, "https://management.azure.com/.default"]) + return await self._get_azure_bearer_token_async(scopes) + + def build_settings(self) -> CycleCloudSessionSettings: + """Resolve credential, transport, and auth settings into a session config.""" + credential_path = self._resolve_credential_path() + credential_data = ( + self._load_credential_file(credential_path) + if credential_path + else CycleCloudCredentialData() + ) + base_url, verify_ssl = self._resolve_transport_settings(credential_data) + auth_mode = self._resolve_auth_mode(credential_data) + return CycleCloudSessionSettings( + base_url=base_url, + verify_ssl=verify_ssl, + auth_mode=auth_mode, + credential_path=credential_path, + ) + + async def resolve_async_auth( + self, + *, + settings: CycleCloudSessionSettings, + ) -> tuple[dict[str, str], httpx.BasicAuth | None, str]: + """Resolve auth settings for an ``httpx.AsyncClient`` transport.""" + if settings.auth_mode == "ssh": + raise CycleCloudConnectionError( + "cyclecloud_auth_mode=ssh is not supported. Configure CycleCloud API credentials instead.", + url=settings.base_url, + ) + + credential_data = ( + self._load_credential_file(settings.credential_path) + if settings.credential_path + else CycleCloudCredentialData() + ) + if credential_data.username and credential_data.password and settings.auth_mode != "bearer": + return {}, httpx.BasicAuth(credential_data.username, credential_data.password), "basic" + + bearer_token = await self._resolve_bearer_token_async( + base_url=settings.base_url, + credential_data=credential_data, + ) + if bearer_token: + return {"Authorization": f"Bearer {bearer_token}"}, None, "bearer" + if settings.auth_mode == "bearer": + raise CycleCloudConnectionError( + "cyclecloud_auth_mode=bearer requested but no bearer token could be resolved.", + url=settings.base_url, + ) + raise CycleCloudConnectionError( + "No CycleCloud auth method resolved. Provide username/password or a bearer token/Azure credential.", + url=settings.base_url, + ) diff --git a/src/orb/providers/azure/infrastructure/dry_run_adapter.py b/src/orb/providers/azure/infrastructure/dry_run_adapter.py new file mode 100644 index 000000000..e1421baf4 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/dry_run_adapter.py @@ -0,0 +1,23 @@ +"""Azure dry-run context manager. + +Activates global ``dry_run_context`` from the infrastructure layer so that +Azure handler code can call ``is_dry_run_active()`` to skip real ARM calls. +""" + +from collections.abc import Generator +from contextlib import contextmanager + +from orb.infrastructure.mocking.dry_run_context import dry_run_context + + +@contextmanager +def azure_dry_run_context() -> Generator[None, None, None]: + """Context manager that activates dry-run mode for Azure operations. + + Usage:: + + with azure_dry_run_context(): + result = await strategy.execute_operation(op) + """ + with dry_run_context(active=True): + yield diff --git a/src/orb/providers/azure/infrastructure/error_codes.py b/src/orb/providers/azure/infrastructure/error_codes.py new file mode 100644 index 000000000..2ff7a74f1 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/error_codes.py @@ -0,0 +1,34 @@ +"""Azure provider error-code payloads.""" + +from __future__ import annotations + +from typing import TypedDict + + +class ProviderErrorEntry(TypedDict, total=False): + """Normalized Azure infrastructure error payload.""" + + error_code: str + error_message: str + instance_id: str + resource_id: str + node_array: str + cc_state: str + lifecycle: str | None + launch_template_id: str | None + launch_template_version: str | None + subnet_id: str | None + instance_type: str | None + instance_requirements: object + status_code: str | None + status_level: str | None + + +def collect_provider_error_codes(errors: list[ProviderErrorEntry]) -> list[str]: + """Return unique canonical error codes from normalized Azure errors.""" + error_codes: list[str] = [] + for error in errors: + error_code = error.get("error_code") + if error_code and error_code not in error_codes: + error_codes.append(str(error_code)) + return error_codes diff --git a/src/orb/providers/azure/infrastructure/error_utils.py b/src/orb/providers/azure/infrastructure/error_utils.py new file mode 100644 index 000000000..9a75432e4 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/error_utils.py @@ -0,0 +1,153 @@ +"""Helpers for normalizing Azure SDK / ARM error payloads.""" + +from __future__ import annotations + +from typing import Any + + +def _json_response_body(response: Any) -> dict[str, Any] | None: + """Best-effort response JSON extraction for heterogeneous SDK responses. + + Uses getattr because Azure SDK response objects do not share a typed protocol + for optional JSON bodies across sync and async transports. + """ + if response is None: + return None + json_method = getattr(response, "json", None) + if not callable(json_method): + return None + try: + body = json_method() + except Exception: + return None + return body if isinstance(body, dict) else None + + +def _normalise_error_details(details: Any) -> list[dict[str, Any]]: + """Normalise Azure nested error details from dicts or SDK error objects. + + Uses getattr because nested Azure error items may be plain dictionaries or + SDK model instances depending on where the exception was raised. + """ + if not isinstance(details, list): + return [] + + normalised: list[dict[str, Any]] = [] + for item in details: + if isinstance(item, dict): + normalised.append(item) + continue + code = getattr(item, "code", None) + message = getattr(item, "message", None) + if code is None and message is None: + continue + normalised.append( + { + "code": str(code) if code not in (None, "") else None, + "message": str(message) if message not in (None, "") else None, + } + ) + return normalised + + +def extract_azure_error_details(exc: Exception) -> dict[str, Any]: + """Extract Azure SDK error details from common exception shapes. + + Uses getattr throughout because this normalises errors from the full + azure-core hierarchy (HttpResponseError, ServiceRequestError, ODataV4Error, + etc.) plus non-Azure exceptions. No single base class exposes all of + error / response / error_code / status_code / message uniformly. + """ + error = getattr(exc, "error", None) + response = getattr(exc, "response", None) + exception_details = getattr(exc, "details", None) + response_body = _json_response_body(response) + response_error = response_body.get("error") if isinstance(response_body, dict) else None + + raw_error_code = ( + getattr(exc, "error_code", None) + or getattr(error, "code", None) + or getattr(exc, "code", None) + or ( + exception_details.get("raw_error_code") if isinstance(exception_details, dict) else None + ) + or (response_error.get("code") if isinstance(response_error, dict) else None) + ) + status_code = getattr(exc, "status_code", None) + if status_code is None and response is not None: + status_code = getattr(response, "status_code", None) + if status_code is None and isinstance(exception_details, dict): + status_code = exception_details.get("status_code") + + message = ( + getattr(error, "message", None) + or getattr(exc, "message", None) + or (exception_details.get("error_message") if isinstance(exception_details, dict) else None) + or (response_error.get("message") if isinstance(response_error, dict) else None) + or str(exc) + ) + details = _normalise_error_details(getattr(error, "details", None)) + if not details and isinstance(exception_details, dict): + details = _normalise_error_details(exception_details.get("details")) + if not details and isinstance(response_error, dict): + details = _normalise_error_details(response_error.get("details")) + + return { + "raw_error_code": str(raw_error_code) if raw_error_code not in (None, "") else None, + "status_code": status_code, + "message": str(message), + "details": details, + } + + +def canonical_azure_error_code(exc: Exception) -> str: + """Return a stable Azure provisioning error code.""" + details = extract_azure_error_details(exc) + raw_error_code = details["raw_error_code"] + status_code = details["status_code"] + message = details["message"].lower() + + if raw_error_code: + return raw_error_code + if status_code == 429: + return "TooManyRequests" + if status_code == 404: + return "ResourceNotFound" + if status_code == 403: + return "AuthorizationFailed" + if "quota" in message or "exceed" in message: + return "QuotaExceeded" + if "allocationfailed" in message or "insufficient" in message: + return "AllocationFailed" + if "validation" in message or "invalid" in message: + return "InvalidRequest" + return type(exc).__name__ + + +_QUOTA_CODES = frozenset({"QuotaExceeded", "OperationNotAllowed", "ResourceQuotaExceeded"}) +_VALIDATION_CODES = frozenset({"InvalidRequest", "InvalidParameter", "BadRequest"}) + + +def classify_azure_error(exc: Exception) -> tuple[str, str]: + """Classify an Azure exception as ``("quota"|"validation"|"other", error_code)``. + + Canonical Azure error codes take precedence. Message-based string matching + is only consulted when the canonical mapping fell back to ``type(exc).__name__`` + — without that guard, tag or resource names containing "quota" or "exceed" + can misclassify unrelated errors. + """ + error_code = canonical_azure_error_code(exc) + + if error_code in _QUOTA_CODES: + return ("quota", error_code) + if error_code in _VALIDATION_CODES: + return ("validation", error_code) + + if error_code == type(exc).__name__: + message = extract_azure_error_details(exc)["message"].lower() + if "quota" in message or "exceeded" in message: + return ("quota", error_code) + if "validation" in message or "invalid" in message: + return ("validation", error_code) + + return ("other", error_code) diff --git a/src/orb/providers/azure/infrastructure/handlers/__init__.py b/src/orb/providers/azure/infrastructure/handlers/__init__.py new file mode 100644 index 000000000..6a1689fcf --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/__init__.py @@ -0,0 +1 @@ +"""Azure infrastructure handlers package.""" diff --git a/src/orb/providers/azure/infrastructure/handlers/_network_identity.py b/src/orb/providers/azure/infrastructure/handlers/_network_identity.py new file mode 100644 index 000000000..765d67827 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/_network_identity.py @@ -0,0 +1,55 @@ +"""Shared handler helpers for optional Azure network identity enrichment.""" + +from __future__ import annotations + +from typing import Awaitable, Callable, Optional + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentity, +) + + +def empty_network_identity() -> AzureNetworkIdentity: + """Return the empty network-identity shape used when enrichment fails.""" + return { + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + + +def network_identity_soft_failure_types() -> tuple[type[BaseException], ...]: + """Return enrichment failures that should not hide visible Azure machines.""" + azure_error_type: Optional[type[BaseException]] + error_types: list[type[BaseException]] = [AttributeError, TypeError] + try: + from azure.core.exceptions import AzureError + + azure_error_type = AzureError + except ImportError: + azure_error_type = None + if azure_error_type is not None: + error_types.append(azure_error_type) + return tuple(dict.fromkeys(error_types)) + + +async def resolve_network_identity_or_empty_async( + *, + logger: LoggingPort, + target_label: str, + resolver: Callable[[], Awaitable[AzureNetworkIdentity]], +) -> AzureNetworkIdentity: + """Resolve optional network identity without hiding otherwise visible machines.""" + try: + return await resolver() + except network_identity_soft_failure_types() as exc: + logger.warning( + "Failed to resolve network identity for %s: %s", + target_label, + exc, + ) + return empty_network_identity() diff --git a/src/orb/providers/azure/infrastructure/handlers/azure_handler.py b/src/orb/providers/azure/infrastructure/handlers/azure_handler.py new file mode 100644 index 000000000..ce6a4f169 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/azure_handler.py @@ -0,0 +1,259 @@ +"""Azure handler base class. + +All Azure infrastructure handlers extend this ABC and expose the native async +create, status, and release operations used by Azure services. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Optional, TypeAlias, TypedDict + +from orb.domain.base.ports import LoggingPort +from orb.domain.request.aggregate import Request +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.exceptions.azure_exceptions import ( + AzureValidationError, + TerminationError, +) +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.infrastructure.cyclecloud_session import CycleCloudRequestContext + + +class _AzureAcquireHostsRequiredResult(TypedDict): + """Fields returned by every Azure create handler.""" + + success: bool + resource_ids: list[str] + instances: list[dict[str, Any]] + + +class AzureAcquireHostsResult(_AzureAcquireHostsRequiredResult, total=False): + """Normalized result returned by Azure create handlers.""" + + error_message: str | None + provider_data: dict[str, Any] + + +class AzureStatusProviderData(TypedDict, total=False): + """Provider-owned metadata surfaced on Azure status results. + + Provider-specific fields (``availability_zone``, ``location``, etc.) live + here per the ``metadata vs provider_data`` architecture rule. The + HostFactory scheduler reads ``cloud_host_id`` from this dict to emit the + Symphony wire ``cloudHostId`` field. + """ + + resource_id: str + cloud_host_id: str | None + vm_name: str + vmss_name: str + vm_id: str + vmss_instance_id: str + node_id: str + node_name: str + cluster_name: str + node_array: str + cc_state: str + hostname: str + resource_group: str + location: str + availability_zone: str | None + nic_id: str | None + nic_name: str | None + vnet_id: str | None + fleet_errors: list[dict[str, Any]] + + +class AzureHandlerStatusResult(TypedDict, total=False): + """Normalized status record returned by Azure handlers.""" + + instance_id: str + name: str + resource_id: str + status: str + private_ip: str | None + public_ip: str | None + launch_time: str | None + instance_type: str | None + subnet_id: str | None + vpc_id: str | None + price_type: str | None + tags: dict[str, str] + provider_type: str + error: str + provider_data: AzureStatusProviderData + + +RAISE_ON_STATUS_ERROR_METADATA_KEY = "raise_on_status_error" + + +def azure_raise_on_status_error(request: Request) -> bool: + """Read the explicit Azure status error policy from request metadata.""" + metadata = request.metadata or {} + value = metadata.get(RAISE_ON_STATUS_ERROR_METADATA_KEY, False) + if isinstance(value, bool): + return value + raise AzureValidationError( + (f"Azure status metadata '{RAISE_ON_STATUS_ERROR_METADATA_KEY}' must be a boolean"), + error_code="InvalidParameter", + ) + + +@dataclass(frozen=True) +class AzureReleaseContext: + """Provider-owned runtime context required for Azure termination flows.""" + + resource_group: str | None = None + resource_id: str | None = None + cyclecloud_request_context: CycleCloudRequestContext = field( + default_factory=CycleCloudRequestContext + ) + + +class AzureSubmittedDeletion(TypedDict, total=False): + """One submitted or attempted Azure deletion target.""" + + requested_id: str + vm_name: str + error: str + + +class AzurePendingResourceCleanupMetadata(TypedDict, total=False): + """Durable VMSS cleanup metadata persisted for follow-up reconciliation.""" + + resource_group: str + vmss_name: str + machine_ids: list[str] + delete_vmss_when_empty: bool + member_delete_submitted: bool + delete_submitted: bool + delete_retry_pending: bool + last_delete_error: str + + +class AzureVmssReleaseProviderData(TypedDict, total=False): + """Provider data returned when a VMSS termination request is submitted.""" + + resource_group: str + vmss_name: str + operation_status: str + submitted_deletions: list[AzureSubmittedDeletion] + failed_deletions: list[AzureSubmittedDeletion] + resolved_instance_ids: list[str] + pending_resource_cleanup: AzurePendingResourceCleanupMetadata + + +class AzureSingleVmReleaseProviderData(TypedDict, total=False): + """Provider data returned when SingleVM termination requests are submitted.""" + + resource_group: str + operation_status: str + submitted_deletions: list[AzureSubmittedDeletion] + + +class AzureCycleCloudReleaseProviderData(TypedDict, total=False): + """Provider data returned when CycleCloud termination requests are submitted.""" + + cluster_name: str + terminate_operation_location: str + operation_status: str + + +AzureReleaseProviderData: TypeAlias = ( + AzureVmssReleaseProviderData + | AzureSingleVmReleaseProviderData + | AzureCycleCloudReleaseProviderData +) + + +class AzureReleaseHostsResult(TypedDict, total=False): + """Normalized termination submission result returned by Azure handlers.""" + + provider_data: AzureReleaseProviderData + + +@injectable +class AzureHandler(ABC): + """Abstract base handler for Azure provisioning operations. + + Concrete implementations (``VMSSHandler``, ``SingleVMHandler``) + implement async methods for their specific Azure API surface. + """ + + def __init__( + self, + azure_client: AzureClient, + logger: LoggingPort, + ) -> None: + self.azure_client = azure_client + self._logger = logger + + def _resolve_release_resource_group( + self, + *, + machine_ids: list[str], + context: Optional[AzureReleaseContext], + ) -> str: + """Resolve the resource group for Azure termination submissions.""" + release_context = context or AzureReleaseContext() + resource_group = release_context.resource_group or self.azure_client.resource_group + if resource_group: + return resource_group + raise TerminationError( + "resource_group is required for release_hosts", + resource_ids=machine_ids, + ) + + @staticmethod + def _resolve_subnet_id(template: AzureTemplate) -> str | None: + """Return the subnet ARM ID from network_config or subnet_ids.""" + if template.network_config and template.network_config.subnet_id: + return template.network_config.subnet_id + + subnet_ids = [ + subnet_id + for subnet_id in template.subnet_ids + if subnet_id and subnet_id != "default-subnet" + ] + if len(subnet_ids) > 1: + raise AzureValidationError( + "Azure templates support a single subnet for VM network interfaces; " + "set network_config.subnet_id or provide exactly one subnet_id", + details={ + "template_id": template.template_id, + "subnet_ids": subnet_ids, + }, + error_code="InvalidParameter", + ) + if subnet_ids: + return subnet_ids[0] + return None + + # ------------------------------------------------------------------ + # Core operations + # ------------------------------------------------------------------ + + @abstractmethod + async def acquire_hosts_async( + self, request: Request, template: AzureTemplate + ) -> AzureAcquireHostsResult: + """Provision resources without blocking the event loop.""" + + @abstractmethod + async def check_hosts_status_async(self, request: Request) -> list[AzureHandlerStatusResult]: + """Return instance details without blocking the event loop.""" + + @abstractmethod + async def release_hosts_async( + self, + machine_ids: list[str], + resource_id: str, + context: Optional[AzureReleaseContext] = None, + ) -> Optional[AzureReleaseHostsResult]: + """Delete / deallocate cloud resources without blocking the event loop.""" + + @classmethod + def get_example_templates(cls) -> list[dict[str, Any]]: + """Return example template dicts for documentation / wizard use.""" + return [] diff --git a/src/orb/providers/azure/infrastructure/handlers/azure_status.py b/src/orb/providers/azure/infrastructure/handlers/azure_status.py new file mode 100644 index 000000000..91c710f9e --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/azure_status.py @@ -0,0 +1,48 @@ +"""Shared Azure VM status mapping. + +Provides a single source of truth for mapping Azure VM power-state and +provisioning-state codes to ORB domain status strings. Used by all +handlers that work with Azure SDK VM objects (SingleVM, VMSS) and by +the machine conversion service. +""" + +from __future__ import annotations + +from orb.providers.azure.infrastructure.sdk_shapes import AzureStatusWithCodeProtocol + +# Azure VM state map. PowerState/* entries are common to all +# VM-based handlers; ProvisioningState/* entries provide a fallback when +# a VM has not yet reached a power state (e.g. still being created). +AZURE_VM_STATE_MAP: dict[str, str] = { + # Power states + "PowerState/starting": "pending", + "PowerState/running": "running", + "PowerState/stopping": "stopping", + "PowerState/stopped": "stopped", + "PowerState/deallocating": "shutting-down", + "PowerState/deallocated": "stopped", + # Provisioning states (fallback when no PowerState is present) + "ProvisioningState/creating": "pending", + "ProvisioningState/succeeded": "running", + "ProvisioningState/failed": "failed", + "ProvisioningState/deleting": "shutting-down", +} + + +def resolve_power_state(statuses: list[AzureStatusWithCodeProtocol]) -> str: + """Extract the ORB domain status from a list of Azure InstanceViewStatus objects. + + Tries PowerState/* first (most accurate for running VMs), then falls + back to ProvisioningState/* for VMs that are still being created or + torn down. + """ + for status in statuses: + code = str(status.code or "") + if code.startswith("PowerState/"): + return AZURE_VM_STATE_MAP.get(code, "unknown") + # Fallback to provisioning state + for status in statuses: + code = str(status.code or "") + if code.startswith("ProvisioningState/"): + return AZURE_VM_STATE_MAP.get(code, "unknown") + return "unknown" diff --git a/src/orb/providers/azure/infrastructure/handlers/cyclecloud_handler.py b/src/orb/providers/azure/infrastructure/handlers/cyclecloud_handler.py new file mode 100644 index 000000000..d7989aaa0 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/cyclecloud_handler.py @@ -0,0 +1,1048 @@ +"""CycleCloud Handler - provisions nodes via the Azure CycleCloud REST API. + +This handler is used when ``provider_api == "CycleCloud"`` in the template. +It manages nodes within an existing CycleCloud cluster, adding nodes to a +specified node array and tracking them through the CycleCloud lifecycle. + +CycleCloud REST API reference: + https://learn.microsoft.com/en-us/azure/cyclecloud/api + +Key CycleCloud concepts: +- Clusters contain one or more node arrays (partitions) +- Each node array has a VM type, image, and autoscale configuration +- Nodes are added/removed by adjusting the target count or via direct API calls +- Node states: Off → Acquiring → Preparing → Ready → (Deallocating → Deallocated) +""" + +from __future__ import annotations + +import json +from collections.abc import Awaitable, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any, Optional + +import httpx + +from orb.domain.base.ports import LoggingPort +from orb.domain.request.aggregate import Request +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import ( + AuthenticationError, + CycleCloudClusterNotFoundError, + CycleCloudConnectionError, + CycleCloudNodeError, + TerminationError, +) +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncAzureCredentialAccessTokenProvider, +) +from orb.providers.azure.infrastructure.cyclecloud_session import ( + AsyncCycleCloudSessionContext, + CycleCloudRequestContext, +) +from orb.providers.azure.infrastructure.cyclecloud_session_builder import ( + CycleCloudSessionBuilder, +) +from orb.providers.azure.infrastructure.error_codes import ( + ProviderErrorEntry, + collect_provider_error_codes, +) +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + AzureAcquireHostsResult, + AzureHandler, + AzureHandlerStatusResult, + AzureReleaseContext, + AzureReleaseHostsResult, + AzureStatusProviderData, +) + +# CycleCloud node state → domain status mapping +_CC_STATE_MAP: dict[str, str] = { + "Off": "stopped", + "Acquiring": "pending", + "Preparing": "pending", + "Starting": "pending", + "Started": "running", + "Software Configuration": "pending", + "Ready": "running", + "Deallocating": "shutting-down", + "Deallocated": "stopped", + "Terminated": "terminated", + "Failed": "failed", +} + + +def resolve_cc_state(state: str) -> str: + """Map a CycleCloud node state to a domain status string.""" + return _CC_STATE_MAP.get(state, "unknown") + + +@dataclass(frozen=True) +class CycleCloudNode: + """Normalized CycleCloud node payload used inside the handler.""" + + name: str + node_id: str + node_array: str + state: str + private_ip: Optional[str] + public_ip: Optional[str] + machine_type: str + create_time: Optional[str] + subnet_id: Optional[str] + hostname: Optional[str] + error_message: Optional[str] + error_code: Optional[str] + + +@dataclass(frozen=True) +class _CycleCloudAcquireRequest: + """Normalized request payload inputs for CycleCloud node creation.""" + + cluster_name: str + node_array: str + vm_size: str + count: int + cyclecloud_request_id: str + node_params: dict[str, Any] + + +@dataclass(frozen=True) +class _CycleCloudStatusRequest: + """Durable request-scoped context needed for CycleCloud status checks.""" + + request_context: CycleCloudRequestContext + cluster_name: str + node_array: Optional[str] + node_ids: list[str] + cyclecloud_request_id: str + + +def _optional_text(value: Any) -> Optional[str]: + """Return a normalized optional string from an external JSON value.""" + if value in (None, ""): + return None + return str(value) + + +def _first_text(node: dict[str, Any], *keys: str, default: str = "") -> str: + """Return the first populated string value from an external node payload.""" + for key in keys: + value = _optional_text(node.get(key)) + if value is not None: + return value + return default + + +def _first_optional_text(node: dict[str, Any], *keys: str) -> Optional[str]: + """Return the first populated optional string value from an external node payload.""" + for key in keys: + value = _optional_text(node.get(key)) + if value is not None: + return value + return None + + +def _collect_cyclecloud_status_results( + *, + logger: LoggingPort, + cluster_name: str, + nodes: list[dict[str, Any]], + node_array: Optional[str], + node_ids: list[str], +) -> list[AzureHandlerStatusResult]: + """Filter and normalize CycleCloud node payloads for status responses.""" + results: list[AzureHandlerStatusResult] = [] + for node in nodes: + parsed_node = _parse_cyclecloud_node(node) + node_name = parsed_node.name + node_id = parsed_node.node_id or node_name + + if node_array and parsed_node.node_array != node_array: + continue + if node_ids and node_name not in node_ids and node_id not in node_ids: + continue + + cc_state = parsed_node.state + status = resolve_cc_state(cc_state) + if status == "unknown": + logger.warning("Unmapped CycleCloud node state: %s", cc_state) + + fleet_errors = _extract_cyclecloud_node_errors( + node, + cluster_name=cluster_name, + node_array=parsed_node.node_array, + ) + + results.append( + _build_cyclecloud_status_result( + cluster_name=cluster_name, + parsed_node=parsed_node, + cc_state=cc_state, + status=status, + fleet_errors=fleet_errors, + ) + ) + return results + + +def _build_cyclecloud_status_result( + *, + cluster_name: str, + parsed_node: CycleCloudNode, + cc_state: str, + status: str, + fleet_errors: list[ProviderErrorEntry], +) -> AzureHandlerStatusResult: + """Build a typed Azure status result for one CycleCloud node.""" + node_name = parsed_node.name or parsed_node.hostname or parsed_node.node_id + node_id = parsed_node.node_id or node_name + provider_data: AzureStatusProviderData = { + "resource_id": cluster_name, + "cloud_host_id": node_id, + "cluster_name": cluster_name, + "node_array": parsed_node.node_array, + "node_id": node_id, + "cc_state": cc_state, + "fleet_errors": [dict(error) for error in fleet_errors], + } + if parsed_node.name: + provider_data["node_name"] = parsed_node.name + if parsed_node.hostname: + provider_data["hostname"] = parsed_node.hostname + return { + "instance_id": node_name, + "name": node_name, + "resource_id": cluster_name, + "status": status, + "private_ip": parsed_node.private_ip, + "public_ip": parsed_node.public_ip, + "launch_time": parsed_node.create_time, + "instance_type": parsed_node.machine_type, + "subnet_id": parsed_node.subnet_id, + "vpc_id": None, + "tags": {}, + "price_type": None, + "provider_type": "azure", + "provider_data": provider_data, + } + + +def _parse_cyclecloud_node(node: dict[str, Any]) -> CycleCloudNode: + """Normalize a CycleCloud node-list payload into a typed object. + + Learn docs underdocument this response shape; in practice CycleCloud has + returned both lower-camel and PascalCase node fields. Normalize those + variants once at the provider boundary. + """ + return CycleCloudNode( + name=_first_text(node, "name", "Name"), + node_id=_first_text(node, "nodeId", "NodeId", "id"), + node_array=_first_text(node, "nodeArray", "NodeArray", "Template"), + state=_first_text(node, "state", "State", "status", "Status", default="Unknown"), + private_ip=_first_optional_text(node, "privateIp", "PrivateIp", "ipAddress", "IpAddress"), + public_ip=_first_optional_text(node, "publicIp", "PublicIp"), + machine_type=_first_text(node, "machineType", "MachineType", default="unknown"), + create_time=_first_optional_text(node, "createTime", "CreateTime"), + subnet_id=_first_optional_text(node, "subnetId", "SubnetId"), + hostname=_first_optional_text(node, "hostname", "Hostname"), + error_message=_first_optional_text( + node, + "message", + "Message", + "error", + "Error", + "statusMessage", + "StatusMessage", + "failureMessage", + "FailureMessage", + ), + error_code=_first_optional_text(node, "errorCode", "ErrorCode"), + ) + + +def _extract_cyclecloud_node_errors( + node: dict[str, Any], + *, + cluster_name: str, + node_array: str, +) -> list[ProviderErrorEntry]: + """Extract structured node errors from CycleCloud node payloads.""" + parsed_node = _parse_cyclecloud_node(node) + if parsed_node.state == "Unknown" and ( + node.get("status") not in (None, "") or node.get("id") not in (None, "") + ): + parsed_node = _parse_cyclecloud_node_result(node) + + message = parsed_node.error_message + error_code = parsed_node.error_code or ("NodeFailed" if parsed_node.state == "Failed" else None) + + if not error_code and not message: + return [] + if parsed_node.state != "Failed" and not message: + return [] + + node_error: ProviderErrorEntry = { + "error_code": str(error_code or "CycleCloudNodeError"), + "error_message": str(message or f"CycleCloud node entered state {parsed_node.state}"), + "resource_id": cluster_name, + "node_array": node_array, + "cc_state": parsed_node.state, + } + if parsed_node.name or parsed_node.node_id: + node_error["instance_id"] = parsed_node.name or parsed_node.node_id + return [node_error] + + +def _parse_cyclecloud_node_result(node: dict[str, Any]) -> CycleCloudNode: + """Normalize a CycleCloud node-management result payload into a typed object.""" + return CycleCloudNode( + name=str(node.get("name") or ""), + node_id=str(node.get("id") or ""), + node_array="", + state=str(node.get("status") or "Unknown"), + private_ip=None, + public_ip=None, + machine_type="unknown", + create_time=None, + subnet_id=None, + hostname=None, + error_message=_optional_text(node.get("message") or node.get("error")), + error_code=_optional_text(node.get("errorCode")), + ) + + +def _response_json_or_error( + *, + response: httpx.Response, + url: str, + decode_error_types: tuple[type[BaseException], ...], +) -> Any: + """Parse a CycleCloud JSON response body or raise a normalized connection error.""" + body: Any = {} + if response.content: + try: + body = response.json() + except decode_error_types as exc: + raise CycleCloudConnectionError( + f"CycleCloud API returned invalid JSON from {url}: {exc}", + url=url, + ) from exc + return body + + +def _response_metadata(response: httpx.Response, *, url: str) -> dict[str, Any]: + """Normalize response metadata returned from sync or async HTTP transports.""" + return { + "headers": dict(response.headers), + "status_code": response.status_code, + "url": url, + } + + +def _http_error_details(body_json: dict[str, Any] | None) -> dict[str, Any]: + """Extract the stable CycleCloud error details shape from a parsed JSON body.""" + error_code = None + if isinstance(body_json, dict): + error_code = ( + body_json.get("code") + or (body_json.get("error") or {}).get("code") + or body_json.get("type") + ) + return { + "body_json": body_json, + "error_code": error_code, + } + + +@injectable +class CycleCloudHandler(AzureHandler): + """Handler that manages nodes in an Azure CycleCloud cluster. + + ``provider_api = "CycleCloud"`` + + This handler communicates with the CycleCloud REST API to: + - Add nodes to a cluster's node array (``acquire_hosts``) + - Query node status (``check_hosts_status``) + - Remove/terminate nodes (``release_hosts``) + + The CycleCloud API credentials and URL are resolved from the template + or from the Azure provider configuration. + """ + + # ------------------------------------------------------------------ + # Internal: CycleCloud REST API helpers + # ------------------------------------------------------------------ + + def _get_provider_cyclecloud_config(self) -> Optional[AzureProviderConfig]: + loaded_cfg = self.azure_client.get_provider_config() + if isinstance(loaded_cfg, AzureProviderConfig): + return loaded_cfg + return None + + def _get_cc_request_timeout(self) -> tuple[int, int]: + provider_cfg = self._get_provider_cyclecloud_config() + if provider_cfg is None: + return 30, 60 + return provider_cfg.connect_timeout, provider_cfg.read_timeout + + async def _build_async_cc_session( + self, + *, + cc_url: Optional[str], + verify_ssl: Optional[bool], + template: Optional[AzureTemplate] = None, + request_context: Optional[CycleCloudRequestContext] = None, + ) -> AsyncCycleCloudSessionContext: + provider_cfg = self._get_provider_cyclecloud_config() + token_provider = None + try: + credential = await self.azure_client.get_async_credential() + except AuthenticationError: + credential = None + if credential is not None: + token_provider = AsyncAzureCredentialAccessTokenProvider(credential) + session_builder = CycleCloudSessionBuilder( + cc_url=cc_url, + verify_ssl=verify_ssl, + template=template, + request_context=request_context, + provider_cfg=provider_cfg, + async_token_provider=token_provider, + ) + settings = session_builder.build_settings() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + auth_headers, auth, resolved_auth_mode = await session_builder.resolve_async_auth( + settings=settings, + ) + headers.update(auth_headers) + connect_timeout, read_timeout = self._get_cc_request_timeout() + client = httpx.AsyncClient( + verify=settings.verify_ssl, + headers=headers, + auth=auth, + timeout=httpx.Timeout( + connect=connect_timeout, + read=read_timeout, + write=read_timeout, + pool=connect_timeout, + ), + ) + return AsyncCycleCloudSessionContext( + client=client, + base_url=settings.base_url, + auth_mode=resolved_auth_mode, + credential_path=settings.credential_path, + verify_ssl=settings.verify_ssl, + ) + + @asynccontextmanager + async def _async_cc_session_scope( + self, + *, + cc_url: Optional[str], + verify_ssl: Optional[bool], + template: Optional[AzureTemplate] = None, + request_context: Optional[CycleCloudRequestContext] = None, + ): + session_context = await self._build_async_cc_session( + cc_url=cc_url, + verify_ssl=verify_ssl, + template=template, + request_context=request_context, + ) + try: + yield session_context + finally: + await session_context.client.aclose() + + async def _cc_request_async( + self, + client: httpx.AsyncClient, + method: str, + url: str, + *, + include_metadata: bool = False, + **kwargs: Any, + ) -> Any: + try: + response = await client.request(method, url, **kwargs) + response.raise_for_status() + body = _response_json_or_error( + response=response, + url=url, + decode_error_types=(json.JSONDecodeError,), + ) + if not include_metadata: + return body + return _response_metadata(response, url=str(response.request.url)) | {"body": body} + except httpx.ConnectError as exc: + raise CycleCloudConnectionError( + f"Cannot connect to CycleCloud at {url}: {exc}", + url=url, + ) from exc + except httpx.HTTPStatusError as exc: + response = exc.response + body = response.text + body_json: dict[str, Any] | None = None + try: + parsed = response.json() + if isinstance(parsed, dict): + body_json = parsed + except json.JSONDecodeError: + body_json = None + + raise CycleCloudConnectionError( + f"CycleCloud API error (HTTP {response.status_code}): {body}", + url=url, + details={ + "status_code": response.status_code, + "body": body, + **_http_error_details(body_json), + }, + ) from exc + except httpx.HTTPError as exc: + raise CycleCloudConnectionError( + f"CycleCloud API request failed: {exc}", + url=url, + ) from exc + + async def _resolve_release_node_targets_via_fetch_async( + self, + *, + fetch_nodes: Callable[[], Awaitable[dict[str, Any]]], + machine_ids: list[str], + ) -> dict[str, list[str]]: + """Async variant of release-target resolution using fetched cluster nodes.""" + try: + nodes_response = await fetch_nodes() + except CycleCloudConnectionError: + return {"names": machine_ids} + return self._resolve_release_node_targets_from_nodes( + nodes=nodes_response.get("nodes", []), + machine_ids=machine_ids, + ) + + async def _resolve_release_node_targets_async( + self, + *, + client: httpx.AsyncClient, + base_url: str, + cluster_name: str, + machine_ids: list[str], + ) -> dict[str, list[str]]: + """Resolve stored machine IDs to the strongest CycleCloud identifier set available.""" + return await self._resolve_release_node_targets_via_fetch_async( + fetch_nodes=lambda: self._cc_request_async( + client, + "GET", + f"{base_url}/clusters/{cluster_name}/nodes", + ), + machine_ids=machine_ids, + ) + + def _resolve_release_node_targets_from_nodes( + self, + *, + nodes: list[dict[str, Any]], + machine_ids: list[str], + ) -> dict[str, list[str]]: + """Resolve stored machine IDs using fetched CycleCloud node payloads.""" + resolved_ids: list[str] = [] + resolved_names: list[str] = [] + seen_ids: set[str] = set() + seen_names: set[str] = set() + + for machine_id in machine_ids: + matched = False + for node in nodes: + parsed_node = _parse_cyclecloud_node(node) + if machine_id in {parsed_node.name, parsed_node.node_id}: + if parsed_node.node_id and parsed_node.node_id not in seen_ids: + resolved_ids.append(parsed_node.node_id) + seen_ids.add(parsed_node.node_id) + if parsed_node.name and parsed_node.name not in seen_names: + resolved_names.append(parsed_node.name) + seen_names.add(parsed_node.name) + matched = True + break + if not matched and machine_id and machine_id not in seen_names: + resolved_names.append(machine_id) + seen_names.add(machine_id) + + if resolved_ids or resolved_names != machine_ids: + self._logger.info( + "Resolved CycleCloud release ids %s -> node_ids=%s node_names=%s", + machine_ids, + resolved_ids, + resolved_names, + ) + + if resolved_ids: + return {"ids": resolved_ids} + return {"names": resolved_names} + + @staticmethod + def _build_release_result( + *, + cluster_name: str, + terminate_response: dict[str, Any], + ) -> AzureReleaseHostsResult: + """Build the normalized CycleCloud release submission result.""" + return { + "provider_data": { + "cluster_name": cluster_name, + "terminate_operation_location": ( + terminate_response.get("headers", {}).get("Location") + ), + "operation_status": "submitted", + } + } + + # ------------------------------------------------------------------ + # acquire_hosts + # ------------------------------------------------------------------ + + def _prepare_acquire_request( + self, + *, + request: Request, + template: AzureTemplate, + ) -> _CycleCloudAcquireRequest: + """Build the CycleCloud create-nodes request payload.""" + cluster_name = template.cluster_name + node_array = template.node_array + vm_size = template.vm_size + count = request.requested_count + + if not cluster_name: + raise CycleCloudNodeError( + "cluster_name is required for CycleCloud provisioning.", + cluster_name=cluster_name or "", + node_array=node_array, + ) + + definition: dict[str, Any] = {} + if not template.uses_vm_size_mix: + definition["machineType"] = vm_size + + cyclecloud_request_id = str(request.request_id) + node_params: dict[str, Any] = { + "requestId": cyclecloud_request_id, + "sets": [ + { + "count": count, + "nodearray": node_array, + "definition": definition, + } + ], + } + + if template.node_attributes: + node_params["sets"][0]["definition"].update(template.node_attributes) + + subnet_id = self._resolve_subnet_id(template) + if subnet_id: + node_params["sets"][0]["definition"]["SubnetId"] = subnet_id + + return _CycleCloudAcquireRequest( + cluster_name=cluster_name, + node_array=node_array, + vm_size=vm_size, + count=count, + cyclecloud_request_id=cyclecloud_request_id, + node_params=node_params, + ) + + async def _validate_cluster_exists_async( + self, + *, + cluster_name: str, + fetch_cluster_status: Callable[[], Awaitable[dict[str, Any]]], + ) -> None: + """Async variant of CycleCloud cluster existence validation.""" + try: + self._log_cluster_state( + cluster_name=cluster_name, + cluster_status=await fetch_cluster_status(), + ) + except CycleCloudConnectionError as exc: + if exc.details and exc.details.get("status_code") == 404: + raise CycleCloudClusterNotFoundError( + f"CycleCloud cluster '{cluster_name}' not found.", + cluster_name=cluster_name, + ) from exc + raise + + def _log_cluster_state(self, *, cluster_name: str, cluster_status: dict[str, Any]) -> None: + """Log the current CycleCloud cluster state after a successful status fetch.""" + cluster_state = cluster_status.get("state", "Unknown") + self._logger.debug("CycleCloud cluster '%s' state: %s", cluster_name, cluster_state) + + def _build_acquire_result( + self, + *, + request_data: _CycleCloudAcquireRequest, + template: AzureTemplate, + base_url: str, + credential_path: Optional[str], + verify_ssl: bool, + auth_mode: Optional[str], + create_response: dict[str, Any], + ) -> AzureAcquireHostsResult: + """Build the normalized CycleCloud acquire response payload.""" + result = create_response.get("body") or {} + operation_id = result.get("operationId", "") + operation_location = create_response.get("headers", {}).get("Location") + created_sets = result.get("sets", []) + fleet_errors: list[ProviderErrorEntry] = [] + added_count = 0 + + for node_set in created_sets: + added = node_set.get("added", 0) + added_count += int(added or 0) + for node in node_set.get("nodes", []): + node_errors = _extract_cyclecloud_node_errors( + node, + cluster_name=request_data.cluster_name, + node_array=request_data.node_array, + ) + for error in node_errors: + if error not in fleet_errors: + fleet_errors.append(error) + + self._logger.info( + "CycleCloud node request accepted for cluster '%s': operation_id=%s, added=%d", + request_data.cluster_name, + operation_id, + added_count, + ) + + return { + "success": True, + "resource_ids": [request_data.cyclecloud_request_id], + "instances": [], + "error_message": None, + "provider_data": { + "cluster_name": request_data.cluster_name, + "node_array": request_data.node_array, + "operation_id": operation_id, + "operation_location": operation_location, + "added_count": added_count, + "submitted_count": request_data.count, + "operation_status": "submitted", + "fulfillment_final": True, + "resource_group": template.resource_group.value, + "location": template.location.value, + "error_codes": collect_provider_error_codes(fleet_errors), + "fleet_errors": fleet_errors, + "cyclecloud_url": base_url, + "cyclecloud_credential_path": credential_path, + "cyclecloud_verify_ssl": verify_ssl, + "cyclecloud_auth_mode": auth_mode, + "cyclecloud_aad_scope": template.cyclecloud_aad_scope, + }, + } + + def _prepare_status_request(self, request: Request) -> Optional[_CycleCloudStatusRequest]: + """Validate and normalize the durable status-check request context.""" + resource_ids = request.resource_ids + if not resource_ids: + self._logger.warning("check_hosts_status called with no resource_ids") + return None + + request_context = CycleCloudRequestContext.from_mapping(request.metadata or {}) + cluster_name = request_context.cluster_name + cyclecloud_request_id = resource_ids[0] + + if not cluster_name: + message = "cluster_name is required for CycleCloud status check" + self._logger.error(message) + raise CycleCloudConnectionError( + message, + url=request_context.cyclecloud_url, + details={"request_id": request.request_id}, + ) + + if not cyclecloud_request_id: + message = f"CycleCloud request identity is required for status check in cluster '{cluster_name}'" + self._logger.error(message) + raise CycleCloudConnectionError( + message, + url=request_context.cyclecloud_url, + details={"resource_ids": resource_ids}, + ) + + return _CycleCloudStatusRequest( + request_context=request_context, + cluster_name=cluster_name, + node_array=request_context.node_array, + node_ids=list(request_context.node_ids), + cyclecloud_request_id=cyclecloud_request_id, + ) + + def _log_status_check_failure( + self, + *, + cluster_name: str, + cyclecloud_request_id: str, + exc: CycleCloudConnectionError, + ) -> None: + """Normalize logging for status-check failures across sync and async paths.""" + if "Cannot connect to CycleCloud" in str(exc): + self._logger.error( + "Failed to build CycleCloud session for status check (cluster '%s'): %s", + cluster_name, + exc, + ) + return + self._logger.error( + "Failed to get node status for cluster '%s' and request_id '%s': %s", + cluster_name, + cyclecloud_request_id, + exc, + ) + + def _build_status_results( + self, + *, + status_request: _CycleCloudStatusRequest, + nodes_response: dict[str, Any], + ) -> list[AzureHandlerStatusResult]: + """Build normalized status results from a CycleCloud nodes response.""" + results = _collect_cyclecloud_status_results( + logger=self._logger, + cluster_name=status_request.cluster_name, + nodes=nodes_response.get("nodes", []), + node_array=status_request.node_array, + node_ids=status_request.node_ids, + ) + + self._logger.debug( + "CycleCloud status check for cluster '%s': %d node(s) found", + status_request.cluster_name, + len(results), + ) + return results + + async def _submit_release_request_async( + self, + *, + cluster_name: str, + machine_ids: list[str], + resolve_node_targets: Callable[[], Awaitable[dict[str, list[str]]]], + submit_terminate: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]], + ) -> AzureReleaseHostsResult: + """Async variant of CycleCloud node termination submission.""" + self._logger.info( + "Terminating %d node(s) from CycleCloud cluster '%s': %s", + len(machine_ids), + cluster_name, + machine_ids, + ) + return await self._submit_release_request_result_async( + cluster_name=cluster_name, + machine_ids=machine_ids, + node_targets=await resolve_node_targets(), + submit_terminate=submit_terminate, + ) + + async def _submit_release_request_result_async( + self, + *, + cluster_name: str, + machine_ids: list[str], + node_targets: dict[str, list[str]], + submit_terminate: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]], + ) -> AzureReleaseHostsResult: + """Submit a prepared CycleCloud termination payload through the async client.""" + try: + terminate_payload: dict[str, Any] = dict(node_targets) + terminate_response = await submit_terminate(terminate_payload) + self._logger.debug("Terminate request sent for CycleCloud nodes: %s", terminate_payload) + self._logger.info( + "Successfully submitted termination for %d node(s) from cluster '%s'", + len(machine_ids), + cluster_name, + ) + return self._build_release_result( + cluster_name=cluster_name, + terminate_response=terminate_response, + ) + except CycleCloudConnectionError as exc: + raise TerminationError( + f"Failed to terminate nodes from CycleCloud cluster '{cluster_name}': {exc}", + resource_ids=machine_ids, + ) from exc + + async def acquire_hosts_async( + self, request: Request, template: AzureTemplate + ) -> AzureAcquireHostsResult: + """Async variant of ``acquire_hosts`` using ``httpx.AsyncClient``.""" + acquire_request = self._prepare_acquire_request(request=request, template=template) + + self._logger.info( + "Adding %d node(s) to CycleCloud cluster '%s' node array '%s' (vm_size=%s)", + acquire_request.count, + acquire_request.cluster_name, + acquire_request.node_array, + acquire_request.vm_size, + ) + + async with self._async_cc_session_scope( + cc_url=template.cyclecloud_url, + verify_ssl=template.cyclecloud_verify_ssl, + template=template, + ) as session_context: + client = session_context.client + base_url = session_context.base_url + + await self._validate_cluster_exists_async( + cluster_name=acquire_request.cluster_name, + fetch_cluster_status=lambda: self._cc_request_async( + client, + "GET", + f"{base_url}/clusters/{acquire_request.cluster_name}/status", + ), + ) + + try: + create_response = await self._cc_request_async( + client, + "POST", + f"{base_url}/clusters/{acquire_request.cluster_name}/nodes/create", + include_metadata=True, + json=acquire_request.node_params, + ) + except CycleCloudConnectionError as exc: + raise CycleCloudNodeError( + f"Failed to add nodes to cluster '{acquire_request.cluster_name}': {exc}", + cluster_name=acquire_request.cluster_name, + node_array=acquire_request.node_array, + ) from exc + + return self._build_acquire_result( + request_data=acquire_request, + template=template, + base_url=base_url, + credential_path=session_context.credential_path, + verify_ssl=session_context.verify_ssl, + auth_mode=session_context.auth_mode, + create_response=create_response, + ) + + async def check_hosts_status_async(self, request: Request) -> list[AzureHandlerStatusResult]: + """Async variant of ``check_hosts_status`` using ``httpx.AsyncClient``.""" + status_request = self._prepare_status_request(request) + if status_request is None: + return [] + + try: + async with self._async_cc_session_scope( + cc_url=status_request.request_context.cyclecloud_url, + verify_ssl=status_request.request_context.cyclecloud_verify_ssl, + request_context=status_request.request_context, + ) as session_context: + nodes_response = await self._cc_request_async( + session_context.client, + "GET", + f"{session_context.base_url}/clusters/{status_request.cluster_name}/nodes", + params={"request_id": status_request.cyclecloud_request_id}, + ) + except CycleCloudConnectionError as exc: + self._log_status_check_failure( + cluster_name=status_request.cluster_name, + cyclecloud_request_id=status_request.cyclecloud_request_id, + exc=exc, + ) + raise + + return self._build_status_results( + status_request=status_request, + nodes_response=nodes_response, + ) + + async def release_hosts_async( + self, + machine_ids: list[str], + resource_id: str, + context: Optional[AzureReleaseContext] = None, + ) -> Optional[AzureReleaseHostsResult]: + """Async variant of ``release_hosts`` using ``httpx.AsyncClient``.""" + release_context = context or AzureReleaseContext() + request_context = release_context.cyclecloud_request_context + cluster_name = str(request_context.cluster_name or resource_id) + + try: + async with self._async_cc_session_scope( + cc_url=request_context.cyclecloud_url, + verify_ssl=request_context.cyclecloud_verify_ssl, + request_context=request_context, + ) as session_context: + return await self._submit_release_request_async( + cluster_name=cluster_name, + machine_ids=machine_ids, + resolve_node_targets=lambda: self._resolve_release_node_targets_async( + client=session_context.client, + base_url=session_context.base_url, + cluster_name=cluster_name, + machine_ids=machine_ids, + ), + submit_terminate=lambda terminate_payload: self._cc_request_async( + session_context.client, + "POST", + f"{session_context.base_url}/clusters/{cluster_name}/nodes/terminate", + include_metadata=True, + json=terminate_payload, + ), + ) + except CycleCloudConnectionError as exc: + raise TerminationError( + f"Failed to build CycleCloud session for release_hosts: {exc}", + resource_ids=machine_ids, + ) from exc + + # ------------------------------------------------------------------ + # Example templates + # ------------------------------------------------------------------ + + @classmethod + def get_example_templates(cls) -> list[dict[str, Any]]: + """Return example CycleCloud template configurations.""" + return [ + { + "template_id": "azure-cyclecloud-hpc", + "name": "Azure CycleCloud HPC Cluster", + "description": "Add HPC nodes to an existing CycleCloud cluster", + "provider_type": "azure", + "provider_api": AzureProviderApi.CYCLECLOUD.value, + "vm_size": "Standard_HB120rs_v3", + "resource_group": "my-resource-group", + "location": "eastus2", + "cluster_name": "my-hpc-cluster", + "node_array": "hpc", + "cyclecloud_url": "https://cyclecloud.example.com", + "max_instances": 100, + }, + { + "template_id": "azure-cyclecloud-htc", + "name": "Azure CycleCloud HTC Cluster", + "description": "Add HTC (high-throughput) nodes to a CycleCloud cluster", + "provider_type": "azure", + "provider_api": AzureProviderApi.CYCLECLOUD.value, + "vm_size": "Standard_D4s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "cluster_name": "my-htc-cluster", + "node_array": "htc", + "cyclecloud_url": "https://cyclecloud.example.com", + "max_instances": 500, + }, + ] diff --git a/src/orb/providers/azure/infrastructure/handlers/single_vm_handler.py b/src/orb/providers/azure/infrastructure/handlers/single_vm_handler.py new file mode 100644 index 000000000..d4470f3c3 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/single_vm_handler.py @@ -0,0 +1,653 @@ +"""SingleVM Handler - provisions individual VMs via the Azure Compute SDK. + +This handler is used when ``provider_api == "SingleVM"`` in the template. +It creates standalone Virtual Machines rather than VMSS instances, which +is suitable for long-lived singleton workloads. +""" + +from __future__ import annotations + +import uuid +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any, Optional, cast + +from orb.domain.request.aggregate import Request +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import ( + LaunchError, + TerminationError, +) +from orb.providers.azure.infrastructure.error_utils import ( + canonical_azure_error_code, + extract_azure_error_details, +) +from orb.providers.azure.infrastructure.handlers._network_identity import ( + resolve_network_identity_or_empty_async, +) +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + AzureAcquireHostsResult, + AzureHandler, + AzureHandlerStatusResult, + AzureReleaseContext, + AzureReleaseHostsResult, + AzureSingleVmReleaseProviderData, + AzureStatusProviderData, + AzureSubmittedDeletion, + azure_raise_on_status_error, +) +from orb.providers.azure.infrastructure.handlers.azure_status import resolve_power_state +from orb.providers.azure.infrastructure.sdk_shapes import ( + AzureVmRuntimeStatusProtocol, + AzureVmWithIdentityProtocol, + instance_view_statuses, +) +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentity, +) + +if TYPE_CHECKING: + from orb.domain.base.ports import LoggingPort + from orb.providers.azure.infrastructure.azure_client import AzureClient + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + + +def _azure_resource_not_found_error_type() -> type[Exception]: + """Resolve the Azure SDK's not-found exception lazily.""" + from azure.core.exceptions import ResourceNotFoundError + + return ResourceNotFoundError + + +def _primary_error_message(error_details: dict[str, Any]) -> str: + return str( + error_details.get("error_message") + or error_details.get("message") + or "Failed to submit SingleVM deployment" + ) + + +def _nested_error_messages(error_details: dict[str, Any]) -> list[str]: + nested_details = error_details.get("details") + if not isinstance(nested_details, list): + return [] + + messages: list[str] = [] + for detail in nested_details: + if not isinstance(detail, dict): + continue + message = detail.get("message") + if message in (None, ""): + continue + code = detail.get("code") + messages.append(f"{code}: {message}" if code not in (None, "") else str(message)) + return messages + + +def _format_azure_error_message(error_details: dict[str, Any]) -> str: + """Return the most actionable Azure create failure message.""" + nested_messages = _nested_error_messages(error_details) + if nested_messages: + return "; ".join(nested_messages) + return _primary_error_message(error_details) + + +def _looks_like_uuid(value: str) -> bool: + try: + uuid.UUID(str(value)) + return True + except (ValueError, AttributeError, TypeError): + return False + + +def _build_vm_name_lookup(vms: Iterable[AzureVmWithIdentityProtocol]) -> dict[str, str]: + """Build a lookup that resolves VM names and VM IDs to VM names.""" + lookup: dict[str, str] = {} + for vm in vms: + vm_name = vm.name + if not vm_name: + continue + resolved_name = str(vm_name) + lookup[resolved_name] = resolved_name + + vm_id = vm.vm_id + if vm_id: + lookup[str(vm_id)] = resolved_name + return lookup + + +@injectable +class SingleVMHandler(AzureHandler): + """Handler that creates and manages individual Azure VMs. + + ``provider_api = "SingleVM"`` + """ + + def __init__( + self, + azure_client: AzureClient, + logger: LoggingPort, + *, + azure_native_spec_service: AzureNativeSpecService | None = None, + ) -> None: + """Initialize handler with deployment and optional native-spec service.""" + super().__init__(azure_client=azure_client, logger=logger) + from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, + ) + + self.azure_deployment_service = AzureDeploymentService( + azure_client=self.azure_client, + logger=self._logger, + ) + self.azure_native_spec_service = azure_native_spec_service + + @staticmethod + def _build_status_result( + *, + vm: AzureVmRuntimeStatusProtocol, + resource_group: str, + status: str, + network_identity: AzureNetworkIdentity, + ) -> AzureHandlerStatusResult: + """Build a typed status result for one Azure VM.""" + hw = vm.hardware_profile + availability_zone = vm.zones[0] if vm.zones else None + provider_data: AzureStatusProviderData = { + "resource_id": str(vm.name), + "cloud_host_id": vm.vm_id or vm.name, + "vm_name": str(vm.name), + "resource_group": resource_group, + "location": str(vm.location), + "availability_zone": availability_zone, + "nic_id": network_identity["nic_id"], + "nic_name": network_identity["nic_name"], + "vnet_id": network_identity["vnet_id"], + } + if vm.vm_id is not None: + provider_data["vm_id"] = vm.vm_id + return { + "instance_id": str(vm.name), + "name": str(vm.name), + "status": status, + "private_ip": network_identity["private_ip"], + "public_ip": network_identity["public_ip"], + "launch_time": None, + "instance_type": str(hw.vm_size) if hw and hw.vm_size else None, + "subnet_id": network_identity["subnet_id"], + "vpc_id": network_identity["vnet_id"], + "tags": vm.tags or {}, + "price_type": None, + "provider_type": "azure", + "provider_data": provider_data, + } + + async def acquire_hosts_async( + self, request: Request, template: AzureTemplate + ) -> AzureAcquireHostsResult: + """Async create for individual VMs using async ARM deployment submission.""" + resource_group = template.resource_group.value + location = template.location.value + count = request.requested_count + + subnet_id = self._resolve_subnet_id(template) + if not subnet_id: + raise LaunchError( + message=( + "No subnet specified. Add 'subnet_id' (full ARM resource ID) " + "to the template under subnet_ids or network_config, e.g.: " + "/subscriptions//resourceGroups//providers/" + "Microsoft.Network/virtualNetworks//subnets/" + ), + template_id=template.template_id, + ) + + nsg_id = ( + template.network_config.network_security_group_id if template.network_config else None + ) + accel_net = bool( + template.network_config.accelerated_networking if template.network_config else False + ) + backend_pool_ids = ( + template.network_config.load_balancer_backend_pool_ids + if template.network_config + else [] + ) + inbound_nat_pool_ids = ( + template.network_config.load_balancer_inbound_nat_pool_ids + if template.network_config + else [] + ) + app_gateway_pool_ids = ( + template.network_config.application_gateway_backend_pool_ids + if template.network_config + else [] + ) + public_ip_enabled = bool( + template.network_config.public_ip_enabled if template.network_config else False + ) + + resolved_ssh_keys = list(template.ssh_public_keys) + if template.ssh_key_name and not resolved_ssh_keys: + from orb.providers.azure.infrastructure.services.ssh_key_resolver import ( + AzureComputeSshKeyClientProtocol, + resolve_ssh_keys_async, + ) + + resolved_ssh_keys = await resolve_ssh_keys_async( + ssh_key_name=template.ssh_key_name, + ssh_public_keys=template.ssh_public_keys, + resource_group=template.resource_group.value, + # cast: SDK's SshPublicKeyResource structurally satisfies our + # AzureSshPublicKeyResourceProtocol, but pyright requires + # invariance through the Awaitable wrapper on the operations + # protocol — so the inferred return type differs even though + # the concrete shapes match. + compute_client=cast( + AzureComputeSshKeyClientProtocol, + await self.azure_client.get_async_compute_client(), + ), + ) + if resolved_ssh_keys != list(template.ssh_public_keys): + template = template.model_copy(update={"ssh_public_keys": resolved_ssh_keys}) + + candidate_vm_sizes = template.candidate_vm_sizes + vm_definitions: list[dict[str, Any]] = [] + for _ in range(count): + vm_name = f"vm-{template.template_id}-{uuid.uuid4().hex[:8]}" + vm_definitions.append( + { + "vm_name": vm_name, + "nic_name": f"nic-{vm_name}", + "public_ip_name": f"pip-{vm_name}" if public_ip_enabled else None, + } + ) + + selected_vm_size: Optional[str] = None + submitted_deployment_name: Optional[str] = None + last_error_details: dict[str, Any] = {} + for candidate_vm_size in candidate_vm_sizes: + try: + resolved_vm_definitions: list[dict[str, Any]] = [] + for vm_definition in vm_definitions: + nic_id = self.azure_deployment_service.resource_id_expression( + "Microsoft.Network/networkInterfaces", + vm_definition["nic_name"], + ) + from orb.providers.azure.infrastructure.services.arm_payload_mapper import ( + ArmPayloadMapper, + ) + + vm_params = ArmPayloadMapper.single_vm_payload( + template, + vm_definition["vm_name"], + nic_id, + vm_size_override=candidate_vm_size, + ) + if self.azure_native_spec_service: + merged_params = ( + self.azure_native_spec_service.process_provider_api_spec_with_merge( + template=template, + request=request, + default_payload=vm_params, + extra_context={ + "vm_name": vm_definition["vm_name"], + "nic_id": nic_id, + "vm_size": candidate_vm_size, + }, + ) + ) + if merged_params: + vm_params = merged_params + resolved_vm_definitions.append({**vm_definition, "vm_payload": vm_params}) + + deployment_name = self.azure_deployment_service.build_deployment_name( + "vm", + str(request.request_id), + template.template_id, + candidate_vm_size, + ) + deployment_template = ( + self.azure_deployment_service.build_single_vm_deployment_template( + location=location, + subnet_id=subnet_id, + vm_definitions=resolved_vm_definitions, + enable_accelerated_networking=accel_net, + nsg_id=nsg_id, + load_balancer_backend_pool_ids=backend_pool_ids, + load_balancer_inbound_nat_pool_ids=inbound_nat_pool_ids, + application_gateway_backend_pool_ids=app_gateway_pool_ids, + ) + ) + submitted_deployment_name = ( + await self.azure_deployment_service.submit_template_deployment_async( + resource_group=resource_group, + deployment_name=deployment_name, + template=deployment_template, + ) + ) + selected_vm_size = candidate_vm_size + break + except Exception as exc: + error_details = extract_azure_error_details(exc) + last_error_details = { + "error_code": self._classify_provisioning_error(exc), + "error_message": error_details["message"], + "resource_group": resource_group, + "instance_type": candidate_vm_size, + "status_code": error_details["status_code"], + "raw_error_code": error_details["raw_error_code"], + "details": error_details["details"], + } + if not self._is_capacity_error(exc): + break + + if submitted_deployment_name is None or selected_vm_size is None: + raise LaunchError( + message=_format_azure_error_message(last_error_details), + template_id=template.template_id, + error_code=last_error_details.get("error_code"), + details=last_error_details, + ) + + created_ids = [vm_definition["vm_name"] for vm_definition in vm_definitions] + operation_tracking = [ + { + "vm_name": vm_definition["vm_name"], + "nic_name": vm_definition["nic_name"], + "public_ip_name": vm_definition["public_ip_name"], + "selected_vm_size": selected_vm_size, + } + for vm_definition in vm_definitions + ] + self._logger.info( + "Submitted create deployment '%s' for %d VM(s)", + submitted_deployment_name, + len(created_ids), + ) + return { + "success": True, + "resource_ids": created_ids, + "instances": [], + "error_message": None, + "provider_data": { + "resource_group": resource_group, + "location": location, + "submitted_count": len(created_ids), + "operation_status": "submitted", + "fulfillment_final": True, + "deployment_name": submitted_deployment_name, + "error_codes": [], + "fleet_errors": [], + "submitted_vms": operation_tracking, + }, + } + + async def check_hosts_status_async(self, request: Request) -> list[AzureHandlerStatusResult]: + """Async status query for individual VM IDs using the Azure async Compute SDK.""" + resource_ids: list[str] = request.resource_ids or [] + raise_on_status_error = azure_raise_on_status_error(request) + resource_group = (request.metadata or {}).get( + "resource_group" + ) or self.azure_client.resource_group + if not resource_group: + message = "Cannot resolve resource_group for status check" + self._logger.error(message) + if resource_ids: + raise RuntimeError(message) + return [] + + results: list[AzureHandlerStatusResult] = [] + status_errors: list[str] = [] + compute = await self.azure_client.get_async_compute_client() + resolved_vm_names = await self._resolve_vm_names_async(resource_group, resource_ids) + + for vm_name in resolved_vm_names: + try: + vm = await compute.virtual_machines.get( + resource_group_name=resource_group, + vm_name=vm_name, + expand="instanceView", + ) + network_identity = await resolve_network_identity_or_empty_async( + logger=self._logger, + target_label=f"VM '{vm_name}'", + resolver=lambda: self.azure_client.resolve_network_identity_from_vm_async(vm), + ) + statuses = instance_view_statuses(vm.instance_view) + results.append( + self._build_status_result( + # cast: SDK 38 exposes hardware_profile/instance_view/vm_id + # via __flattened_items __getattr__, invisible to static + # checkers though documented as the SDK's public surface. + vm=cast(AzureVmRuntimeStatusProtocol, vm), + resource_group=resource_group, + status=( + resolve_power_state(statuses) if statuses is not None else "unknown" + ), + network_identity=network_identity, + ) + ) + except Exception as exc: + error_message = f"Failed to get status for VM '{vm_name}': {exc}" + self._logger.error(error_message) + status_errors.append(error_message) + all_requested_vms_failed = bool(resolved_vm_names) and not results + if status_errors and (raise_on_status_error or all_requested_vms_failed): + raise RuntimeError("; ".join(status_errors)) + return results + + @staticmethod + def _raise_release_failures( + *, + machine_ids: list[str], + resource_group: str, + submitted_deletions: list[AzureSubmittedDeletion], + failed_deletions: list[AzureSubmittedDeletion], + ) -> None: + """Raise an aggregated termination error when any delete submissions fail.""" + if not failed_deletions: + return + + failed_requested_ids = [ + requested_id + for deletion in failed_deletions + if (requested_id := deletion.get("requested_id")) is not None + ] + raise TerminationError( + (f"Failed to submit deletion for {len(failed_deletions)} of {len(machine_ids)} VM(s)"), + resource_ids=failed_requested_ids, + details={ + "resource_group": resource_group, + "submitted_deletions": submitted_deletions, + "failed_deletions": failed_deletions, + }, + ) + + @staticmethod + def _build_release_result( + resource_group: str, + submitted_deletions: list[AzureSubmittedDeletion], + ) -> AzureReleaseHostsResult: + """Build the typed release result payload for SingleVM deletes.""" + provider_data: AzureSingleVmReleaseProviderData = { + "resource_group": resource_group, + "operation_status": "submitted", + "submitted_deletions": submitted_deletions, + } + return {"provider_data": provider_data} + + async def release_hosts_async( + self, + machine_ids: list[str], + resource_id: str, + context: Optional[AzureReleaseContext] = None, + ) -> Optional[AzureReleaseHostsResult]: + """Async delete submission for individual VMs using the Azure async Compute SDK.""" + resource_group = self._resolve_release_resource_group( + context=context, + machine_ids=machine_ids, + ) + compute = await self.azure_client.get_async_compute_client() + vm_names = await self._resolve_vm_names_async(resource_group, machine_ids) + submitted_deletions: list[AzureSubmittedDeletion] = [] + failed_deletions: list[AzureSubmittedDeletion] = [] + for original_id, vm_name in zip(machine_ids, vm_names, strict=True): + try: + self._logger.info("Deleting VM '%s' (requested id='%s')", vm_name, original_id) + await compute.virtual_machines.begin_delete( + resource_group_name=resource_group, + vm_name=vm_name, + ) + submitted_deletions.append({"requested_id": str(original_id), "vm_name": vm_name}) + except Exception as exc: + self._logger.error("Failed to delete VM '%s': %s", vm_name, exc) + failed_deletions.append( + {"requested_id": str(original_id), "vm_name": vm_name, "error": str(exc)} + ) + + self._raise_release_failures( + machine_ids=machine_ids, + resource_group=resource_group, + submitted_deletions=submitted_deletions, + failed_deletions=failed_deletions, + ) + return self._build_release_result(resource_group, submitted_deletions) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _finalize_resolved_vm_names( + *, + resource_group: str, + machine_ids: list[str], + resolved: list[Optional[str]], + logger: LoggingPort, + ) -> list[str]: + """Apply unresolved lookups, preserve input order, and log any remapping.""" + ordered_resolved = [ + resolved_name if resolved_name is not None else str(machine_id) + for machine_id, resolved_name in zip(machine_ids, resolved, strict=True) + ] + if ordered_resolved != [str(mid) for mid in machine_ids]: + logger.debug( + "Resolved SingleVM IDs in resource_group '%s': %s -> %s", + resource_group, + machine_ids, + ordered_resolved, + ) + return ordered_resolved + + async def _resolve_vm_names_async( + self, resource_group: str, machine_ids: list[str] + ) -> list[str]: + """Resolve a list of mixed identifiers to canonical Azure VM names. + + Each input is treated as a vm_name first and looked up via a direct + ``virtual_machines.get`` (cheap when callers already have names). Inputs + that look like Azure ``vm_id`` GUIDs, or that 404 on the direct lookup, + are deferred to a single ``virtual_machines.list`` over the resource + group and matched by vm_id. The output preserves input order. + + Best-effort: any unhandled error falls back to returning the input list + as-is so callers can still attempt downstream operations with the + original identifiers. + """ + if not machine_ids: + return [] + + try: + compute = await self.azure_client.get_async_compute_client() + resolved: list[Optional[str]] = [None] * len(machine_ids) + unresolved_indices: list[int] = [] + + for index, machine_id in enumerate(machine_ids): + machine_id_str = str(machine_id) + if _looks_like_uuid(machine_id_str): + unresolved_indices.append(index) + continue + + try: + vm = await compute.virtual_machines.get( + resource_group_name=resource_group, + vm_name=machine_id_str, + ) + resolved[index] = str(vm.name or machine_id_str) + except _azure_resource_not_found_error_type(): + unresolved_indices.append(index) + + if unresolved_indices: + pager = compute.virtual_machines.list(resource_group_name=resource_group) + # cast: SDK 38 exposes vm_id via __flattened_items __getattr__, + # invisible to static checkers though documented as the SDK's + # public surface. + lookup = _build_vm_name_lookup( + [cast(AzureVmWithIdentityProtocol, vm) async for vm in pager] + ) + + for index in unresolved_indices: + machine_id = str(machine_ids[index]) + resolved[index] = lookup.get(machine_id, machine_id) + + return self._finalize_resolved_vm_names( + resource_group=resource_group, + machine_ids=machine_ids, + resolved=resolved, + logger=self._logger, + ) + except Exception as exc: + self._logger.warning( + "Failed to resolve VM names, using provided IDs directly: %s", + exc, + ) + return [str(machine_id) for machine_id in machine_ids] + + @staticmethod + def _classify_provisioning_error(exc: Exception) -> str: + """Map common Azure provisioning failures to stable error codes.""" + return canonical_azure_error_code(exc) + + @staticmethod + def _is_capacity_error(exc: Exception) -> bool: + """Return True when trying another candidate VM size is reasonable.""" + return canonical_azure_error_code(exc) in { + "AllocationFailed", + "ZonalAllocationFailed", + "SkuNotAvailable", + "OverconstrainedAllocationRequest", + } + + # NIC and Public IP cleanup is handled natively by Azure via + # deleteOption: "Delete" on the NIC and Public IP references in the + # ARM template. When a VM is deleted, Azure cascades the deletion + # through NIC → Public IP automatically. For Flexible VMSS this is + # the default behaviour. No ORB-managed rollback methods are needed. + # See: https://learn.microsoft.com/en-us/azure/virtual-machines/delete + + @classmethod + def get_example_templates(cls) -> list[dict[str, Any]]: + """Return example SingleVM template configurations.""" + return [ + { + "template_id": "azure-singlevm-linux", + "name": "Azure Single VM Linux", + "description": "Individual Ubuntu 22.04 VM on Standard_D2s_v5", + "provider_type": "azure", + "provider_api": AzureProviderApi.SINGLE_VM.value, + "vm_size": "Standard_D2s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "max_instances": 5, + }, + ] diff --git a/src/orb/providers/azure/infrastructure/handlers/vmss_handler.py b/src/orb/providers/azure/infrastructure/handlers/vmss_handler.py new file mode 100644 index 000000000..172301c22 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/handlers/vmss_handler.py @@ -0,0 +1,1240 @@ +"""VMSS Handler - provisions VM Scale Sets via the Azure Compute SDK. + +This is the primary handler for the Azure provider + +It handles: +- Creating a VMSS with Flexible orchestration (default) or Uniform +- Listing instances in a VMSS for status checks +- Deleting VMSS instances and the scale set itself + +Important limitation: +- Azure Flexible VMSS does not expose an AWS-ASG-style "detach these exact + instances and decrement desired capacity for them" flow. +- Scaling in first can let Azure choose different victims than the caller + requested, so VMSS termination in this provider uses exact-instance deletion + and, when needed, a narrow follow-up that deletes the VMSS after it is empty. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional, TypedDict, cast + +from orb.domain.request.aggregate import Request +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import ( + AzureProviderApi, + AzureVMSSOrchestrationMode, +) +from orb.providers.azure.exceptions.azure_exceptions import ( + AzureValidationError, + QuotaExceededError, + TerminationError, + VMSSCreationError, + VMSSNotFoundError, +) +from orb.providers.azure.infrastructure.error_codes import ProviderErrorEntry +from orb.providers.azure.infrastructure.error_utils import ( + classify_azure_error, +) +from orb.providers.azure.infrastructure.handlers._network_identity import ( + resolve_network_identity_or_empty_async, +) +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + AzureAcquireHostsResult, + AzureHandler, + AzureHandlerStatusResult, + AzureReleaseContext, + AzureReleaseHostsResult, + AzureStatusProviderData, + AzureSubmittedDeletion, + AzureVmssReleaseProviderData, + azure_raise_on_status_error, +) +from orb.providers.azure.infrastructure.handlers.azure_status import resolve_power_state +from orb.providers.azure.infrastructure.sdk_shapes import ( + AzureVmRuntimeStatusProtocol, + AzureVmWithIdentityProtocol, + instance_view_statuses, +) +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentity, +) +from orb.providers.azure.infrastructure.vmss_cleanup import PendingVmssCleanup + +if TYPE_CHECKING: + from azure.mgmt.compute.models import OrchestrationMode as SdkOrchestrationMode + + from orb.domain.base.ports import LoggingPort + from orb.providers.azure.infrastructure.azure_client import AzureClient + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + +def _status_attr(status: Any, attr: str, default: Any = None) -> Any: + """Read Azure status attributes from SDK objects or attribute-based test doubles. + + getattr is necessary here: callers pass heterogeneous status-like objects + (SDK InstanceViewStatus, plain dicts wrapped in SimpleNamespace, etc.) and + the requested attribute varies per call-site. + """ + return getattr(status, attr, default) + + +@dataclass(frozen=True) +class _AzureVmIdentity: + """Normalized identity fields across Azure VM and VMSS VM SDK shapes.""" + + instance_id: str + vm_id: str + vm_name: Optional[str] + + +@dataclass(frozen=True) +class _VmssReleasePlan: + """Precomputed VMSS member-release inputs for the async VMSS release flow.""" + + resource_group: str + vmss_name: str + orchestration_mode: AzureVMSSOrchestrationMode + current_members: list[AzureHandlerStatusResult] + resolved_instance_ids: list[str] + resolved_vm_names: list[str] + delete_vmss_when_empty: bool + + +class _VmssCleanupSubmissionState(TypedDict): + """Immediate VMSS delete-submission state carried into release provider data.""" + + delete_submitted: bool + delete_retry_pending: bool + last_delete_error: Optional[str] + + +def _build_vmss_delete_instance_ids(instance_ids: list[str]) -> Any: + """Build the VMSS delete payload using the SDK model when available.""" + try: + from azure.mgmt.compute.models import VirtualMachineScaleSetVMInstanceRequiredIDs + except ImportError: + return {"instance_ids": instance_ids} + return VirtualMachineScaleSetVMInstanceRequiredIDs(instance_ids=instance_ids) + + +def _read_vm_identity(vm: AzureVmWithIdentityProtocol) -> _AzureVmIdentity: + """Normalize VM identity across VMSS VM and regular VM SDK objects. + + Microsoft documents `VirtualMachineScaleSetVM` with `name`, `instance_id`, + and `vm_id`, while regular `VirtualMachine` objects expose `name` and + `vm_id` but not `instance_id`. When Azure returns a regular VM object, + `name` becomes the stable machine identifier when `instance_id` is absent. + """ + vm_name = vm.name + instance_id = str(_status_attr(vm, "instance_id", "") or vm_name or "") + vm_id = str(vm.vm_id or instance_id) + return _AzureVmIdentity( + instance_id=instance_id, + vm_id=vm_id, + vm_name=vm_name, + ) + + +@injectable +class VMSSHandler(AzureHandler): + """Handler that creates and manages VMSS resources. + + ``provider_api = "VMSS"`` or ``"VMSSUniform"`` + """ + + def __init__( + self, + azure_client: AzureClient, + logger: LoggingPort, + *, + azure_native_spec_service: AzureNativeSpecService | None = None, + azure_resource_manager: AzureResourceManager | None = None, + ) -> None: + """Initialize handler with explicit optional infrastructure services.""" + super().__init__(azure_client=azure_client, logger=logger) + self.azure_native_spec_service = azure_native_spec_service + self.azure_resource_manager = azure_resource_manager + + async def acquire_hosts_async( + self, request: Request, template: AzureTemplate + ) -> AzureAcquireHostsResult: + """Async VMSS create using the Azure async Compute SDK.""" + vmss_name = template.vmss_name or f"vmss-{template.template_id}-{uuid.uuid4().hex[:8]}" + resource_group = template.resource_group.value + location = template.location.value + + self._logger.info( + "Creating VMSS '%s' in resource group '%s' (location=%s, capacity=%d)", + vmss_name, + resource_group, + location, + request.requested_count, + ) + + try: + resolved_template = template + if not resolved_template.network_config: + subnet_id = self._resolve_subnet_id(resolved_template) + if subnet_id: + from orb.providers.azure.domain.template.value_objects import AzureNetworkConfig + + resolved_template = resolved_template.model_copy( + update={"network_config": AzureNetworkConfig(subnet_id=subnet_id)} + ) + self._logger.debug("Auto-created network_config from subnet_id: %s", subnet_id) + else: + raise AzureValidationError( + "No subnet specified. Add 'subnet_id' (full ARM resource ID) " + "to the template in subnet_ids or network_config, e.g.: " + "/subscriptions//resourceGroups//providers/" + "Microsoft.Network/virtualNetworks//subnets/", + details={"template_id": template.template_id}, + error_code="InvalidParameter", + ) + + from orb.providers.azure.infrastructure.services.arm_payload_mapper import ( + ArmPayloadMapper, + ) + from orb.providers.azure.infrastructure.services.ssh_key_resolver import ( + AzureComputeSshKeyClientProtocol, + resolve_ssh_keys_async, + ) + + arm_payload = ArmPayloadMapper.vmss_payload(resolved_template) + if self.azure_native_spec_service: + merged_payload = ( + self.azure_native_spec_service.process_provider_api_spec_with_merge( + template=resolved_template, + request=request, + default_payload=arm_payload, + extra_context={"vmss_name": vmss_name}, + ) + ) + if merged_payload: + arm_payload = merged_payload + + arm_payload.setdefault("sku", {}) + arm_payload["sku"]["capacity"] = request.requested_count + + compute = await self.azure_client.get_async_compute_client() + if resolved_template.ssh_key_name and not resolved_template.ssh_public_keys: + resolved_keys = await resolve_ssh_keys_async( + ssh_key_name=resolved_template.ssh_key_name, + ssh_public_keys=resolved_template.ssh_public_keys, + resource_group=resolved_template.resource_group.value, + # cast: SDK's SshPublicKeyResource structurally satisfies our + # AzureSshPublicKeyResourceProtocol, but pyright requires + # invariance through the Awaitable wrapper on the operations + # protocol — so the inferred return type differs even though + # the concrete shapes match. + compute_client=cast(AzureComputeSshKeyClientProtocol, compute), + ) + vm_profile = arm_payload["properties"]["virtualMachineProfile"] + vm_profile["osProfile"]["linuxConfiguration"] = { + "disablePasswordAuthentication": True, + "ssh": { + "publicKeys": [ + { + "path": f"/home/{template.admin_username}/.ssh/authorized_keys", + "keyData": key, + } + for key in resolved_keys + ], + }, + } + + vmss_operations: Any = compute.virtual_machine_scale_sets + await vmss_operations.begin_create_or_update( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + parameters=arm_payload, + ) + + self._logger.info( + "Submitted native VMSS create for '%s'; tracking will continue via status checks", + vmss_name, + ) + return { + "success": True, + "resource_ids": [vmss_name], + "instances": [], + "error_message": None, + "provider_data": { + "vmss_name": vmss_name, + "resource_group": resource_group, + "location": location, + "provisioning_state": "creating", + "operation_status": "submitted", + "error_codes": [], + "fulfillment_final": True, + }, + } + except AzureValidationError: + raise + except Exception as exc: + error_msg = f"Failed to create VMSS '{vmss_name}': {exc}" + self._logger.error(error_msg) + category, error_code = classify_azure_error(exc) + if category == "quota": + raise QuotaExceededError(error_msg, error_code=error_code) from exc + if category == "validation": + raise AzureValidationError(error_msg, error_code=error_code) from exc + raise VMSSCreationError( + message=error_msg, + template_id=template.template_id, + vmss_name=vmss_name, + error_code=error_code, + ) from exc + + async def check_hosts_status_async(self, request: Request) -> list[AzureHandlerStatusResult]: + """Async status query for VMSS members using the Azure async Compute SDK.""" + resource_ids = request.resource_ids + if not resource_ids: + self._logger.warning("check_hosts_status called with no resource_ids") + return [] + + all_instances: list[AzureHandlerStatusResult] = [] + status_errors: list[str] = [] + raise_on_status_error = azure_raise_on_status_error(request) + + resource_group = (request.metadata or {}).get( + "resource_group" + ) or self.azure_client.resource_group + for vmss_name in resource_ids: + if not resource_group: + error_message = f"Cannot resolve resource_group for VMSS '{vmss_name}'" + self._logger.error(error_message) + status_errors.append(error_message) + continue + + try: + instances = await self._list_vmss_instances_async( + resource_group, vmss_name, include_instance_view=True + ) + all_instances.extend(instances) + except Exception as exc: + error_message = f"Failed to list instances for VMSS '{vmss_name}': {exc}" + self._logger.error(error_message) + status_errors.append(error_message) + + all_requested_vmss_failed = bool(resource_ids) and not all_instances + if status_errors and (raise_on_status_error or all_requested_vmss_failed): + raise RuntimeError("; ".join(status_errors)) + + return all_instances + + async def release_hosts_async( + self, + machine_ids: list[str], + resource_id: str, + context: Optional[AzureReleaseContext] = None, + ) -> Optional[AzureReleaseHostsResult]: + """Async delete submission for VMSS members using the Azure async Compute SDK.""" + if not machine_ids: + return None + + release_plan = await self._build_release_plan_async( + machine_ids=machine_ids, + resource_id=resource_id, + context=context, + ) + resource_group = release_plan.resource_group + vmss_name = release_plan.vmss_name + compute = await self.azure_client.get_async_compute_client() + + self._log_release_submission(vmss_name=vmss_name, machine_ids=machine_ids) + try: + cleanup_submission_state: _VmssCleanupSubmissionState = { + "delete_submitted": False, + "delete_retry_pending": False, + "last_delete_error": None, + } + if release_plan.orchestration_mode == AzureVMSSOrchestrationMode.FLEXIBLE: + submitted_deletions: list[AzureSubmittedDeletion] = [] + failed_deletions: list[AzureSubmittedDeletion] = [] + for requested_id, vm_name in zip( + machine_ids, + release_plan.resolved_vm_names, + strict=True, + ): + try: + await compute.virtual_machines.begin_delete( + resource_group_name=resource_group, + vm_name=str(vm_name), + ) + submitted_deletions.append( + { + "requested_id": str(requested_id), + "vm_name": str(vm_name), + } + ) + except Exception as exc: + self._logger.error( + "Failed to delete VMSS flexible member '%s' from '%s': %s", + vm_name, + vmss_name, + exc, + ) + failed_deletions.append( + { + "requested_id": str(requested_id), + "vm_name": str(vm_name), + "error": str(exc), + } + ) + self._raise_flexible_release_failures( + machine_ids=machine_ids, + resource_group=resource_group, + vmss_name=vmss_name, + submitted_deletions=submitted_deletions, + failed_deletions=failed_deletions, + ) + if release_plan.delete_vmss_when_empty: + cleanup_submission_state = await self._submit_vmss_delete_if_emptying_async( + resource_group=resource_group, + vmss_name=vmss_name, + ) + return self._build_flexible_release_result( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + delete_vmss_when_empty=release_plan.delete_vmss_when_empty, + cleanup_submission_state=cleanup_submission_state, + submitted_deletions=submitted_deletions, + failed_deletions=failed_deletions, + ) + + await compute.virtual_machine_scale_sets.begin_delete_instances( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + vm_instance_i_ds=_build_vmss_delete_instance_ids( + release_plan.resolved_instance_ids + ), + ) + if release_plan.delete_vmss_when_empty: + cleanup_submission_state = await self._submit_vmss_delete_if_emptying_async( + resource_group=resource_group, + vmss_name=vmss_name, + ) + return self._build_uniform_release_result( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + delete_vmss_when_empty=release_plan.delete_vmss_when_empty, + cleanup_submission_state=cleanup_submission_state, + resolved_instance_ids=release_plan.resolved_instance_ids, + ) + except TerminationError: + raise + except Exception as exc: + raise TerminationError( + f"Failed to delete instances from VMSS '{vmss_name}': {exc}", + resource_ids=machine_ids, + ) from exc + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _log_release_submission(self, *, vmss_name: str, machine_ids: list[str]) -> None: + """Log the start of a VMSS member-delete submission.""" + self._logger.info( + "Deleting %d instance(s) from VMSS '%s'", + len(machine_ids), + vmss_name, + ) + + @staticmethod + def _raise_flexible_release_failures( + *, + machine_ids: list[str], + resource_group: str, + vmss_name: str, + submitted_deletions: list[AzureSubmittedDeletion], + failed_deletions: list[AzureSubmittedDeletion], + ) -> None: + """Raise when any flexible VMSS member delete submission fails.""" + if not failed_deletions: + return + + failed_requested_ids = [ + requested_id + for deletion in failed_deletions + if (requested_id := deletion.get("requested_id")) is not None + ] + raise TerminationError( + ( + f"Failed to submit deletion for {len(failed_deletions)} of " + f"{len(machine_ids)} VMSS member(s)" + ), + resource_ids=failed_requested_ids, + details={ + "resource_group": resource_group, + "vmss_name": vmss_name, + "submitted_deletions": submitted_deletions, + "failed_deletions": failed_deletions, + }, + ) + + async def _build_release_plan_async( + self, + *, + machine_ids: list[str], + resource_id: str, + context: Optional[AzureReleaseContext], + ) -> _VmssReleasePlan: + """Precompute everything the release submission needs for one VMSS. + + Resolves the resource group, queries the VMSS to determine its + orchestration mode, lists current members once, and produces the + delete identifiers in the form the SDK expects for that mode: + + - Uniform VMSS deletes by ``instance_id`` — populated from + ``current_members``; ``resolved_vm_names`` is empty. + - Flexible VMSS deletes individual VMs by ``vm_name`` — populated + from ``current_members``; ``resolved_instance_ids`` mirrors the + input ``machine_ids`` (Flexible accepts either, but we forward + the originals for traceability). + + Splitting plan-building from submission keeps the submission path + flat and lets tests exercise the plan shape without mocking the + delete calls. + """ + resource_group = self._resolve_release_resource_group( + machine_ids=machine_ids, + context=context, + ) + vmss_name = resource_id + orchestration_mode = await self._get_vmss_orchestration_mode_async( + resource_group, vmss_name + ) + current_members = await self._list_vmss_instances_async( + resource_group=resource_group, + vmss_name=vmss_name, + include_instance_view=False, + orchestration_mode=orchestration_mode, + ) + resolved_instance_ids = ( + await self._resolve_vmss_instance_ids_async( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + current_members=current_members, + ) + if orchestration_mode != AzureVMSSOrchestrationMode.FLEXIBLE + else [str(machine_id) for machine_id in machine_ids] + ) + resolved_vm_names = ( + self._resolve_flexible_vm_names_from_members( + machine_ids=machine_ids, + current_members=current_members, + logger=self._logger, + vmss_name=vmss_name, + ) + if orchestration_mode == AzureVMSSOrchestrationMode.FLEXIBLE + else [] + ) + return _VmssReleasePlan( + resource_group=resource_group, + vmss_name=vmss_name, + orchestration_mode=orchestration_mode, + current_members=current_members, + resolved_instance_ids=resolved_instance_ids, + resolved_vm_names=resolved_vm_names, + delete_vmss_when_empty=self._should_delete_vmss_when_empty( + orchestration_mode=orchestration_mode, + machine_ids=machine_ids, + current_members=current_members, + resolved_instance_ids=resolved_instance_ids, + resolved_vm_names=resolved_vm_names, + ), + ) + + @staticmethod + def _build_pending_resource_cleanup( + *, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + member_delete_submitted: bool = False, + delete_submitted: bool = False, + delete_retry_pending: bool = False, + last_delete_error: Optional[str] = None, + ) -> dict[str, Any]: + return PendingVmssCleanup.create( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + delete_vmss_when_empty=True, + member_delete_submitted=member_delete_submitted, + delete_submitted=delete_submitted, + delete_retry_pending=delete_retry_pending, + last_delete_error=last_delete_error, + ).to_metadata() + + def _release_cleanup_provider_data( + self, + *, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + delete_vmss_when_empty: bool, + member_delete_submitted: bool = False, + delete_submitted: bool = False, + delete_retry_pending: bool = False, + last_delete_error: Optional[str] = None, + ) -> AzureVmssReleaseProviderData: + if not delete_vmss_when_empty: + return {} + + return { + "pending_resource_cleanup": cast( + Any, + self._build_pending_resource_cleanup( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + member_delete_submitted=member_delete_submitted, + delete_submitted=delete_submitted, + delete_retry_pending=delete_retry_pending, + last_delete_error=last_delete_error, + ), + ) + } + + def _build_uniform_release_result( + self, + *, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + delete_vmss_when_empty: bool, + cleanup_submission_state: _VmssCleanupSubmissionState, + resolved_instance_ids: list[str], + ) -> AzureReleaseHostsResult: + """Build provider data for a uniform VMSS member-delete submission.""" + provider_data: AzureVmssReleaseProviderData = { + "resource_group": resource_group, + "vmss_name": vmss_name, + "operation_status": "submitted", + "resolved_instance_ids": resolved_instance_ids, + **self._release_cleanup_provider_data( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + delete_vmss_when_empty=delete_vmss_when_empty, + member_delete_submitted=True, + delete_submitted=cleanup_submission_state["delete_submitted"], + delete_retry_pending=cleanup_submission_state["delete_retry_pending"], + last_delete_error=cleanup_submission_state["last_delete_error"], + ), + } + return {"provider_data": provider_data} + + def _build_flexible_release_result( + self, + *, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + delete_vmss_when_empty: bool, + cleanup_submission_state: _VmssCleanupSubmissionState, + submitted_deletions: list[AzureSubmittedDeletion], + failed_deletions: list[AzureSubmittedDeletion], + ) -> AzureReleaseHostsResult: + """Build provider data for a flexible VMSS member-delete submission.""" + provider_data: AzureVmssReleaseProviderData = { + "resource_group": resource_group, + "vmss_name": vmss_name, + "operation_status": "submitted", + "submitted_deletions": submitted_deletions, + **self._release_cleanup_provider_data( + resource_group=resource_group, + vmss_name=vmss_name, + machine_ids=machine_ids, + delete_vmss_when_empty=delete_vmss_when_empty, + member_delete_submitted=True, + delete_submitted=cleanup_submission_state["delete_submitted"], + delete_retry_pending=cleanup_submission_state["delete_retry_pending"], + last_delete_error=cleanup_submission_state["last_delete_error"], + ), + } + if failed_deletions: + provider_data["failed_deletions"] = failed_deletions + return {"provider_data": provider_data} + + async def _submit_vmss_delete_if_emptying_async( + self, + *, + resource_group: str, + vmss_name: str, + ) -> _VmssCleanupSubmissionState: + """Async best-effort VMSS delete submission when this return should empty the scale set.""" + try: + compute = await self.azure_client.get_async_compute_client() + await compute.virtual_machine_scale_sets.begin_delete( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + return self._vmss_delete_submission_result(vmss_name=vmss_name) + except Exception as exc: + return self._vmss_delete_submission_result(vmss_name=vmss_name, exc=exc) + + def _should_delete_vmss_when_empty( + self, + *, + orchestration_mode: AzureVMSSOrchestrationMode, + machine_ids: list[str], + current_members: list[AzureHandlerStatusResult], + resolved_instance_ids: list[str], + resolved_vm_names: list[str] | None = None, + ) -> bool: + """Return whether deleting these exact members would leave the VMSS empty.""" + if not machine_ids: + return False + if not current_members: + return False + + current_member_ids = { + str(instance_id) + for instance_id in (member.get("instance_id") for member in current_members) + if instance_id not in (None, "") + } + if len(current_member_ids) != len(current_members): + return False + + if orchestration_mode == AzureVMSSOrchestrationMode.FLEXIBLE: + requested_ids = { + str(machine_id) + for machine_id in (resolved_vm_names or machine_ids) + if machine_id not in (None, "") + } + return bool(requested_ids) and requested_ids == current_member_ids + + requested_ids = { + str(instance_id) + for instance_id in resolved_instance_ids + if instance_id not in (None, "") + } + return bool(requested_ids) and requested_ids == current_member_ids + + @staticmethod + def _resolve_flexible_vm_names_from_members( + *, + machine_ids: list[str], + current_members: list[AzureHandlerStatusResult], + logger: LoggingPort, + vmss_name: str, + ) -> list[str]: + """Resolve Flexible VMSS requested IDs to Azure VM names.""" + if not machine_ids: + return [] + + lookup: dict[str, str] = {} + for member in current_members: + provider_data = member.get("provider_data") or {} + + vm_name = ( + provider_data.get("vm_name") or member.get("name") or member.get("instance_id") + ) + if not vm_name: + continue + resolved_vm_name = str(vm_name) + + for candidate in ( + member.get("instance_id"), + member.get("name"), + provider_data.get("vmss_instance_id"), + provider_data.get("vm_id"), + provider_data.get("vm_name"), + ): + if candidate not in (None, ""): + lookup[str(candidate)] = resolved_vm_name + + unresolved_ids = [ + str(machine_id) for machine_id in machine_ids if str(machine_id) not in lookup + ] + if unresolved_ids: + raise TerminationError( + f"Could not resolve {len(unresolved_ids)} requested Flexible VMSS member ID(s)", + resource_ids=unresolved_ids, + details={ + "vmss_name": vmss_name, + "unresolved_ids": unresolved_ids, + "available_member_ids": sorted(lookup), + }, + ) + + resolved = [lookup[str(machine_id)] for machine_id in machine_ids] + if resolved != [str(machine_id) for machine_id in machine_ids]: + logger.debug( + "Resolved Flexible VMSS machine IDs for '%s': %s -> %s", + vmss_name, + machine_ids, + resolved, + ) + return resolved + + @staticmethod + def _resolve_vmss_instance_ids_from_members( + *, + machine_ids: list[str], + current_members: list[AzureHandlerStatusResult], + logger: LoggingPort, + vmss_name: str, + ) -> list[str]: + """Resolve mixed IDs (vm_id/vm_name/instance_id) using already-fetched VMSS members.""" + if not machine_ids: + return [] + + lookup: dict[str, str] = {} + for vm in current_members: + vmss_instance_id = str(vm.get("instance_id", "") or "") + if not vmss_instance_id: + continue + lookup[vmss_instance_id] = vmss_instance_id + + provider_data = vm.get("provider_data") or {} + vm_id = provider_data.get("vm_id") + if vm_id: + lookup[str(vm_id)] = vmss_instance_id + + vm_name = provider_data.get("vm_name") + if vm_name: + lookup[str(vm_name)] = vmss_instance_id + + unresolved_ids = [ + str(machine_id) for machine_id in machine_ids if str(machine_id) not in lookup + ] + if unresolved_ids: + raise TerminationError( + f"Could not resolve {len(unresolved_ids)} requested VMSS member ID(s)", + resource_ids=unresolved_ids, + details={ + "vmss_name": vmss_name, + "unresolved_ids": unresolved_ids, + "available_member_ids": sorted(lookup), + }, + ) + + resolved = [lookup[str(machine_id)] for machine_id in machine_ids] + if resolved != [str(mid) for mid in machine_ids]: + logger.debug( + "Resolved VMSS machine IDs for '%s': %s -> %s", + vmss_name, + machine_ids, + resolved, + ) + return resolved + + async def _resolve_vmss_instance_ids_async( + self, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + current_members: Optional[list[AzureHandlerStatusResult]] = None, + ) -> list[str]: + """Async resolve mixed IDs (vm_id/vm_name/instance_id) to VMSS instance IDs.""" + if not machine_ids: + return [] + if current_members is None: + current_members = await self._list_vmss_instances_async( + resource_group=resource_group, + vmss_name=vmss_name, + include_instance_view=False, + orchestration_mode=AzureVMSSOrchestrationMode.UNIFORM, + ) + return self._resolve_vmss_instance_ids_from_members( + machine_ids=machine_ids, + current_members=current_members, + logger=self._logger, + vmss_name=vmss_name, + ) + + async def _list_vmss_instances_async( + self, + resource_group: str, + vmss_name: str, + include_instance_view: bool = False, + orchestration_mode: Optional[AzureVMSSOrchestrationMode] = None, + ) -> list[AzureHandlerStatusResult]: + """Async list of normalised instance dicts for a VMSS.""" + compute = await self.azure_client.get_async_compute_client() + if orchestration_mode is None: + orchestration_mode = await self._get_vmss_orchestration_mode_async( + resource_group, vmss_name + ) + + if orchestration_mode == AzureVMSSOrchestrationMode.FLEXIBLE: + return await self._list_flexible_vmss_instances_async( + resource_group=resource_group, + vmss_name=vmss_name, + include_instance_view=include_instance_view, + ) + + try: + pager = compute.virtual_machine_scale_set_vms.list( + resource_group_name=resource_group, + virtual_machine_scale_set_name=vmss_name, + expand=self._vmss_expand_arg(include_instance_view), + ) + vms = [vm async for vm in pager] + except Exception as exc: + raise VMSSNotFoundError( + f"Could not list VMs in VMSS '{vmss_name}': {exc}", + vmss_name=vmss_name, + ) from exc + + return [ + # cast: SDK 38 exposes hardware_profile/instance_view/vm_id via the + # __flattened_items __getattr__ shim, invisible to static checkers + # though documented as the SDK's public surface. + await self._normalise_vm_async( + cast(AzureVmRuntimeStatusProtocol, vm), vmss_name, resource_group + ) + for vm in vms + ] + + async def get_vmss_resource_errors_async( + self, + resource_group: str, + vmss_name: str, + ) -> list[ProviderErrorEntry]: + """Return VMSS-level provisioning errors via the async Azure SDK.""" + compute = await self.azure_client.get_async_compute_client() + try: + vmss = await compute.virtual_machine_scale_sets.get( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + except Exception as exc: + self._logger.warning( + "Failed to fetch VMSS resource errors for '%s' in resource group '%s': %s", + vmss_name, + resource_group, + exc, + ) + return [] + + errors: list[ProviderErrorEntry] = [] + provisioning_state = str(vmss.provisioning_state or "") + # getattr: Azure may surface instance-view statuses dynamically here. + statuses = getattr(vmss, "statuses", None) or [] + + errors.extend( + self._extract_vm_errors( + statuses, + instance_id=vmss_name, + vmss_name=vmss_name, + ) + ) + + if provisioning_state.lower() == "failed" and not errors: + errors.append( + { + "error_code": "ProvisioningStateFailed", + "error_message": f"VMSS '{vmss_name}' provisioning failed", + "instance_id": vmss_name, + "resource_id": vmss_name, + "status_code": "ProvisioningState/failed", + "status_level": "Error", + } + ) + + return errors + + async def _get_vmss_orchestration_mode_async( + self, + resource_group: str, + vmss_name: str, + ) -> AzureVMSSOrchestrationMode: + compute = await self.azure_client.get_async_compute_client() + vmss = await compute.virtual_machine_scale_sets.get( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + return self._coerce_vmss_orchestration_mode(vmss.orchestration_mode) + + async def _list_flexible_vmss_instances_async( + self, + resource_group: str, + vmss_name: str, + include_instance_view: bool = False, + ) -> list[AzureHandlerStatusResult]: + compute = await self.azure_client.get_async_compute_client() + try: + pager = compute.virtual_machines.list( + **self._flexible_vmss_list_kwargs( + resource_group=resource_group, + vmss_name=vmss_name, + include_instance_view=include_instance_view, + ) + ) + vms = [vm async for vm in pager] + except Exception as exc: + raise VMSSNotFoundError( + f"Could not list flexible VMs for VMSS '{vmss_name}': {exc}", + vmss_name=vmss_name, + ) from exc + + return [ + # cast: SDK 38 exposes hardware_profile/instance_view/vm_id via the + # __flattened_items __getattr__ shim, invisible to static checkers + # though documented as the SDK's public surface. + await self._normalise_vm_async( + cast(AzureVmRuntimeStatusProtocol, vm), vmss_name, resource_group + ) + for vm in vms + ] + + async def _normalise_vm_async( + self, vm: AzureVmRuntimeStatusProtocol, vmss_name: str, resource_group: str + ) -> AzureHandlerStatusResult: + """Async variant of ``_normalise_vm`` with async network enrichment.""" + vm_identity = _read_vm_identity(vm) + network_identity = await resolve_network_identity_or_empty_async( + logger=self._logger, + target_label=f"VMSS member '{vm_identity.instance_id}' in '{vmss_name}'", + resolver=lambda: self.azure_client.resolve_network_identity_from_vm_async(vm), + ) + return self._build_normalized_vm_status( + vm=vm, + vm_identity=vm_identity, + vmss_name=vmss_name, + resource_group=resource_group, + network_identity=network_identity, + ) + + @staticmethod + def _vmss_expand_arg(include_instance_view: bool) -> str | None: + """Return the SDK expand value used for instance-view enrichment.""" + return "instanceView" if include_instance_view else None + + def _flexible_vmss_list_kwargs( + self, + *, + resource_group: str, + vmss_name: str, + include_instance_view: bool, + ) -> dict[str, Any]: + """Build the VM list filter Azure expects for Flexible VMSS membership.""" + vmss_resource_id = ( + f"/subscriptions/{self.azure_client.subscription_id}" + f"/resourceGroups/{resource_group}" + f"/providers/Microsoft.Compute/virtualMachineScaleSets/{vmss_name}" + ) + list_kwargs: dict[str, Any] = { + "resource_group_name": resource_group, + "filter": f"'virtualMachineScaleSet/id' eq '{vmss_resource_id}'", + } + expand = self._vmss_expand_arg(include_instance_view) + if expand is not None: + list_kwargs["expand"] = expand + return list_kwargs + + def _build_normalized_vm_status( + self, + *, + vm: AzureVmRuntimeStatusProtocol, + vm_identity: _AzureVmIdentity, + vmss_name: str, + resource_group: str, + network_identity: AzureNetworkIdentity, + ) -> AzureHandlerStatusResult: + """Build the normalized VMSS member status once network identity is resolved.""" + status = "unknown" + instance_view = vm.instance_view + vm_statuses = instance_view_statuses(instance_view) + if vm_statuses is not None: + status = resolve_power_state(vm_statuses) + fleet_errors = self._extract_vm_errors( + vm_statuses, + instance_id=vm_identity.instance_id, + vmss_name=vmss_name, + ) + else: + fleet_errors = [] + + hw = vm.hardware_profile + instance_type = hw.vm_size if hw else None + if not instance_type: + instance_type = "unknown" + + location = vm.location + zones = vm.zones + availability_zone = zones[0] if zones else None + + launch_time = None + if vm_statuses is not None: + for status_entry in vm_statuses: + timestamp = status_entry.time + if timestamp is not None: + launch_time = str(timestamp) + break + + vnet_id = network_identity["vnet_id"] + provider_data: AzureStatusProviderData = { + "resource_id": vmss_name, + "cloud_host_id": vm_identity.vm_id or vm_identity.instance_id, + "vmss_name": vmss_name, + "resource_group": resource_group, + "vmss_instance_id": vm_identity.instance_id, + "vm_id": vm_identity.vm_id, + "availability_zone": availability_zone, + "nic_id": network_identity["nic_id"], + "nic_name": network_identity["nic_name"], + "vnet_id": vnet_id, + "fleet_errors": [dict(error) for error in fleet_errors], + } + if vm_identity.vm_name: + provider_data["vm_name"] = vm_identity.vm_name + if location: + provider_data["location"] = str(location) + return { + "instance_id": vm_identity.instance_id, + "name": vm_identity.vm_name or vm_identity.instance_id, + "status": status, + "private_ip": network_identity["private_ip"], + "public_ip": network_identity["public_ip"], + "launch_time": launch_time, + "instance_type": instance_type, + "subnet_id": network_identity["subnet_id"], + "vpc_id": vnet_id, + "tags": vm.tags or {}, + "price_type": None, + "provider_type": "azure", + "provider_data": provider_data, + } + + @staticmethod + def _coerce_vmss_orchestration_mode( + raw_mode: Optional[SdkOrchestrationMode], + ) -> AzureVMSSOrchestrationMode: + """Coerce Azure's orchestration-mode field into the provider enum. + + The SDK returns ``Optional[OrchestrationMode]`` (a plain ``Enum``; its + ``str()`` is the qualified ``"OrchestrationMode.FLEXIBLE"`` form rather + than the wire value, so we read ``.value`` directly). + """ + if raw_mode is None: + return AzureVMSSOrchestrationMode.FLEXIBLE + return AzureVMSSOrchestrationMode(raw_mode.value) + + def _vmss_delete_submission_result( + self, + *, + vmss_name: str, + exc: Exception | None = None, + ) -> _VmssCleanupSubmissionState: + """Build the normalized delete-submission state for follow-up cleanup.""" + if exc is None: + self._logger.info( + "Submitted immediate delete for VMSS '%s' because return should empty it", + vmss_name, + ) + return { + "delete_submitted": True, + "delete_retry_pending": False, + "last_delete_error": None, + } + + self._logger.info( + "Immediate delete submission for VMSS '%s' did not succeed; async cleanup will retry: %s", + vmss_name, + exc, + ) + return { + "delete_submitted": False, + "delete_retry_pending": True, + "last_delete_error": str(exc), + } + + def _extract_vm_errors( + self, + statuses: list[Any], + *, + instance_id: str, + vmss_name: str, + ) -> list[ProviderErrorEntry]: + """Extract provisioning failures from Azure instance view statuses.""" + errors: list[ProviderErrorEntry] = [] + + for status in statuses: + code = str(_status_attr(status, "code", "") or "") + level = str(_status_attr(status, "level", "") or "") + message = _status_attr(status, "message") + display_status = _status_attr(status, "display_status") + + is_failure = ( + code.lower().startswith("provisioningstate/failed") or level.lower() == "error" + ) + if not is_failure: + continue + + vm_error: ProviderErrorEntry = { + "error_code": "ProvisioningStateFailed", + "error_message": str( + message or display_status or code or "Azure provisioning failed" + ), + "instance_id": instance_id, + "resource_id": vmss_name, + "status_code": code, + "status_level": level or None, + } + errors.append(vm_error) + + return errors + + # ------------------------------------------------------------------ + # Example templates + # ------------------------------------------------------------------ + + @classmethod + def get_example_templates(cls) -> list[dict[str, Any]]: + """Return example VMSS template configurations.""" + return [ + { + "template_id": "azure-vmss-linux-basic", + "name": "Azure VMSS Linux Basic", + "description": "VMSS with Ubuntu 22.04 LTS on Standard_D4s_v5", + "provider_type": "azure", + "provider_api": AzureProviderApi.VMSS.value, + "vm_size": "Standard_D4s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "max_instances": 2, + }, + { + "template_id": "azure-vmss-spot", + "name": "Azure VMSS Spot Instances", + "description": "VMSS with Spot VMs for cost-effective workloads", + "provider_type": "azure", + "provider_api": AzureProviderApi.VMSS.value, + "vm_size": "Standard_D4s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "priority": "Spot", + "eviction_policy": "Deallocate", + "billing_profile_max_price": -1.0, + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "max_instances": 5, + }, + ] diff --git a/src/orb/providers/azure/infrastructure/sdk_shapes.py b/src/orb/providers/azure/infrastructure/sdk_shapes.py new file mode 100644 index 000000000..ec3c0f2ec --- /dev/null +++ b/src/orb/providers/azure/infrastructure/sdk_shapes.py @@ -0,0 +1,102 @@ +"""Shared structural typing helpers for Azure SDK object shapes. + +These protocols describe the small attribute subsets ORB consumes from the +Azure SDK models. Keeping them here centralizes the cast boundary for Azure's +generated model types instead of redefining local one-off protocols in each +handler/service. +""" + +from __future__ import annotations + +from typing import Any, Optional, Protocol, cast + + +class AzureStatusWithCodeProtocol(Protocol): + """Azure status-like object exposing a ``code`` attribute.""" + + @property + def code(self) -> str | None: + """Return the Azure status code.""" + ... + + +class AzureVmWithNameProtocol(Protocol): + """Azure VM-like object exposing a ``name`` attribute.""" + + @property + def name(self) -> str | None: + """Return the Azure VM resource name.""" + ... + + +class AzureVmWithIdentityProtocol(Protocol): + """Azure VM-like object exposing ``name`` and ``vm_id`` attributes.""" + + @property + def name(self) -> Optional[str]: + """Return the Azure VM resource name.""" + ... + + @property + def vm_id(self) -> Optional[str]: + """Return Azure's stable VM identifier.""" + ... + + +class AzureVmHardwareProfileProtocol(Protocol): + """Azure VM hardware profile surface used for status normalization.""" + + @property + def vm_size(self) -> Optional[str]: + """Return the Azure VM size.""" + ... + + +class AzureVmRuntimeStatusProtocol(AzureVmWithIdentityProtocol, Protocol): + """Azure VM surface shared by standalone VMs and VMSS member VMs. + + Azure's SDK exposes standalone ``VirtualMachine`` and VMSS + ``VirtualMachineScaleSetVM`` as separate generated classes. ORB only needs + this structural subset when normalizing status. + """ + + @property + def hardware_profile(self) -> Optional[AzureVmHardwareProfileProtocol]: + """Return the VM hardware profile when Azure included one.""" + ... + + @property + def instance_view(self) -> object | None: + """Return the VM instance view when Azure included one.""" + ... + + @property + def location(self) -> Optional[str]: + """Return the Azure location for the VM.""" + ... + + @property + def zones(self) -> Optional[list[str]]: + """Return the availability zones attached to the VM.""" + ... + + @property + def tags(self) -> Optional[dict[str, str]]: + """Return the user-supplied tag dict applied to the VM, if any.""" + ... + + +class AzureInstanceViewWithStatusesProtocol(Protocol): + """Azure instance-view-like object exposing ``statuses``.""" + + @property + def statuses(self) -> list[Any]: + """Return Azure instance-view status entries.""" + ... + + +def instance_view_statuses(instance_view: object | None) -> list[Any] | None: + """Return Azure instance-view statuses when the object shape supports them.""" + if instance_view is None: + return None + return cast(AzureInstanceViewWithStatusesProtocol, instance_view).statuses diff --git a/src/orb/providers/azure/infrastructure/services/__init__.py b/src/orb/providers/azure/infrastructure/services/__init__.py new file mode 100644 index 000000000..7a9f511d9 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/__init__.py @@ -0,0 +1 @@ +"""Azure infrastructure services package.""" diff --git a/src/orb/providers/azure/infrastructure/services/arm_payload_mapper.py b/src/orb/providers/azure/infrastructure/services/arm_payload_mapper.py new file mode 100644 index 000000000..41c707f86 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/arm_payload_mapper.py @@ -0,0 +1,344 @@ +"""ARM payload mapper. + +Converts ``AzureTemplate`` domain objects into Azure Resource Manager (ARM) +payloads for VMSS and SingleVM resources. This keeps infrastructure +serialisation concerns (API versions, ARM resource type strings, payload +structure) out of the domain aggregate. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import ( + AzureCachingType, + AzureOSDiskType, + AzurePriority, + AzureVMSSOrchestrationMode, +) + +# Flexible VMSS network profiles require a Microsoft.Network API version +# in the payload. +_VMSS_FLEX_NETWORK_API_VERSION = "2022-11-01" + + +class ArmPayloadMapper: + """Stateless mapper from ``AzureTemplate`` to ARM resource payloads.""" + + # ------------------------------------------------------------------ + # VMSS payload + # ------------------------------------------------------------------ + + @staticmethod + def vmss_payload(template: AzureTemplate) -> dict[str, Any]: + """Build the ARM resource payload for a VMSS create/update. + + Equivalent to the former ``AzureTemplate.to_azure_api_format()``. + The broad return type is necessary: this is arbitrary ARM JSON that can be + merged with user-provided native specs before it reaches the Azure SDK. + """ + properties: dict[str, Any] = { + "orchestrationMode": template.orchestration_mode.value, + "singlePlacementGroup": template.single_placement_group, + "upgradePolicy": {"mode": template.upgrade_policy_mode.value}, + } + + if template.orchestration_mode == AzureVMSSOrchestrationMode.UNIFORM: + properties["overprovision"] = template.overprovision + + platform_fault_domain_count = template.platform_fault_domain_count + if template.orchestration_mode == AzureVMSSOrchestrationMode.FLEXIBLE: + platform_fault_domain_count = platform_fault_domain_count or 1 + if platform_fault_domain_count is not None: + properties["platformFaultDomainCount"] = platform_fault_domain_count + + # --- Virtual machine profile --- + vm_profile: dict[str, Any] = {"hardwareProfile": {}} + + storage_profile = _build_storage_profile(template) + vm_profile["storageProfile"] = storage_profile + + # Network profile + if template.network_config: + nic_config = template.network_config.to_arm_dict() + vm_profile["networkProfile"] = { + "networkInterfaceConfigurations": [nic_config], + "networkApiVersion": _VMSS_FLEX_NETWORK_API_VERSION, + } + + # OS profile + os_profile: dict[str, Any] = { + "computerNamePrefix": (template.vmss_name or template.template_id)[:9], + "adminUsername": template.admin_username, + } + if template.custom_data: + os_profile["customData"] = template.custom_data + if template.ssh_public_keys: + os_profile["linuxConfiguration"] = _build_linux_ssh_config(template) + vm_profile["osProfile"] = os_profile + + _apply_priority(template, vm_profile) + _apply_security_profile(template, vm_profile) + + # Extension profile + if template.extension_profile: + vm_profile["extensionProfile"] = { + "extensions": template.extension_profile, + } + + properties["virtualMachineProfile"] = vm_profile + + # --- Spot restore policy --- + if template.spot_restore_enabled: + spot_restore: dict[str, Any] = {"enabled": True} + if template.spot_restore_timeout: + spot_restore["restoreTimeout"] = template.spot_restore_timeout + properties["spotRestorePolicy"] = spot_restore + + if template.spot_percentage is not None: + properties["priorityMixPolicy"] = { + "baseRegularPriorityCount": template.base_regular_priority_count, + "regularPriorityPercentageAboveBase": 100 - template.spot_percentage, + } + + # --- Identity --- + identity = _build_identity(template) + + # --- VMSS SDK body (VirtualMachineScaleSet) --- + resource: dict[str, Any] = { + "location": template.location.value, + "sku": { + "name": "Mix" if template.uses_vm_size_mix else template.vm_size, + "capacity": template.max_instances, + }, + "properties": properties, + "tags": template.tags if template.tags else {}, + } + + if template.uses_vm_size_mix: + sku_profile: dict[str, Any] = {"vmSizes": template.build_vm_size_profile()} + if template.vmss_allocation_strategy: + sku_profile["allocationStrategy"] = template.vmss_allocation_strategy.to_arm_value() + properties["skuProfile"] = sku_profile + + if template.zones: + resource["zones"] = template.zones + if identity: + resource["identity"] = identity + if template.proximity_placement_group_id: + properties["proximityPlacementGroup"] = { + "id": template.proximity_placement_group_id.value, + } + if template.capacity_reservation_group_id: + vm_profile["capacityReservation"] = { + "capacityReservationGroup": { + "id": template.capacity_reservation_group_id.value, + }, + } + if template.node_attributes: + properties.update(template.node_attributes) + + return resource + + # ------------------------------------------------------------------ + # SingleVM payload + # ------------------------------------------------------------------ + + @staticmethod + def single_vm_payload( + template: AzureTemplate, + vm_name: str, + nic_id: str, + *, + vm_size_override: Optional[str] = None, + ) -> dict[str, Any]: + """Build ARM VM create parameters for a single VM deployment. + + Equivalent to the former ``SingleVMHandler._build_vm_params()``. + """ + params: dict[str, Any] = { + "location": template.location.value, + "properties": { + "hardwareProfile": {"vmSize": vm_size_override or template.vm_size}, + "storageProfile": _build_storage_profile( + template, + default_os_disk_storage_account_type=AzureOSDiskType.STANDARD_LRS.value, + include_default_os_disk_caching=False, + ), + "osProfile": { + "computerName": vm_name[:15], # Azure host name limit + "adminUsername": template.admin_username, + }, + "networkProfile": { + "networkInterfaces": [ + { + "id": nic_id, + "properties": { + "primary": True, + "deleteOption": "Delete", + }, + } + ] + }, + }, + "tags": template.tags or {}, + } + + # SSH keys + if template.ssh_public_keys: + params["properties"]["osProfile"]["linuxConfiguration"] = _build_linux_ssh_config( + template + ) + + # Custom data + if template.custom_data: + params["properties"]["osProfile"]["customData"] = template.custom_data + + _apply_priority(template, params["properties"]) + + _apply_security_profile(template, params["properties"]) + + # Capacity reservation + if template.capacity_reservation_group_id: + params["properties"]["capacityReservation"] = { + "capacityReservationGroup": { + "id": template.capacity_reservation_group_id.value, + } + } + + # Identity + identity = _build_identity(template) + if identity: + params["identity"] = identity + + # Availability zones + if template.zones: + params["zones"] = template.zones + + return params + + +# ------------------------------------------------------------------ +# Shared helpers +# ------------------------------------------------------------------ + + +def _build_storage_profile( + template: AzureTemplate, + *, + default_os_disk_storage_account_type: str = AzureOSDiskType.PREMIUM_LRS.value, + include_default_os_disk_caching: bool = True, +) -> dict[str, Any]: + """Build the storageProfile section common to Azure compute payloads.""" + storage_profile: dict[str, Any] = {} + + if template.image: + storage_profile["imageReference"] = template.image.to_arm_dict() + elif template.image_id: + storage_profile["imageReference"] = {"id": template.image_id} + + if template.os_disk: + storage_profile["osDisk"] = template.os_disk.to_arm_dict() + else: + default_os_disk: dict[str, Any] = { + "createOption": "FromImage", + "deleteOption": "Delete", + "managedDisk": { + "storageAccountType": default_os_disk_storage_account_type, + }, + } + if include_default_os_disk_caching: + default_os_disk["caching"] = AzureCachingType.READ_WRITE.value + storage_profile["osDisk"] = default_os_disk + + if template.data_disks: + storage_profile["dataDisks"] = [d.to_arm_dict() for d in template.data_disks] + + _apply_disk_encryption(template, storage_profile) + + return storage_profile + + +def _apply_disk_encryption(template: AzureTemplate, storage_profile: dict[str, Any]) -> None: + """Apply disk encryption set to OS and data disks if configured.""" + if not template.disk_encryption_set_id: + return + + os_disk_managed = storage_profile.get("osDisk", {}).get("managedDisk") + if isinstance(os_disk_managed, dict): + os_disk_managed["diskEncryptionSet"] = { + "id": template.disk_encryption_set_id.value, + } + + for data_disk in storage_profile.get("dataDisks", []): + managed_disk = data_disk.get("managedDisk") + if isinstance(managed_disk, dict): + managed_disk["diskEncryptionSet"] = { + "id": template.disk_encryption_set_id.value, + } + + +def _build_linux_ssh_config(template: AzureTemplate) -> dict[str, Any]: + """Build the linuxConfiguration SSH block.""" + return { + "disablePasswordAuthentication": True, + "ssh": { + "publicKeys": [ + { + "path": f"/home/{template.admin_username}/.ssh/authorized_keys", + "keyData": key, + } + for key in template.ssh_public_keys + ], + }, + } + + +def _apply_priority(template: AzureTemplate, target: dict[str, Any]) -> None: + """Apply priority / spot / billing profile to a VM profile dict.""" + if template.priority == AzurePriority.REGULAR: + return + target["priority"] = template.priority.value + if template.eviction_policy: + target["evictionPolicy"] = template.eviction_policy.value + if template.billing_profile_max_price is not None: + target["billingProfile"] = { + "maxPrice": template.billing_profile_max_price, + } + + +def _apply_security_profile(template: AzureTemplate, target: dict[str, Any]) -> None: + """Apply security profile (TrustedLaunch, Confidential, encryption at host).""" + if not template.security_type: + return + security_profile: dict[str, Any] = { + "securityType": template.security_type.value, + } + uefi: dict[str, Any] = {} + if template.secure_boot_enabled is not None: + uefi["secureBootEnabled"] = template.secure_boot_enabled + if template.vtpm_enabled is not None: + uefi["vTpmEnabled"] = template.vtpm_enabled + if uefi: + security_profile["uefiSettings"] = uefi + if template.encryption_at_host is not None: + security_profile["encryptionAtHost"] = template.encryption_at_host + target["securityProfile"] = security_profile + + +def _build_identity(template: AzureTemplate) -> dict[str, Any] | None: + """Build the identity block for the ARM resource.""" + if template.system_assigned_identity and template.user_assigned_identity_ids: + return { + "type": "SystemAssigned, UserAssigned", + "userAssignedIdentities": {uid: {} for uid in template.user_assigned_identity_ids}, + } + if template.system_assigned_identity: + return {"type": "SystemAssigned"} + if template.user_assigned_identity_ids: + return { + "type": "UserAssigned", + "userAssignedIdentities": {uid: {} for uid in template.user_assigned_identity_ids}, + } + return None diff --git a/src/orb/providers/azure/infrastructure/services/arm_resource_id_parser.py b/src/orb/providers/azure/infrastructure/services/arm_resource_id_parser.py new file mode 100644 index 000000000..eba476d89 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/arm_resource_id_parser.py @@ -0,0 +1,101 @@ +"""Canonical ARM resource ID parsing helpers for Azure infrastructure.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +@dataclass(frozen=True) +class ParsedArmResourceId: + """Validated ARM resource identifier components.""" + + subscription_id: str + resource_group: str + provider_namespace: str + resource_path_segments: tuple[str, ...] + + @property + def resource_name(self) -> str: + """Return the leaf resource name from the ARM resource path.""" + return self.resource_path_segments[-1] + + @property + def resource_type(self) -> str: + """Return the ARM resource type segment (e.g. 'virtualMachines').""" + return self.resource_path_segments[-2] + + def parent_resource_id(self) -> Optional[str]: + """Return the parent ARM resource ID for child resources.""" + if len(self.resource_path_segments) <= 2: + return None + + return ( + f"/subscriptions/{self.subscription_id}" + f"/resourceGroups/{self.resource_group}" + f"/providers/{self.provider_namespace}/" + "/".join(self.resource_path_segments[:-2]) + ) + + +class ArmResourceIdParser: + """Parse and project canonical Azure ARM resource identifiers.""" + + @staticmethod + def parse(arm_id: str) -> Optional[ParsedArmResourceId]: + """Parse an ARM resource ID only when it matches the canonical shape.""" + raw_arm_id = str(arm_id).strip() + if not raw_arm_id: + return None + + stripped = raw_arm_id.strip("/") + if not stripped: + return None + + parts = stripped.split("/") + if any(not segment for segment in parts): + return None + + if len(parts) < 8: + return None + + if parts[0].lower() != "subscriptions" or parts[2].lower() != "resourcegroups": + return None + + if parts[4].lower() != "providers": + return None + + subscription_id = parts[1] + resource_group = parts[3] + provider_namespace = parts[5] + resource_path_segments = tuple(parts[6:]) + + if not subscription_id or not resource_group or not provider_namespace: + return None + + if len(resource_path_segments) < 2 or len(resource_path_segments) % 2 != 0: + return None + + return ParsedArmResourceId( + subscription_id=subscription_id, + resource_group=resource_group, + provider_namespace=provider_namespace, + resource_path_segments=resource_path_segments, + ) + + def extract_resource_group_and_name(self, arm_id: str) -> Optional[tuple[str, str]]: + """Extract ``(resource_group, resource_name)`` from an ARM resource ID.""" + parsed_arm_id = self.parse(arm_id) + if parsed_arm_id is None: + return None + return parsed_arm_id.resource_group, parsed_arm_id.resource_name + + def subnet_id_to_vnet_id(self, subnet_id: Optional[str]) -> Optional[str]: + """Return the parent VNet ARM ID from a subnet ARM ID.""" + if not subnet_id: + return None + parsed_arm_id = self.parse(subnet_id) + if parsed_arm_id is None: + return None + if parsed_arm_id.resource_type.lower() != "subnets": + return None + return parsed_arm_id.parent_resource_id() diff --git a/src/orb/providers/azure/infrastructure/services/azure_deployment_service.py b/src/orb/providers/azure/infrastructure/services/azure_deployment_service.py new file mode 100644 index 000000000..ebba228bf --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/azure_deployment_service.py @@ -0,0 +1,222 @@ +"""Azure ARM deployment helpers.""" + +from __future__ import annotations + +import hashlib +import re +from typing import Any, Optional + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.infrastructure.azure_client import AzureClient + + +class AzureDeploymentService: + """Build and submit Azure ARM deployments for provider-owned resources.""" + + _TEMPLATE_SCHEMA = ( + "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#" + ) + _DEPLOYMENT_API_VERSION = "2025-04-01" + _DEPLOYMENT_PROVIDER_NAMESPACE = "Microsoft.Resources" + _DEPLOYMENT_RESOURCE_TYPE = "deployments" + _DEPLOYMENT_MODE = "Incremental" + _VM_API_VERSION = "2023-09-01" + _NETWORK_API_VERSION = "2023-09-01" + + def __init__(self, azure_client: AzureClient, logger: LoggingPort) -> None: + """Initialize with an Azure client and logger.""" + self.azure_client = azure_client + self._logger = logger + + @classmethod + def build_deployment_name(cls, *parts: str) -> str: + """Build a deterministic Azure deployment name within ARM limits.""" + normalized = [ + re.sub(r"[^A-Za-z0-9._()-]+", "-", str(part)).strip("-") + for part in parts + if str(part).strip() + ] + name = "-".join(part for part in normalized if part) or "orb-deployment" + if len(name) <= 64: + return name + + digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:8] + prefix = name[:55].rstrip("-") + return f"{prefix}-{digest}" + + @staticmethod + def resource_id_expression(resource_type: str, resource_name: str) -> str: + """Return an ARM expression for a sibling resource ID.""" + return f"[resourceId('{resource_type}', '{resource_name}')]" + + async def submit_template_deployment_async( + self, + *, + resource_group: str, + deployment_name: str, + template: dict[str, Any], + parameters: Optional[dict[str, Any]] = None, + ) -> str: + """Submit an ARM deployment without waiting for completion via the async SDK.""" + deployment_payload = { + "properties": { + "mode": self._DEPLOYMENT_MODE, + "template": template, + "parameters": parameters or {}, + } + } + resource_client = await self.azure_client.get_async_resource_client() + resource_operations: Any = resource_client.resources + await resource_operations.begin_create_or_update( + resource_group_name=resource_group, + resource_provider_namespace=self._DEPLOYMENT_PROVIDER_NAMESPACE, + parent_resource_path="", + resource_type=self._DEPLOYMENT_RESOURCE_TYPE, + resource_name=deployment_name, + api_version=self._DEPLOYMENT_API_VERSION, + parameters=deployment_payload, + ) + + self._logger.info( + "Submitted ARM deployment '%s' in resource group '%s'", + deployment_name, + resource_group, + ) + return deployment_name + + async def get_deployment_status_async( + self, + *, + resource_group: str, + deployment_name: str, + ) -> Optional[dict[str, Any]]: + """Return provisioning metadata for a submitted ARM deployment via the async SDK.""" + resource_client = await self.azure_client.get_async_resource_client() + deployment = await resource_client.resources.get( + resource_group_name=resource_group, + resource_provider_namespace=self._DEPLOYMENT_PROVIDER_NAMESPACE, + parent_resource_path="", + resource_type=self._DEPLOYMENT_RESOURCE_TYPE, + resource_name=deployment_name, + api_version=self._DEPLOYMENT_API_VERSION, + ) + properties = deployment.properties + if not isinstance(properties, dict): + properties = {} + + error = properties.get("error") + error_code = None + error_message = None + if isinstance(error, dict): + error_code = error.get("code") + error_message = error.get("message") + + return { + "provisioning_state": properties.get("provisioningState"), + "error_code": error_code, + "error_message": error_message, + } + + def build_single_vm_deployment_template( + self, + *, + location: str, + subnet_id: str, + vm_definitions: list[dict[str, Any]], + enable_accelerated_networking: bool = False, + nsg_id: Optional[str] = None, + load_balancer_backend_pool_ids: Optional[list[str]] = None, + load_balancer_inbound_nat_pool_ids: Optional[list[str]] = None, + application_gateway_backend_pool_ids: Optional[list[str]] = None, + ) -> dict[str, Any]: + """Build a deployment template containing Public IP, NIC, and VM resources.""" + resources: list[dict[str, Any]] = [] + for vm_definition in vm_definitions: + vm_name = vm_definition["vm_name"] + nic_name = vm_definition["nic_name"] + vm_payload = vm_definition["vm_payload"] + public_ip_name = vm_definition.get("public_ip_name") + + nic_depends_on: list[str] = [] + ip_config_properties: dict[str, Any] = { + "subnet": {"id": subnet_id}, + "privateIPAllocationMethod": "Dynamic", + } + + if public_ip_name: + public_ip_id = self.resource_id_expression( + "Microsoft.Network/publicIPAddresses", + public_ip_name, + ) + resources.append( + { + "type": "Microsoft.Network/publicIPAddresses", + "apiVersion": self._NETWORK_API_VERSION, + "name": public_ip_name, + "location": location, + "sku": {"name": "Standard"}, + "properties": { + "publicIPAllocationMethod": "Static", + "deleteOption": "Delete", + }, + } + ) + nic_depends_on.append(public_ip_id) + ip_config_properties["publicIPAddress"] = { + "id": public_ip_id, + "deleteOption": "Delete", + } + + if load_balancer_backend_pool_ids: + ip_config_properties["loadBalancerBackendAddressPools"] = [ + {"id": pool_id} for pool_id in load_balancer_backend_pool_ids + ] + if load_balancer_inbound_nat_pool_ids: + ip_config_properties["loadBalancerInboundNatPools"] = [ + {"id": pool_id} for pool_id in load_balancer_inbound_nat_pool_ids + ] + if application_gateway_backend_pool_ids: + ip_config_properties["applicationGatewayBackendAddressPools"] = [ + {"id": pool_id} for pool_id in application_gateway_backend_pool_ids + ] + + nic_resource: dict[str, Any] = { + "type": "Microsoft.Network/networkInterfaces", + "apiVersion": self._NETWORK_API_VERSION, + "name": nic_name, + "location": location, + "properties": { + "ipConfigurations": [ + { + "name": "ipconfig1", + "properties": ip_config_properties, + } + ], + "enableAcceleratedNetworking": enable_accelerated_networking, + }, + } + if nsg_id: + nic_resource["properties"]["networkSecurityGroup"] = {"id": nsg_id} + if nic_depends_on: + nic_resource["dependsOn"] = nic_depends_on + resources.append(nic_resource) + + vm_resource_payload = { + key: value + for key, value in vm_payload.items() + if key not in {"type", "apiVersion", "name", "dependsOn"} + } + vm_resource = dict(vm_resource_payload) + vm_resource["type"] = "Microsoft.Compute/virtualMachines" + vm_resource["apiVersion"] = self._VM_API_VERSION + vm_resource["name"] = vm_name + vm_resource["dependsOn"] = [ + self.resource_id_expression("Microsoft.Network/networkInterfaces", nic_name) + ] + resources.append(vm_resource) + + return { + "$schema": self._TEMPLATE_SCHEMA, + "contentVersion": "1.0.0.0", + "resources": resources, + } diff --git a/src/orb/providers/azure/infrastructure/services/azure_native_spec_service.py b/src/orb/providers/azure/infrastructure/services/azure_native_spec_service.py new file mode 100644 index 000000000..962a40687 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/azure_native_spec_service.py @@ -0,0 +1,113 @@ +"""Azure-specific native spec processing.""" + +from __future__ import annotations + +from typing import Any, Optional + +from pydantic import BaseModel, ConfigDict, Field + +from orb.application.services.native_spec_service import NativeSpecService +from orb.domain.base.ports.configuration_port import ConfigurationPort +from orb.domain.request.aggregate import Request +from orb.infrastructure.di.injectable import injectable +from orb.infrastructure.utilities.common.deep_merge import deep_merge +from orb.infrastructure.utilities.file.json_utils import read_json_file +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate + + +class _AzureNativeSpecConfig(BaseModel): + """Azure native spec extension config owned by the Azure native spec subsystem.""" + + spec_file_base_path: str = Field(default="specs/azure") + + +def _default_azure_native_spec_config() -> _AzureNativeSpecConfig: + """Create the default Azure native spec extension config.""" + return _AzureNativeSpecConfig() + + +class _AzureProviderExtensionsConfig(BaseModel): + """Azure provider extensions consumed by the Azure native spec subsystem.""" + + model_config = ConfigDict(extra="ignore") + + native_spec: _AzureNativeSpecConfig = Field(default_factory=_default_azure_native_spec_config) + + +@injectable +class AzureNativeSpecService: + """Azure-specific native spec processing for VMSS and SingleVM payloads.""" + + def __init__(self, native_spec_service: NativeSpecService, config_port: ConfigurationPort): + self.native_spec_service = native_spec_service + self.config_port = config_port + + def process_provider_api_spec_with_merge( + self, + template: AzureTemplate, + request: Request, + default_payload: dict[str, Any], + extra_context: Optional[dict[str, Any]] = None, + ) -> Optional[dict[str, Any]]: + """Render and optionally merge an Azure provider API spec.""" + if not self.native_spec_service.is_native_spec_enabled(): + return None + + native_spec = self._resolve_provider_api_spec(template) + if not native_spec: + return None + + context = self._build_azure_context(template, request) + if extra_context: + context.update(extra_context) + + rendered_native_spec = self.native_spec_service.render_spec(native_spec, context) + + native_config = self.config_port.get_native_spec_config() or {} + merge_mode = native_config.get("merge_mode", "merge") + + if merge_mode == "replace": + return rendered_native_spec + if merge_mode == "merge": + return deep_merge(default_payload, rendered_native_spec) + return rendered_native_spec + + def _resolve_provider_api_spec(self, template: AzureTemplate) -> Optional[dict[str, Any]]: + """Resolve provider API spec from inline data or file path.""" + if template.provider_api_spec: + return template.provider_api_spec + if template.provider_api_spec_file: + return self._load_spec_file(template.provider_api_spec_file) + return None + + def _load_spec_file(self, file_path: str) -> dict[str, Any]: + """Load Azure native spec file.""" + provider_config = self.config_port.get_provider_config() + azure_defaults = None + if provider_config: + azure_defaults = provider_config.provider_defaults.get("azure") + + base_path = _AzureNativeSpecConfig.model_fields["spec_file_base_path"].default + if azure_defaults and azure_defaults.extensions: + azure_extensions = _AzureProviderExtensionsConfig.model_validate( + azure_defaults.extensions + ) + base_path = azure_extensions.native_spec.spec_file_base_path + + return read_json_file(f"{base_path}/{file_path}") + + def _build_azure_context(self, template: AzureTemplate, request: Request) -> dict[str, Any]: + """Build Azure-specific native spec rendering context.""" + package_info = self.config_port.get_package_info() or {} + + return { + "request_id": str(request.request_id), + "requested_count": request.requested_count, + "template_id": template.template_id, + "provider_api": template.provider_api.value, + "resource_group": template.resource_group.value, + "location": template.location.value, + "vm_size": template.vm_size, + "package_name": package_info.get("name", "open-resource-broker"), + "package_version": package_info.get("version", "unknown"), + } diff --git a/src/orb/providers/azure/infrastructure/services/azure_network_identity_resolver.py b/src/orb/providers/azure/infrastructure/services/azure_network_identity_resolver.py new file mode 100644 index 000000000..6c6724f70 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/azure_network_identity_resolver.py @@ -0,0 +1,321 @@ +"""Azure VM/NIC network identity enrichment helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Optional, Protocol, TypedDict, cast + +from orb.domain.base.ports import LoggingPort + +from .arm_resource_id_parser import ArmResourceIdParser + + +class AzureResourceRefProtocol(Protocol): + """Minimal ARM resource reference carrying an Azure resource ID.""" + + id: Optional[str] + + +class AzureNicReferencePropertiesProtocol(Protocol): + """Subset of NIC reference properties used for primary-NIC ordering.""" + + primary: Optional[bool] + + +class AzureNicReferenceProtocol(AzureResourceRefProtocol, Protocol): + """NIC reference shape exposed from a VM network profile.""" + + properties: Optional[AzureNicReferencePropertiesProtocol] + + +class AzureNetworkProfileProtocol(Protocol): + """VM network profile surface needed for NIC reference enumeration.""" + + network_interfaces: list[AzureNicReferenceProtocol] + + +class AzureIpConfigurationPropertiesProtocol(Protocol): + """Fallback property bag exposed by Azure IP configuration objects.""" + + private_ip_address: Optional[str] + subnet: Optional[AzureResourceRefProtocol] + public_ip_address: Optional[AzureResourceRefProtocol] + + +class AzureIpConfigurationProtocol(Protocol): + """IP configuration fields used to resolve private/public network identity.""" + + private_ip_address: Optional[str] + subnet: Optional[AzureResourceRefProtocol] + public_ip_address: Optional[AzureResourceRefProtocol] + properties: Optional[AzureIpConfigurationPropertiesProtocol] + + +class AzureNicProtocol(Protocol): + """NIC surface used to enumerate IP configurations.""" + + ip_configurations: list[AzureIpConfigurationProtocol] + + +class AzurePublicIpProtocol(Protocol): + """Public IP resource surface used to read the resolved IP address.""" + + ip_address: Optional[str] + + +class AzureVmNetworkIdentityProtocol(Protocol): + """VM surface used to enter the network-identity resolution flow.""" + + network_profile: Optional[AzureNetworkProfileProtocol] + + +class AzureNetworkIdentity(TypedDict): + """Resolved network identity fields used by Azure status normalization.""" + + private_ip: str | None + public_ip: str | None + subnet_id: str | None + vnet_id: str | None + nic_id: str | None + nic_name: str | None + + +@dataclass(frozen=True) +class _AzureResourceLookup: + """Parsed ARM lookup fields needed for Azure SDK get calls.""" + + resource_id: str + resource_group: str + name: str + + +class AzureNetworkIdentityResolver: + """Resolve private/public IP and subnet/VNet identity from Azure NIC references.""" + + def __init__( + self, + *, + async_network_client_getter: Callable[[], Awaitable[Any]], + logger: LoggingPort, + arm_resource_id_parser: ArmResourceIdParser, + network_lookup_error_types: Callable[[], tuple[type[BaseException], ...]], + ) -> None: + self._async_network_client_getter = async_network_client_getter + self._logger = logger + self._arm_resource_id_parser = arm_resource_id_parser + self._network_lookup_error_types = network_lookup_error_types + + @staticmethod + def empty_network_identity() -> AzureNetworkIdentity: + """Return the empty network identity shape used when enrichment fails.""" + return { + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + + @staticmethod + def _network_interface_refs_from_profile( + network_profile: Optional[AzureNetworkProfileProtocol], + ) -> list[AzureNicReferenceProtocol]: + if network_profile is None: + return [] + return list(network_profile.network_interfaces or []) + + @staticmethod + def _is_primary_nic_ref(nic_ref: AzureNicReferenceProtocol) -> bool: + nic_properties = nic_ref.properties + if nic_properties is None: + return False + return bool(nic_properties.primary) + + @staticmethod + def _subnet_from_ip_config( + ip_config: AzureIpConfigurationProtocol, + ) -> Optional[AzureResourceRefProtocol]: + if ip_config.subnet is not None: + return ip_config.subnet + ip_properties = ip_config.properties + if ip_properties is None: + return None + return ip_properties.subnet + + @staticmethod + def _public_ip_ref_from_ip_config( + ip_config: AzureIpConfigurationProtocol, + ) -> Optional[AzureResourceRefProtocol]: + if ip_config.public_ip_address is not None: + return ip_config.public_ip_address + ip_properties = ip_config.properties + if ip_properties is None: + return None + return ip_properties.public_ip_address + + @staticmethod + def _public_ip_address_from_resource(public_ip_resource: Any) -> Optional[str]: + return cast(AzurePublicIpProtocol, public_ip_resource).ip_address + + def _ordered_nic_refs( + self, + nic_refs: list[AzureNicReferenceProtocol], + ) -> list[AzureNicReferenceProtocol]: + """Return NIC refs ordered with the primary NIC first.""" + return sorted( + nic_refs, + key=lambda ref: not self._is_primary_nic_ref(ref), + ) + + def _nic_lookup_from_ref( + self, + nic_ref: AzureNicReferenceProtocol, + ) -> _AzureResourceLookup | None: + """Extract the NIC ARM ID, resource group, and resource name from a ref.""" + nic_id = nic_ref.id + if not nic_id: + return None + + nic_lookup = self._arm_resource_id_parser.extract_resource_group_and_name(str(nic_id)) + if not nic_lookup: + return None + + nic_rg, nic_name = nic_lookup + return _AzureResourceLookup( + resource_id=str(nic_id), + resource_group=nic_rg, + name=nic_name, + ) + + @staticmethod + def _private_ip_from_ip_config(ip_config: AzureIpConfigurationProtocol) -> Optional[str]: + """Resolve the private IP from direct fields or nested property bags.""" + private_ip = ip_config.private_ip_address + if not private_ip and ip_config.properties is not None: + private_ip = ip_config.properties.private_ip_address + return private_ip + + def _build_network_identity( + self, + *, + nic_id: str, + nic_name: str, + ip_config: AzureIpConfigurationProtocol, + public_ip: Optional[str], + ) -> AzureNetworkIdentity: + """Build the resolved network identity payload for one NIC IP configuration.""" + subnet = self._subnet_from_ip_config(ip_config) + subnet_id = subnet.id if subnet is not None else None + return { + "private_ip": self._private_ip_from_ip_config(ip_config), + "public_ip": public_ip, + "subnet_id": subnet_id, + "vnet_id": self._arm_resource_id_parser.subnet_id_to_vnet_id(subnet_id), + "nic_id": nic_id, + "nic_name": nic_name, + } + + def _public_ip_lookup_from_ip_config( + self, + ip_config: AzureIpConfigurationProtocol, + ) -> _AzureResourceLookup | None: + """Extract the public IP ARM ID, resource group, and resource name from an IP config.""" + public_ip_ref = self._public_ip_ref_from_ip_config(ip_config) + public_ip_id = public_ip_ref.id if public_ip_ref is not None else None + if not public_ip_id: + return None + + public_ip_lookup = self._arm_resource_id_parser.extract_resource_group_and_name( + str(public_ip_id) + ) + if not public_ip_lookup: + return None + + pip_rg, pip_name = public_ip_lookup + return _AzureResourceLookup( + resource_id=str(public_ip_id), + resource_group=pip_rg, + name=pip_name, + ) + + async def _resolve_public_ip_with_client_async( + self, + *, + network_client: Any, + ip_config: AzureIpConfigurationProtocol, + ) -> Optional[str]: + """Resolve a public IP address using the async Azure network client.""" + public_ip_lookup = self._public_ip_lookup_from_ip_config(ip_config) + if public_ip_lookup is None: + return None + + try: + pip = await network_client.public_ip_addresses.get( + resource_group_name=public_ip_lookup.resource_group, + public_ip_address_name=public_ip_lookup.name, + ) + return self._public_ip_address_from_resource(pip) + except self._network_lookup_error_types() as exc: + self._logger.debug( + "Failed to resolve public IP %s: %s", + public_ip_lookup.resource_id, + exc, + ) + return None + + async def _get_async_network_client(self) -> Any: + return await self._async_network_client_getter() + + async def resolve_from_vm_async(self, vm: Any) -> AzureNetworkIdentity: + """Async variant of ``resolve_from_vm`` using the async Network SDK.""" + net_profile = cast(AzureVmNetworkIdentityProtocol, vm).network_profile + nic_refs = self._network_interface_refs_from_profile(net_profile) + return await self.resolve_from_nic_refs_async(nic_refs) + + async def resolve_from_nic_refs_async( + self, + nic_refs: list[AzureNicReferenceProtocol], + ) -> AzureNetworkIdentity: + """Async variant of ``resolve_from_nic_refs`` using the async Network SDK.""" + network_identity = self.empty_network_identity() + if not nic_refs: + return network_identity + + ordered_refs = self._ordered_nic_refs(nic_refs) + network_client = await self._get_async_network_client() + + for nic_ref in ordered_refs: + nic_lookup = self._nic_lookup_from_ref(nic_ref) + if nic_lookup is None: + continue + + try: + nic = await network_client.network_interfaces.get( + resource_group_name=nic_lookup.resource_group, + network_interface_name=nic_lookup.name, + ) + except self._network_lookup_error_types() as exc: + self._logger.debug( + "Failed to resolve NIC %s: %s", + nic_lookup.resource_id, + exc, + ) + continue + + ip_configs = list(cast(AzureNicProtocol, nic).ip_configurations or []) + for ip_cfg in ip_configs: + network_identity.update( + self._build_network_identity( + nic_id=nic_lookup.resource_id, + nic_name=nic_lookup.name, + ip_config=ip_cfg, + public_ip=await self._resolve_public_ip_with_client_async( + network_client=network_client, + ip_config=ip_cfg, + ), + ) + ) + return network_identity + + return network_identity diff --git a/src/orb/providers/azure/infrastructure/services/spot_placement_score_adapter.py b/src/orb/providers/azure/infrastructure/services/spot_placement_score_adapter.py new file mode 100644 index 000000000..062151b34 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/spot_placement_score_adapter.py @@ -0,0 +1,193 @@ +"""Azure spot placement score adapter.""" + +from __future__ import annotations + +import asyncio +from typing import Any, Protocol + +import httpx + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.domain.template.value_objects import AzureLocationName +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.services.spot_placement_planner import ( + PlacementCandidate, + PlacementScore, + SpotPlacementScoreAdapter, +) + + +class AzureSpotPlacementTemplate(Protocol): + """Structural template interface required for Azure spot placement scoring.""" + + vm_size: str + location: AzureLocationName + placement_regions: list[str] + placement_zones: list[str] + zones: list[str] + candidate_vm_sizes: list[str] + + +class AzureSpotPlacementScoreAdapter(SpotPlacementScoreAdapter): + """Azure candidate scoring using the Spot Placement Scores REST API.""" + + _API_VERSION = "2025-02-01-preview" + _SCORE_MAP = { + "low": 0.2, + "medium": 0.6, + "high": 1.0, + "datanotfoundorstale": 0.0, + } + + def __init__( + self, + azure_client: AzureClient, + logger: LoggingPort, + subscription_id: str | None, + base_location: str, + ) -> None: + self._azure_client = azure_client + self._logger = logger + self._subscription_id = subscription_id + self._base_location = base_location + + def score_candidates( + self, requested_count: int, template: AzureSpotPlacementTemplate + ) -> list[PlacementScore]: + """Fetch spot placement scores from sync code only.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run( + self.score_candidates_async( + requested_count=requested_count, + template=template, + ) + ) + raise TypeError( + "score_candidates cannot run inside an active event loop; " + "use score_candidates_async instead" + ) + + async def score_candidates_async( + self, requested_count: int, template: AzureSpotPlacementTemplate + ) -> list[PlacementScore]: + """Fetch and return spot placement scores using async HTTP transport.""" + vm_sizes, regions, zones, candidates = self._candidate_inputs(template) + + if not candidates: + return [] + + raw_scores = await self._fetch_scores_async( + requested_count=requested_count, + regions=regions, + vm_sizes=vm_sizes, + zones=zones, + ) + + return self._build_scores(candidates=candidates, raw_scores=raw_scores) + + def _candidate_inputs( + self, + template: AzureSpotPlacementTemplate, + ) -> tuple[list[str], list[str], list[str], list[PlacementCandidate]]: + """Build spot-placement API inputs and matching planner candidates.""" + vm_sizes = template.candidate_vm_sizes + regions = template.placement_regions or [template.location.value or self._base_location] + zones = template.placement_zones or template.zones or [] + candidates = [ + PlacementCandidate( + candidate_id=f"azure:{region}:{zone or 'regional'}:{vm_size}", + instance_type=vm_size, + region=region, + zone=zone, + ) + for region in regions + for vm_size in vm_sizes + for zone in (zones or [None]) + ] + return vm_sizes, regions, zones, candidates + + async def _fetch_scores_async( + self, + requested_count: int, + regions: list[str], + vm_sizes: list[str], + zones: list[str], + ) -> dict[tuple[str | None, str | None, str], dict[str, Any]]: + if not self._subscription_id: + self._logger.warning( + "Azure subscription_id not available; skipping spot placement scoring" + ) + return {} + + payload = { + "desiredLocations": regions, + "desiredSizes": [{"sku": vm_size} for vm_size in vm_sizes], + "desiredCount": requested_count, + "availabilityZones": bool(zones), + } + + url = ( + "https://management.azure.com/subscriptions/" + f"{self._subscription_id}/providers/Microsoft.Compute/locations/" + f"{self._base_location}/placementScores/spot/generate" + f"?api-version={self._API_VERSION}" + ) + + try: + credential = await self._azure_client.get_async_credential() + token = await credential.get_token("https://management.azure.com/.default") + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + url, + json=payload, + headers={ + "Authorization": f"Bearer {token.token}", + "Content-Type": "application/json", + }, + ) + response.raise_for_status() + response_payload = response.json() + except Exception as exc: + self._logger.warning("Azure spot placement score lookup failed: %s", exc, exc_info=True) + return {} + + placement_scores = response_payload.get("placementScores", []) + return { + ( + entry.get("region"), + entry.get("availabilityZone"), + entry.get("sku"), + ): entry + for entry in placement_scores + } + + def _build_scores( + self, + *, + candidates: list[PlacementCandidate], + raw_scores: dict[tuple[str | None, str | None, str], dict[str, Any]], + ) -> list[PlacementScore]: + scores: list[PlacementScore] = [] + for candidate in candidates: + lookup_key = (candidate.region, candidate.zone, candidate.instance_type) + entry = raw_scores.get(lookup_key, {}) + raw_score = entry.get("score", "Low") + scores.append( + PlacementScore( + candidate=candidate, + raw_score=raw_score, + normalized_score=self._normalize_score(raw_score), + approximate=False, + metadata={ + "is_quota_available": entry.get("isQuotaAvailable"), + "raw_entry": entry, + }, + ) + ) + return scores + + @classmethod + def _normalize_score(cls, raw_score: str) -> float: + return cls._SCORE_MAP.get(str(raw_score).strip().lower(), 0.0) diff --git a/src/orb/providers/azure/infrastructure/services/ssh_key_resolver.py b/src/orb/providers/azure/infrastructure/services/ssh_key_resolver.py new file mode 100644 index 000000000..fdcf018c3 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/services/ssh_key_resolver.py @@ -0,0 +1,146 @@ +"""SSH key resolution service. + +Resolves Azure SSH Public Key resource names to actual key data via +the Compute SDK, keeping infrastructure concerns out of the domain model. +""" + +from __future__ import annotations + +from collections.abc import Awaitable +from typing import Any, Protocol + + +class AzureSshPublicKeyResourceProtocol(Protocol): + """Azure SSH public key resource fields used by ORB.""" + + @property + def public_key(self) -> str | None: + """Return the public key material stored on the resource.""" + ... + + +class AzureSshPublicKeyOperationsProtocol(Protocol): + """Azure SDK SSH public key operations used by ORB.""" + + def get( + self, + resource_group_name: str, + ssh_public_key_name: str, + **kwargs: Any, + ) -> Awaitable[AzureSshPublicKeyResourceProtocol]: + """Fetch an Azure SSH public key resource.""" + ... + + +class AzureComputeSshKeyClientProtocol(Protocol): + """Azure Compute client surface used for SSH key lookup.""" + + @property + def ssh_public_keys(self) -> AzureSshPublicKeyOperationsProtocol: + """Return SSH public key operations.""" + ... + + +def _azure_resource_not_found_error_type() -> type[Exception]: + """Resolve the Azure SDK's not-found exception lazily.""" + from azure.core.exceptions import ResourceNotFoundError + + return ResourceNotFoundError + + +def _azure_http_response_error_type() -> type[Exception]: + """Resolve the Azure SDK's HTTP response exception lazily.""" + from azure.core.exceptions import HttpResponseError + + return HttpResponseError + + +def _validate_ssh_key_request( + *, + ssh_key_name: str | None, + ssh_public_keys: list[str], +) -> list[str]: + """Return inline keys or raise when no resolvable key source is configured.""" + if ssh_public_keys: + return list(ssh_public_keys) + + if not ssh_key_name: + raise ValueError( + "Cannot resolve SSH keys: neither 'ssh_key_name' nor " + "'ssh_public_keys' is set on the template." + ) + return [] + + +def _extract_key_data( + *, + ssh_resource: AzureSshPublicKeyResourceProtocol, + ssh_key_name: str, + resource_group: str, +) -> list[str]: + """Validate resolved Azure SSH key payload and return it in ORB list form.""" + key_data = ssh_resource.public_key + if not key_data: + raise ValueError( + f"Azure SSH Public Key resource '{ssh_key_name}' " + f"in resource group '{resource_group}' exists but " + f"contains no public key data." + ) + return [key_data] + + +def _translate_ssh_key_resolution_error( + *, + exc: Exception, + ssh_key_name: str, + resource_group: str, +) -> ValueError: + """Translate Azure SDK lookup failures into stable provider-facing errors.""" + if isinstance(exc, _azure_resource_not_found_error_type()): + return ValueError( + f"Azure SSH Public Key resource '{ssh_key_name}' " + f"in resource group '{resource_group}' was not found. " + f"Ensure the resource exists and the caller has " + f"'Microsoft.Compute/sshPublicKeys/read' permission." + ) + if isinstance(exc, _azure_http_response_error_type()): + return ValueError( + f"Failed to resolve Azure SSH Public Key '{ssh_key_name}' " + f"in resource group '{resource_group}': {exc}" + ) + raise exc + + +async def resolve_ssh_keys_async( + *, + ssh_key_name: str | None, + ssh_public_keys: list[str], + resource_group: str, + compute_client: AzureComputeSshKeyClientProtocol, +) -> list[str]: + """Async variant of ``resolve_ssh_keys`` for the Azure async SDK clients.""" + inline_keys = _validate_ssh_key_request( + ssh_key_name=ssh_key_name, + ssh_public_keys=ssh_public_keys, + ) + if inline_keys: + return inline_keys + + resolved_ssh_key_name = str(ssh_key_name) + try: + ssh_resource = await compute_client.ssh_public_keys.get( + resource_group_name=resource_group, + ssh_public_key_name=resolved_ssh_key_name, + ) + return _extract_key_data( + ssh_resource=ssh_resource, + ssh_key_name=resolved_ssh_key_name, + resource_group=resource_group, + ) + except Exception as exc: + translated = _translate_ssh_key_resolution_error( + exc=exc, + ssh_key_name=resolved_ssh_key_name, + resource_group=resource_group, + ) + raise translated from exc diff --git a/src/orb/providers/azure/infrastructure/vmss_cleanup.py b/src/orb/providers/azure/infrastructure/vmss_cleanup.py new file mode 100644 index 000000000..b13839b30 --- /dev/null +++ b/src/orb/providers/azure/infrastructure/vmss_cleanup.py @@ -0,0 +1,507 @@ +"""Typed Azure-owned follow-up state and coordination for VMSS empty-delete cleanup.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from threading import RLock +from typing import Any, Mapping, Optional, Protocol, SupportsInt, TypeAlias, cast + + +class _VmssCleanupLogger(Protocol): + """Protocol for logging warnings in the VMSS cleanup coordinator.""" + + def warning(self, message: str, *args: object) -> None: + """A minimal logger protocol for the VMSS cleanup coordinator, focused on warning messages.""" + ... + + +GetVmssMemberCount: TypeAlias = Callable[..., Awaitable[Optional[int]]] +VmssExists: TypeAlias = Callable[..., Awaitable[Optional[bool]]] +BeginDeleteVmss: TypeAlias = Callable[..., Awaitable[None]] + + +# VMSS empty-delete follow-up is driven by request-status polling. Keep this +# patient because Azure can report empty membership before the scale set accepts +# a delete submission. +DEFAULT_VMSS_DELETE_FOLLOW_UP_RETRIES = 24 + + +@dataclass +class PendingVmssCleanup: + """Durable cleanup intent and submission state for one VMSS.""" + + resource_group: str + vmss_name: str + machine_ids: list[str] + delete_vmss_when_empty: bool + member_delete_submitted: bool = False + delete_submitted: bool = False + delete_retry_pending: bool = False + delete_retry_count: int = 0 + delete_retry_exhausted: bool = False + last_delete_error: Optional[str] = None + + @classmethod + def from_metadata(cls, metadata: Mapping[str, object]) -> Optional[PendingVmssCleanup]: + """Create a PendingVmssCleanup instance from a metadata mapping. + + Args: + metadata (Mapping[str, object]): Metadata containing cleanup info. + Returns: + Optional[PendingVmssCleanup]: The constructed instance, or None if invalid. + """ + resource_group = metadata.get("resource_group") + vmss_name = metadata.get("vmss_name") + if vmss_name in (None, ""): + vmss_name = metadata.get("resource_id") + raw_machine_ids = metadata.get("machine_ids", []) + if resource_group in (None, "") or vmss_name in (None, ""): + return None + if not isinstance(raw_machine_ids, list): + return None + delete_retry_count = max( + 0, + int(cast(SupportsInt, metadata.get("delete_retry_count", 0))), + ) + + machine_ids: list[str] = [] + for machine_id in raw_machine_ids: + machine_id_str = str(machine_id) + if machine_id_str and machine_id_str not in machine_ids: + machine_ids.append(machine_id_str) + + return cls( + resource_group=str(resource_group), + vmss_name=str(vmss_name), + machine_ids=machine_ids, + delete_vmss_when_empty=bool(metadata.get("delete_vmss_when_empty", False)), + member_delete_submitted=bool(metadata.get("member_delete_submitted", True)), + delete_submitted=bool(metadata.get("delete_submitted", False)), + delete_retry_pending=bool(metadata.get("delete_retry_pending", False)), + delete_retry_count=delete_retry_count, + delete_retry_exhausted=bool(metadata.get("delete_retry_exhausted", False)), + last_delete_error=( + None + if metadata.get("last_delete_error") in (None, "") + else str(metadata.get("last_delete_error")) + ), + ) + + @classmethod + def create( + cls, + *, + resource_group: str, + vmss_name: str, + machine_ids: list[str], + delete_vmss_when_empty: bool, + member_delete_submitted: bool = False, + delete_submitted: bool = False, + delete_retry_pending: bool = False, + delete_retry_count: int = 0, + delete_retry_exhausted: bool = False, + last_delete_error: Optional[str] = None, + ) -> PendingVmssCleanup: + """Create a new PendingVmssCleanup instance with the given parameters. + + Args: + resource_group (str): The Azure resource group name. + vmss_name (str): The VMSS name. + machine_ids (list[str]): List of machine IDs. + delete_vmss_when_empty (bool): Whether to delete VMSS when empty. + member_delete_submitted (bool): Whether member delete has been submitted. + delete_submitted (bool): Whether delete has been submitted. + delete_retry_pending (bool): Whether a retry is pending. + delete_retry_count (int): Number of failed follow-up delete retry attempts. + delete_retry_exhausted (bool): Whether follow-up delete retries are exhausted. + last_delete_error (Optional[str]): Last error message, if any. + Returns: + PendingVmssCleanup: The constructed instance. + """ + return cls( + resource_group=str(resource_group), + vmss_name=str(vmss_name), + machine_ids=[str(machine_id) for machine_id in machine_ids], + delete_vmss_when_empty=delete_vmss_when_empty, + member_delete_submitted=member_delete_submitted, + delete_submitted=delete_submitted, + delete_retry_pending=delete_retry_pending, + delete_retry_count=max(0, int(delete_retry_count)), + delete_retry_exhausted=delete_retry_exhausted, + last_delete_error=(None if last_delete_error in (None, "") else str(last_delete_error)), + ) + + def combine_for_same_vmss(self, other: PendingVmssCleanup) -> PendingVmssCleanup: + """Merge this cleanup with another for the same VMSS, combining machine IDs and state. + + Args: + other (PendingVmssCleanup): Another cleanup for the same VMSS. + Returns: + PendingVmssCleanup: The merged cleanup instance. + """ + merged_machine_ids = list(self.machine_ids) + for machine_id in other.machine_ids: + if machine_id not in merged_machine_ids: + merged_machine_ids.append(machine_id) + + return PendingVmssCleanup( + resource_group=self.resource_group, + vmss_name=self.vmss_name, + machine_ids=merged_machine_ids, + delete_vmss_when_empty=self.delete_vmss_when_empty or other.delete_vmss_when_empty, + member_delete_submitted=self.member_delete_submitted or other.member_delete_submitted, + delete_submitted=self.delete_submitted or other.delete_submitted, + delete_retry_pending=(self.delete_retry_pending or other.delete_retry_pending) + and not (self.delete_retry_exhausted or other.delete_retry_exhausted), + delete_retry_count=max(self.delete_retry_count, other.delete_retry_count), + delete_retry_exhausted=self.delete_retry_exhausted or other.delete_retry_exhausted, + last_delete_error=other.last_delete_error or self.last_delete_error, + ) + + def mark_delete_submitted(self) -> None: + """Mark this cleanup as having had its delete submitted.""" + self.delete_submitted = True + self.delete_retry_pending = False + self.last_delete_error = None + + def mark_delete_retry_pending(self, exc: Exception, *, max_delete_retries: int) -> None: + """Mark this cleanup as needing a retry, recording the exception message. + + Args: + exc (Exception): The exception that caused the retry. + max_delete_retries (int): Maximum follow-up delete retries before terminal failure. + """ + self.delete_submitted = False + self.delete_retry_count += 1 + self.delete_retry_exhausted = self.delete_retry_count >= max_delete_retries + self.delete_retry_pending = not self.delete_retry_exhausted + self.last_delete_error = str(exc) + + def to_metadata(self) -> dict[str, Any]: + """Convert this cleanup to a metadata dictionary for serialization. + + Returns: + dict[str, Any]: Metadata representing this cleanup. + """ + metadata: dict[str, object] = { + "resource_group": self.resource_group, + "vmss_name": self.vmss_name, + "machine_ids": list(self.machine_ids), + "delete_vmss_when_empty": self.delete_vmss_when_empty, + "member_delete_submitted": self.member_delete_submitted, + "delete_submitted": self.delete_submitted, + "delete_retry_pending": self.delete_retry_pending, + } + if self.delete_retry_count: + metadata["delete_retry_count"] = self.delete_retry_count + if self.delete_retry_exhausted: + metadata["delete_retry_exhausted"] = True + if self.last_delete_error not in (None, ""): + metadata["last_delete_error"] = self.last_delete_error + return metadata + + def to_status_detail(self) -> dict[str, Any]: + """Get a status detail dictionary for reporting purposes. + + Returns: + dict[str, Any]: Status detail for this cleanup. + """ + return self.to_metadata() + + +class VmssCleanupCoordinator: + """Azure-owned coordinator for VMSS empty-delete follow-up and retries.""" + + def __init__( + self, + *, + logger: _VmssCleanupLogger, + get_vmss_member_count: GetVmssMemberCount, + vmss_exists: VmssExists, + begin_delete_vmss: BeginDeleteVmss, + max_delete_retries: int = DEFAULT_VMSS_DELETE_FOLLOW_UP_RETRIES, + ) -> None: + """Initialize the coordinator with necessary dependencies.""" + self._logger = logger + self._get_vmss_member_count = get_vmss_member_count + self._vmss_exists = vmss_exists + self._begin_delete_vmss = begin_delete_vmss + self._max_delete_retries = max(1, int(max_delete_retries)) + self._pending_cleanups: dict[tuple[str, str], PendingVmssCleanup] = {} + self._lock = RLock() + + def clear(self) -> None: + """Clear all pending VMSS cleanups.""" + with self._lock: + self._pending_cleanups.clear() + + def record(self, handler_result: Mapping[str, object] | object) -> None: + """Record a pending cleanup from a handler result, if present. + + Args: + handler_result (Mapping[str, object] | object): Handler result possibly containing cleanup info. + """ + if not isinstance(handler_result, Mapping): + return + + provider_data = handler_result.get("provider_data") + if not isinstance(provider_data, Mapping): + return + + pending_metadata = provider_data.get("pending_resource_cleanup") + if not isinstance(pending_metadata, Mapping): + return + + pending = PendingVmssCleanup.from_metadata(pending_metadata) + if pending is None: + return + + key = (pending.resource_group, pending.vmss_name) + with self._lock: + existing = self._pending_cleanups.get(key) + self._pending_cleanups[key] = ( + pending if existing is None else existing.combine_for_same_vmss(pending) + ) + + def restore_from_request_metadata(self, request_metadata: Mapping[str, object]) -> None: + """Restore pending cleanups from request metadata. + + Args: + request_metadata (Mapping[str, object]): Metadata possibly containing cleanup info. + """ + direct_pending = request_metadata.get("pending_resource_cleanup") + if isinstance(direct_pending, Mapping): + self.record({"provider_data": request_metadata}) + + termination_requests = request_metadata.get("termination_requests") + if not isinstance(termination_requests, list): + return + + for termination_request in termination_requests: + if isinstance(termination_request, Mapping): + self.record({"provider_data": termination_request}) + + def has_pending(self, *, resource_group: Optional[str], resource_ids: list[str]) -> bool: + """Check if there are pending cleanups for the given resource group and IDs. + + Args: + resource_group (Optional[str]): The Azure resource group name. + resource_ids (list[str]): List of resource IDs. + Returns: + bool: True if any pending cleanup exists, False otherwise. + """ + if not resource_group: + return False + + with self._lock: + for resource_id in resource_ids: + if (str(resource_group), str(resource_id)) in self._pending_cleanups: + return True + return False + + def status_metadata( + self, + *, + resource_group: Optional[str], + resource_ids: list[str], + ) -> dict[str, Any]: + """Get metadata about pending cleanups for reporting. + + Args: + resource_group (Optional[str]): The Azure resource group name. + resource_ids (list[str]): List of resource IDs. + Returns: + dict[str, Any]: Metadata about pending cleanups. + """ + follow_up_details: list[dict[str, Any]] = [] + follow_up_failed = False + + if resource_group: + with self._lock: + for resource_id in self._dedupe_resource_ids(resource_ids): + pending = self._pending_cleanups.get((str(resource_group), str(resource_id))) + if pending is not None: + follow_up_details.append(pending.to_status_detail()) + follow_up_failed = follow_up_failed or pending.delete_retry_exhausted + + metadata: dict[str, Any] = { + "termination_follow_up_pending": bool(follow_up_details) and not follow_up_failed, + "termination_follow_up_details": follow_up_details, + } + if follow_up_failed: + metadata["termination_follow_up_failed"] = True + return metadata + + async def reconcile( + self, + *, + resource_group: Optional[str], + resource_ids: list[str], + observed_ids: set[str], + ) -> None: + """Reconcile pending cleanups with observed instance IDs, submitting deletes if needed. + + Args: + resource_group (Optional[str]): The Azure resource group name. + resource_ids (list[str]): List of resource IDs. + observed_ids (set[str]): Set of observed instance IDs. + """ + if not resource_group or not resource_ids: + return + + for vmss_name in self._dedupe_resource_ids(resource_ids): + await self._reconcile_one( + resource_group=str(resource_group), + vmss_name=vmss_name, + observed_ids=observed_ids, + ) + + @staticmethod + def _dedupe_resource_ids(resource_ids: list[str]) -> list[str]: + """Remove duplicates from a list of resource IDs, preserving order. + + Args: + resource_ids (list[str]): List of resource IDs. + Returns: + list[str]: Deduplicated list of resource IDs. + """ + deduped: list[str] = [] + for resource_id in resource_ids: + vmss_name = str(resource_id) + if vmss_name and vmss_name not in deduped: + deduped.append(vmss_name) + return deduped + + async def _reconcile_one( + self, + *, + resource_group: str, + vmss_name: str, + observed_ids: set[str], + ) -> None: + """Reconcile cleanup for a single VMSS, submitting delete if empty and not already submitted. + + Args: + resource_group (str): The Azure resource group name. + vmss_name (str): The VMSS name. + observed_ids (set[str]): Set of observed instance IDs. + """ + key = (resource_group, vmss_name) + + # Snapshot state under the lock. Between this read and later + # mutations another thread may call record() which replaces the + # dict entry, so mutations below must re-check identity. + with self._lock: + pending = self._pending_cleanups.get(key) + if pending is None: + return + requested_ids = set(pending.machine_ids) + delete_submitted = pending.delete_submitted + + if not requested_ids: + with self._lock: + if self._pending_cleanups.get(key) is pending: + self._pending_cleanups.pop(key, None) + return + + if delete_submitted: + await self._clear_if_vmss_is_gone( + resource_group=resource_group, + vmss_name=vmss_name, + ) + return + + if pending.delete_retry_exhausted: + return + + if requested_ids & observed_ids: + return + + try: + if await self._submit_vmss_delete_if_empty(key=key, pending=pending): + return + with self._lock: + if self._pending_cleanups.get(key) is pending: + self._pending_cleanups.pop(key, None) + except Exception as exc: + with self._lock: + current = self._pending_cleanups.get(key) + if current is not None: + current.mark_delete_retry_pending( + exc, + max_delete_retries=self._max_delete_retries, + ) + exhausted = current.delete_retry_exhausted + else: + exhausted = False + self._logger.warning( + "Failed to clean up pending VMSS '%s' in '%s': %s", + vmss_name, + resource_group, + exc, + ) + if exhausted: + self._logger.warning( + "VMSS cleanup delete retries exhausted for '%s' in '%s' " + "after %d failed follow-up attempt(s)", + vmss_name, + resource_group, + self._max_delete_retries, + ) + + async def _clear_if_vmss_is_gone( + self, + *, + resource_group: str, + vmss_name: str, + ) -> None: + """Remove pending cleanup if the VMSS no longer exists in Azure. + + Args: + resource_group (str): The Azure resource group name. + vmss_name (str): The VMSS name. + """ + if await self._vmss_exists(resource_group=resource_group, vmss_name=vmss_name) is False: + with self._lock: + self._pending_cleanups.pop((resource_group, vmss_name), None) + + async def _submit_vmss_delete_if_empty( + self, + *, + key: tuple[str, str], + pending: PendingVmssCleanup, + ) -> bool: + """Submit a VMSS delete if it is empty and deletion is required. + + Args: + key (tuple[str, str]): The (resource_group, vmss_name) key. + pending (PendingVmssCleanup): The pending cleanup object. + Returns: + bool: True if delete was submitted or still pending, False if not needed. + """ + if not pending.delete_vmss_when_empty: + return False + + member_count = await self._get_vmss_member_count( + resource_group=pending.resource_group, + vmss_name=pending.vmss_name, + ) + if member_count is None or member_count > 0: + return True + + with self._lock: + current = self._pending_cleanups.get(key) + if current is None: + return False + if current.delete_submitted: + return True + if current.delete_retry_exhausted: + return True + current.mark_delete_submitted() + + await self._begin_delete_vmss( + resource_group=pending.resource_group, + vmss_name=pending.vmss_name, + ) + return True diff --git a/src/orb/providers/azure/managers/__init__.py b/src/orb/providers/azure/managers/__init__.py new file mode 100644 index 000000000..b4bf88919 --- /dev/null +++ b/src/orb/providers/azure/managers/__init__.py @@ -0,0 +1,5 @@ +"""Azure managers.""" + +from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + +__all__: list[str] = ["AzureResourceManager"] diff --git a/src/orb/providers/azure/managers/azure_resource_manager.py b/src/orb/providers/azure/managers/azure_resource_manager.py new file mode 100644 index 000000000..d612ab101 --- /dev/null +++ b/src/orb/providers/azure/managers/azure_resource_manager.py @@ -0,0 +1,145 @@ +"""Azure Resource Manager. + +Provides the live async Azure resource-management helpers used by ORB. +""" + +from typing import Optional + +from orb.domain.base.ports import LoggingPort +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.exceptions.azure_exceptions import AzureInfrastructureError +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.services.resource_metadata_service import VmssCapacityInfo + + +@injectable +class AzureResourceManager: + """High-level resource management for Azure.""" + + def __init__( + self, + azure_client: AzureClient, + config: AzureProviderConfig, + logger: LoggingPort, + ) -> None: + self._azure_client = azure_client + self._config = config + self._logger = logger + + # ------------------------------------------------------------------ + # VMSS operations + # ------------------------------------------------------------------ + + async def get_vmss_capacity_async( + self, resource_group: str, vmss_name: str + ) -> VmssCapacityInfo: + """Return current VMSS capacity information via the async Azure SDK.""" + try: + compute = await self._azure_client.get_async_compute_client() + vmss = await compute.virtual_machine_scale_sets.get( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + sku = vmss.sku + orchestration_mode = vmss.orchestration_mode or "Flexible" + provisioned_instance_count = await self.get_vmss_member_count_async( + resource_group=resource_group, + vmss_name=vmss_name, + orchestration_mode=str(orchestration_mode), + ) + return { + "vmss_name": vmss_name, + "resource_group": resource_group, + "capacity": int(sku.capacity or 0) if sku else 0, + "vm_size": sku.name if sku else None, + "provisioning_state": vmss.provisioning_state, + "provisioned_instance_count": provisioned_instance_count or 0, + } + except Exception as exc: + self._logger.error("Failed to get VMSS capacity for %s: %s", vmss_name, exc) + raise AzureInfrastructureError(f"Failed to get VMSS capacity: {exc}") from exc + + async def get_vmss_member_count_async( + self, + resource_group: str, + vmss_name: str, + orchestration_mode: Optional[str] = None, + ) -> Optional[int]: + """Return the attached instance count via the async Azure SDK.""" + resolved_orchestration_mode = orchestration_mode + if resolved_orchestration_mode in (None, ""): + compute = await self._azure_client.get_async_compute_client() + vmss = await compute.virtual_machine_scale_sets.get( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + resolved_orchestration_mode = vmss.orchestration_mode or "Flexible" + + return await self._get_vmss_instance_count_async( + resource_group=resource_group, + vmss_name=vmss_name, + orchestration_mode=str(resolved_orchestration_mode), + ) + + async def _get_vmss_instance_count_async( + self, resource_group: str, vmss_name: str, orchestration_mode: str + ) -> Optional[int]: + """Return the attached instance count via the async Azure SDK.""" + compute = await self._azure_client.get_async_compute_client() + + try: + if str(orchestration_mode).lower() == "flexible": + count = 0 + pager = compute.virtual_machines.list(resource_group_name=resource_group) + async for vm in pager: + vmss_ref = vm.virtual_machine_scale_set + vmss_id = vmss_ref.id if vmss_ref else "" + if vmss_id and vmss_id.rstrip("/").endswith( + f"/virtualMachineScaleSets/{vmss_name}" + ): + count += 1 + return count + + count = 0 + pager = compute.virtual_machine_scale_set_vms.list( + resource_group_name=resource_group, + virtual_machine_scale_set_name=vmss_name, + ) + async for _ in pager: + count += 1 + return count + except Exception as exc: + self._logger.warning("Failed to count VMSS instances for %s: %s", vmss_name, exc) + return None + + @staticmethod + def _vmss_lookup_not_found(exc: Exception) -> bool: + """Return whether a VMSS lookup exception represents a not-found result.""" + # getattr: Exception subclasses vary — not all have error_code/status_code. + error_code = getattr(exc, "error_code", None) + if error_code in {"ResourceNotFound", "NotFound", "VMSSNotFoundError"}: + return True + + status_code = getattr(exc, "status_code", None) + if status_code == 404: + return True + + message = str(exc).lower() + return "not found" in message or "could not find" in message + + async def vmss_exists_async(self, resource_group: str, vmss_name: str) -> Optional[bool]: + """Return whether the VMSS still exists, or None if Azure could not be queried.""" + try: + compute = await self._azure_client.get_async_compute_client() + await compute.virtual_machine_scale_sets.get( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + return True + except Exception as exc: + if self._vmss_lookup_not_found(exc): + return False + + self._logger.warning("Failed to determine whether VMSS %s exists: %s", vmss_name, exc) + return None diff --git a/src/orb/providers/azure/py.typed b/src/orb/providers/azure/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/src/orb/providers/azure/registration.py b/src/orb/providers/azure/registration.py new file mode 100644 index 000000000..524210483 --- /dev/null +++ b/src/orb/providers/azure/registration.py @@ -0,0 +1,617 @@ +"""Azure Provider Registration. + +Bootstrap helpers called at application startup to register the Azure +provider with the provider registry, template extension registry, and +DI container. +""" + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Optional, Protocol, TypeVar + +from orb.config import PerformanceConfig +from orb.config.schemas.provider_strategy_schema import ProviderInstanceConfig + +if TYPE_CHECKING: + from orb.domain.base.ports import LoggingPort + from orb.providers.azure.configuration.config import AzureProviderConfig + from orb.providers.azure.infrastructure.adapters.azure_validation_adapter import ( + AzureValidationAdapter, + ) + from orb.providers.azure.infrastructure.azure_client import ( + AzureClient, + AzureClientRuntimeConfig, + ) + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + from orb.providers.registry import ProviderRegistry + +from orb.domain.template.factory import TemplateFactory +from orb.infrastructure.registry.template_extension_registry import TemplateExtensionRegistry +from orb.providers.azure.configuration.template_extension import AzureTemplateExtensionConfig + +T = TypeVar("T") + + +class PerformanceConfigSource(Protocol): + """Configuration source shape used to resolve shared performance settings.""" + + def get_typed(self, config_type: type[T]) -> T: + """Return a typed configuration object.""" + ... + + +def _resolve_performance_config( + config_port: PerformanceConfigSource | None, + logger: "LoggingPort", +) -> PerformanceConfig: + """Resolve shared performance config, falling back to defaults.""" + if config_port is None: + logger.debug("No shared config port available; using default Azure performance config") + return PerformanceConfig.model_validate({}) + + try: + perf_config = config_port.get_typed(PerformanceConfig) + except Exception as exc: + logger.debug("Could not load performance config from shared config port: %s", exc) + return PerformanceConfig.model_validate({}) + + if not isinstance(perf_config, PerformanceConfig): + logger.debug( + "Ignoring unexpected performance config type from shared config port: %s", + type(perf_config).__name__, + ) + return PerformanceConfig.model_validate({}) + + return perf_config + + +def _build_azure_client_runtime_config( + azure_config: "AzureProviderConfig", + logger: "LoggingPort", + *, + performance_config: PerformanceConfig | None = None, + config_port: PerformanceConfigSource | None = None, +) -> "AzureClientRuntimeConfig": + """Assemble the explicit runtime config owned by Azure infrastructure.""" + from orb.providers.azure.infrastructure.azure_client import AzureClientRuntimeConfig + + resolved = performance_config or _resolve_performance_config(config_port, logger) + return AzureClientRuntimeConfig( + azure_config=azure_config, + performance_config=resolved, + ) + + +def _create_azure_client( + runtime_config: "AzureClientRuntimeConfig", + logger: "LoggingPort", +) -> "AzureClient": + """Construct an ``AzureClient`` from explicit runtime config.""" + from orb.providers.azure.infrastructure.azure_client import AzureClient + + return AzureClient(runtime_config=runtime_config, logger=logger) + + +# ------------------------------------------------------------------ +# Factory functions +# ------------------------------------------------------------------ + + +def create_azure_strategy( + provider_config: Mapping[str, Any], + *, + provider_instance_name: str, + performance_config: PerformanceConfig | None = None, + config_port: PerformanceConfigSource | None = None, + azure_native_spec_service: Optional["AzureNativeSpecService"] = None, +) -> "AzureProviderStrategy": + """Create an ``AzureProviderStrategy`` from raw provider config data. + + Azure/GCP are standardized at this boundary for now: provider factories + consume only the provider's config mapping, while instance registration + is responsible for unpacking ``provider_instance.config``. + """ + from orb.infrastructure.adapters.logging_adapter import LoggingAdapter + from orb.providers.azure.configuration.config import AzureProviderConfig + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + + try: + azure_config = AzureProviderConfig(**provider_config) + logger = LoggingAdapter() + runtime_config = _build_azure_client_runtime_config( + azure_config, + logger, + performance_config=performance_config, + config_port=config_port, + ) + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name=provider_instance_name, + azure_client_resolver=lambda: _create_azure_client(runtime_config, logger), + azure_native_spec_service=azure_native_spec_service, + ) + + return strategy + except ImportError as exc: + raise ImportError(f"Azure provider strategy not available: {exc!s}") + except Exception as exc: + raise RuntimeError(f"Failed to create Azure strategy: {exc!s}") + + +def create_azure_config(data: ProviderInstanceConfig | Mapping[str, Any]) -> "AzureProviderConfig": + """Create an ``AzureProviderConfig`` from registry config input.""" + try: + from orb.providers.azure.configuration.config import AzureProviderConfig + + config_data = data.config if isinstance(data, ProviderInstanceConfig) else data + return AzureProviderConfig(**config_data) + except ImportError as exc: + raise ImportError(f"Azure configuration not available: {exc!s}") + except Exception as exc: + raise RuntimeError(f"Failed to create Azure config: {exc!s}") + + +def create_azure_validator( + provider_config: "AzureProviderConfig | Mapping[str, Any] | None" = None, +) -> Optional["AzureValidationAdapter"]: + """Create an Azure template validator.""" + from orb.infrastructure.adapters.logging_adapter import LoggingAdapter + from orb.providers.azure.configuration.config import AzureProviderConfig + from orb.providers.azure.infrastructure.adapters.azure_validation_adapter import ( + AzureValidationAdapter, + ) + + if provider_config is None: + return None + + try: + if isinstance(provider_config, AzureProviderConfig): + azure_config = provider_config + elif isinstance(provider_config, Mapping): + azure_config = AzureProviderConfig(**provider_config) + else: + return None + + return AzureValidationAdapter(config=azure_config, logger=LoggingAdapter()) + except Exception as exc: + raise RuntimeError(f"Failed to create Azure validator: {exc!s}") + + +# ------------------------------------------------------------------ +# Provider registration +# ------------------------------------------------------------------ + + +def _register_named_azure_provider_instance( + registry: "ProviderRegistry", + instance_name: str, + *, + config_port: PerformanceConfigSource | None = None, +) -> None: + """Register a named Azure provider instance with the registry.""" + registry.register_provider_instance( + provider_type="azure", + instance_name=instance_name, + # Keep wrapper unpacking at the registry boundary instead of inside + # provider factories so Azure/GCP share one local contract. + strategy_factory=lambda provider_instance_config: create_azure_strategy( + provider_instance_config.config, + provider_instance_name=instance_name, + config_port=config_port, + ), + config_factory=create_azure_config, + validator_factory=lambda provider_instance_config: create_azure_validator( + provider_instance_config.config + ), + ) + + +def register_azure_provider( + registry: Optional["ProviderRegistry"] = None, + logger: Optional["LoggingPort"] = None, + instance_name: Optional[str] = None, + config_port: PerformanceConfigSource | None = None, +) -> None: + """Register Azure provider with the provider registry.""" + if registry is None: + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + + try: + if instance_name: + _register_named_azure_provider_instance( + registry, + instance_name, + config_port=config_port, + ) + else: + registry.register_provider( + provider_type="azure", + strategy_factory=lambda provider_config: create_azure_strategy( + provider_config, + provider_instance_name="azure-default", + config_port=config_port, + ), + config_factory=create_azure_config, + validator_factory=create_azure_validator, + ) + + register_azure_auth_strategies(logger) + if logger: + logger.info("Azure provider registered successfully") + except Exception as exc: + if logger: + logger.error("Failed to register Azure provider: %s", str(exc), exc_info=True) + raise + + +def register_azure_provider_instance( + provider_instance: Any, + logger: Optional["LoggingPort"] = None, + config_port: PerformanceConfigSource | None = None, +) -> bool: + """Register an Azure provider instance using the canonical registry contract.""" + try: + if logger: + logger.debug("Registering Azure provider instance: %s", provider_instance.name) + + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + + if not registry.is_provider_registered("azure"): + register_azure_provider( + registry=registry, + logger=logger, + config_port=config_port, + ) + + _register_named_azure_provider_instance( + registry, + provider_instance.name, + config_port=config_port, + ) + + if logger: + logger.debug( + "Successfully registered Azure provider instance: %s", provider_instance.name + ) + return True + except Exception as exc: + if logger: + logger.error( + "Failed to register Azure provider instance '%s': %s", + provider_instance.name, + str(exc), + exc_info=True, + ) + return False + + +# ------------------------------------------------------------------ +# DI registration +# ------------------------------------------------------------------ + + +def register_azure_provider_with_di(provider_instance: Any, container: Any) -> bool: + """Register Azure provider instance using DI container context.""" + from orb.domain.base.ports import LoggingPort + from orb.domain.base.ports.configuration_port import ConfigurationPort + + logger = container.get(LoggingPort) + + try: + logger.debug("Registering Azure provider instance: %s", provider_instance.name) + + azure_config = create_azure_config(provider_instance.config) + + _register_azure_components_with_di(container, azure_config, provider_instance.name) + + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + + def azure_strategy_factory(_config: Any) -> Any: + """Factory to create Azure strategy using DI container.""" + return _create_azure_strategy_with_di(container, azure_config, provider_instance.name) + + if not registry.is_provider_registered("azure"): + register_azure_provider( + registry=registry, + logger=logger, + config_port=container.get(ConfigurationPort), + ) + + registry.register_provider_instance( + provider_type="azure", + instance_name=provider_instance.name, + strategy_factory=azure_strategy_factory, + config_factory=lambda _data: azure_config, + ) + + logger.debug("Successfully registered Azure provider instance: %s", provider_instance.name) + return True + + except Exception as exc: + logger.error( + "Failed to register Azure provider instance '%s': %s", + provider_instance.name, + str(exc), + exc_info=True, + ) + return False + + +def _register_azure_components_with_di( + container: Any, azure_config: "AzureProviderConfig", instance_name: str +) -> None: + """Register Azure components with DI container for a specific instance.""" + from orb.domain.base.ports import LoggingPort + from orb.providers.azure.infrastructure.azure_handler_factory import AzureHandlerFactory + + def azure_client_factory(container_instance: Any) -> Any: + """Factory to create an Azure client with the correct config and logger.""" + from orb.domain.base.ports.configuration_port import ConfigurationPort + + logger_port = container_instance.get(LoggingPort) + runtime_config = _build_azure_client_runtime_config( + azure_config, + logger_port, + config_port=container_instance.get(ConfigurationPort), + ) + client = _create_azure_client(runtime_config, logger_port) + logger_port.info( + "Azure client initialized for %s: region=%s", + instance_name, + azure_config.region, + ) + return client + + container.register_factory(f"AzureClient_{instance_name}", azure_client_factory) + + def azure_handler_factory_factory(container_instance: Any) -> Any: + """Factory to create an Azure handler factory for the named provider instance.""" + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + azure_client = container_instance.get(f"AzureClient_{instance_name}") + logger_port = container_instance.get(LoggingPort) + return AzureHandlerFactory( + azure_client=azure_client, + logger=logger_port, + azure_native_spec_service=container_instance.get_optional(AzureNativeSpecService), + azure_resource_manager=AzureResourceManager( + azure_client=azure_client, + config=azure_config, + logger=logger_port, + ), + ) + + container.register_factory( + f"AzureHandlerFactory_{instance_name}", + azure_handler_factory_factory, + ) + + +def _create_azure_strategy_with_di( + container: Any, azure_config: "AzureProviderConfig", instance_name: str +) -> "AzureProviderStrategy": + """Create Azure strategy using DI container.""" + from orb.domain.base.ports import LoggingPort + + logger = container.get(LoggingPort) + + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + + return AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name=instance_name, + azure_client_resolver=lambda: container.get(f"AzureClient_{instance_name}"), + azure_handler_factory_resolver=lambda: container.get( + f"AzureHandlerFactory_{instance_name}" + ), + azure_native_spec_service=container.get_optional(AzureNativeSpecService), + ) + + +# ------------------------------------------------------------------ +# Template extensions +# ------------------------------------------------------------------ + + +def register_azure_extensions(logger: Optional["LoggingPort"] = None) -> None: + """Register Azure template extensions with the global registry.""" + try: + TemplateExtensionRegistry.register_extension("azure", AzureTemplateExtensionConfig) + if logger: + logger.debug("Azure template extensions registered successfully") + except Exception as exc: + error_msg = f"Failed to register Azure template extensions: {exc}" + if logger: + logger.error(error_msg) + raise + + +def register_azure_provider_settings() -> None: + """Register AzureProviderConfig with the provider settings registry.""" + try: + from orb.config.schemas.provider_settings_registry import ProviderSettingsRegistry + from orb.providers.azure.configuration.config import AzureProviderConfig + + ProviderSettingsRegistry.register_provider_settings("azure", AzureProviderConfig) + except ImportError: + pass + except Exception as exc: + raise RuntimeError(f"Failed to register Azure provider settings: {exc!s}") + + +def register_azure_cli_spec() -> None: + """Register Azure CLI spec for provider add/update commands.""" + try: + from orb.infrastructure.registry.cli_spec_registry import CLISpecRegistry + from orb.providers.azure.cli.azure_cli_spec import AzureCLISpec + + CLISpecRegistry.register("azure", AzureCLISpec()) + except ImportError: + pass + except Exception as exc: + raise RuntimeError(f"Failed to register Azure CLI spec: {exc!s}") + + +def register_azure_hostfactory_field_mapping() -> None: + """Register Azure HostFactory field mapping.""" + try: + from orb.infrastructure.scheduler.hostfactory.field_mapping_registry import ( + FieldMappingRegistry, + ) + from orb.providers.azure.scheduler.hostfactory_field_mapping import ( + AzureFieldMapping, + ) + + FieldMappingRegistry.register("azure", AzureFieldMapping()) + except ImportError: + pass + except Exception as exc: + raise RuntimeError(f"Failed to register Azure HostFactory field mapping: {exc!s}") + + +def register_azure_auth_strategies(logger: Optional["LoggingPort"] = None) -> None: + """Register Azure authentication strategies with the auth registry.""" + try: + from orb.infrastructure.auth.registry import get_auth_registry + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + + registry = get_auth_registry() + if not registry.is_registered("azure"): + registry.register_strategy("azure", AzureAuthStrategy) + if logger: + logger.debug("Azure auth strategy registered") + except ImportError as exc: + if logger: + logger.warning("Azure auth strategy not available: %s", exc) + except Exception as exc: + if logger: + logger.error("Failed to register Azure auth strategy: %s", exc, exc_info=True) + raise + + +def register_azure_template_factory( + factory: TemplateFactory, logger: Optional["LoggingPort"] = None +) -> None: + """Register AzureTemplate with the template factory.""" + try: + from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate + + factory.register_provider_template_class("azure", AzureTemplate) + if logger: + logger.info("Azure template class registered with factory") + except ImportError: + if logger: + logger.debug("Azure template class not available, using core template") + except Exception as exc: + if logger: + logger.error("Failed to register Azure template factory: %s", exc, exc_info=True) + + +def get_azure_extension_defaults() -> dict[str, Any]: + """Get default Azure extension configuration.""" + return AzureTemplateExtensionConfig( + os_disk_type="Premium_LRS", + ).to_template_defaults() + + +def initialize_azure_provider( + template_factory: Optional[TemplateFactory] = None, + logger: Optional["LoggingPort"] = None, +) -> None: + """Initialize Azure provider components at application startup.""" + try: + register_azure_extensions(logger) + register_azure_provider_settings() + register_azure_cli_spec() + register_azure_hostfactory_field_mapping() + register_azure_auth_strategies(logger) + if template_factory: + register_azure_template_factory(template_factory, logger) + if logger: + logger.info("Azure provider initialization completed successfully") + except Exception as exc: + if logger: + logger.error("Azure provider initialization failed: %s", exc, exc_info=True) + raise + + +def is_azure_provider_registered() -> bool: + """Check if Azure provider extensions are registered.""" + return TemplateExtensionRegistry.has_extension("azure") + + +def register_azure_services_with_di(container) -> None: + """Register Azure services with the DI container.""" + from orb.domain.base.ports import LoggingPort + + logger = container.get(LoggingPort) + + try: + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + + if not container.is_registered(AzureNativeSpecService): + + def create_azure_native_spec_service(c): + """Factory to create AzureNativeSpecService with dependencies from DI container.""" + from orb.application.services.native_spec_service import NativeSpecService + from orb.domain.base.ports.configuration_port import ConfigurationPort + + return AzureNativeSpecService( + native_spec_service=c.get(NativeSpecService), + config_port=c.get(ConfigurationPort), + ) + + container.register_factory(AzureNativeSpecService, create_azure_native_spec_service) + logger.debug("Azure Native Spec Service registered with DI container") + except Exception as exc: + logger.warning( + "Failed to register Azure services with DI container: %s", exc, exc_info=True + ) + + from orb.infrastructure.registry.template_example_generator_registry import ( + TemplateExampleGeneratorRegistry, + ) + from orb.providers.azure.infrastructure.adapters.template_example_generator_adapter import ( + AzureTemplateExampleGeneratorAdapter, + ) + + TemplateExampleGeneratorRegistry.register("azure", AzureTemplateExampleGeneratorAdapter()) + logger.debug( + "Azure TemplateExampleGeneratorAdapter registered in TemplateExampleGeneratorRegistry" + ) + + +# ------------------------------------------------------------------ +# Auto-register extensions on import +# ------------------------------------------------------------------ + +try: + register_azure_extensions() + register_azure_provider_settings() + register_azure_cli_spec() + register_azure_hostfactory_field_mapping() + register_azure_auth_strategies() +except Exception: + import logging as _logging + + _logging.getLogger(__name__).warning( + "Failed to auto-register Azure extensions on import", exc_info=True + ) diff --git a/src/orb/providers/azure/resilience/__init__.py b/src/orb/providers/azure/resilience/__init__.py new file mode 100644 index 000000000..c0db2167b --- /dev/null +++ b/src/orb/providers/azure/resilience/__init__.py @@ -0,0 +1 @@ +"""Azure resilience utilities.""" diff --git a/src/orb/providers/azure/scheduler/__init__.py b/src/orb/providers/azure/scheduler/__init__.py new file mode 100644 index 000000000..0143f33f7 --- /dev/null +++ b/src/orb/providers/azure/scheduler/__init__.py @@ -0,0 +1 @@ +"""Azure scheduler integrations.""" diff --git a/src/orb/providers/azure/scheduler/hostfactory_field_mapping.py b/src/orb/providers/azure/scheduler/hostfactory_field_mapping.py new file mode 100644 index 000000000..85c17116d --- /dev/null +++ b/src/orb/providers/azure/scheduler/hostfactory_field_mapping.py @@ -0,0 +1,86 @@ +"""Azure implementation of ``FieldMappingPort`` for the HostFactory scheduler.""" + +from __future__ import annotations + +from orb.infrastructure.scheduler.hostfactory.field_mapping_port import FieldMappingPort + + +class AzureFieldMapping: + """Azure-specific field-mapping adapter for the HostFactory scheduler.""" + + _PROVIDER_MAPPINGS: dict[str, str] = { + # Override generic AWS-centric meanings for Azure templates. + "vmType": "vm_size", + "vmTypes": "vm_sizes", + "keyName": "ssh_key_name", + "subnetId": "network_config.subnet_id", + "securityGroupIds": "network_config.network_security_group_id", + # Azure resource targeting. + "resourceGroup": "resource_group", + "subscriptionId": "subscription_id", + # Azure VMSS / compute configuration. + "vmSize": "vm_size", + "vmSizes": "vm_sizes", + "vmSizePreferences": "vm_size_preferences", + "vmssName": "vmss_name", + "orchestrationMode": "orchestration_mode", + "platformFaultDomainCount": "platform_fault_domain_count", + "singlePlacementGroup": "single_placement_group", + # Azure pricing / placement. + "evictionPolicy": "eviction_policy", + "billingProfileMaxPrice": "billing_profile_max_price", + "spotPercentage": "spot_percentage", + "baseRegularPriorityCount": "base_regular_priority_count", + "vmssAllocationStrategy": "vmss_allocation_strategy", + "spotRestoreEnabled": "spot_restore_enabled", + "spotRestoreTimeout": "spot_restore_timeout", + "zoneBalance": "zone_balance", + "proximityPlacementGroupId": "proximity_placement_group_id", + "capacityReservationGroupId": "capacity_reservation_group_id", + # Azure storage / network / security. + "osDisk": "os_disk", + "dataDisks": "data_disks", + "networkConfig": "network_config", + "securityType": "security_type", + "secureBootEnabled": "secure_boot_enabled", + "vtpmEnabled": "vtpm_enabled", + "encryptionAtHost": "encryption_at_host", + "diskEncryptionSetId": "disk_encryption_set_id", + # Azure identity / bootstrap. + "adminUsername": "admin_username", + "sshKeyName": "ssh_key_name", + "sshPublicKeys": "ssh_public_keys", + "userAssignedIdentityIds": "user_assigned_identity_ids", + "systemAssignedIdentity": "system_assigned_identity", + "customData": "custom_data", + "extensionProfile": "extension_profile", + "upgradePolicyMode": "upgrade_policy_mode", + # Azure native spec / metadata. + "providerApiSpec": "provider_api_spec", + "providerApiSpecFile": "provider_api_spec_file", + "nodeAttributes": "node_attributes", + # Azure CycleCloud. + "clusterName": "cluster_name", + "nodeArray": "node_array", + "cyclecloudUrl": "cyclecloud_url", + "cyclecloudCredentialPath": "cyclecloud_credential_path", + "cyclecloudVerifySsl": "cyclecloud_verify_ssl", + "cyclecloudAuthMode": "cyclecloud_auth_mode", + "cyclecloudAadScope": "cyclecloud_aad_scope", + } + + def get_mappings(self) -> dict[str, str]: + """Return Azure-specific HF-field to internal-field name entries.""" + return dict(self._PROVIDER_MAPPINGS) + + def apply_defaults(self, mapped: dict) -> dict: + """Apply Azure-specific defaults after field mapping.""" + mapped.setdefault("max_instances", 1) + return mapped + + def derive_attributes(self, machine_type: str | None) -> dict[str, list[str]] | None: + """Azure does not derive cpu/ram from VM size names in HostFactory output.""" + return None + + +_: FieldMappingPort = AzureFieldMapping() # type: ignore[assignment] diff --git a/src/orb/providers/azure/services/cyclecloud_request_context_service.py b/src/orb/providers/azure/services/cyclecloud_request_context_service.py new file mode 100644 index 000000000..f0ade61be --- /dev/null +++ b/src/orb/providers/azure/services/cyclecloud_request_context_service.py @@ -0,0 +1,14 @@ +"""CycleCloud request-context resolution owned by the Azure provider.""" + +from __future__ import annotations + +from orb.providers.base.strategy import ProviderOperation + + +def resolve_cyclecloud_request_metadata( + *, + operation: ProviderOperation, +) -> dict[str, object]: + """Return request metadata supplied with the current operation.""" + request_metadata = dict(operation.parameters.get("request_metadata") or {}) + return request_metadata diff --git a/src/orb/providers/azure/services/health_check_service.py b/src/orb/providers/azure/services/health_check_service.py new file mode 100644 index 000000000..94fea171b --- /dev/null +++ b/src/orb/providers/azure/services/health_check_service.py @@ -0,0 +1,91 @@ +"""Azure provider health checks.""" + +from __future__ import annotations + +import time + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncDefaultAzureAccessTokenProvider, + DefaultAzureAccessTokenProvider, +) +from orb.providers.base.strategy import ProviderHealthStatus + + +class AzureHealthCheckService: + """Own the Azure provider health-check interaction.""" + + _MANAGEMENT_SCOPE = "https://management.azure.com/.default" + + def __init__(self, config: AzureProviderConfig, logger: LoggingPort) -> None: + self._config = config + self._logger = logger + + def check_health(self) -> ProviderHealthStatus: + """Perform a synchronous health check with short-lived Azure credentials.""" + start_time = time.time() + + try: + from orb.infrastructure.mocking.dry_run_context import is_dry_run_active + + if is_dry_run_active(): + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.healthy( + f"Azure provider healthy (DRY-RUN) - Region: {self._config.region}", + response_time_ms, + ) + + token_provider = DefaultAzureAccessTokenProvider( + client_id=self._config.client_id, + logger=self._logger, + ) + token_provider.get_access_token(self._MANAGEMENT_SCOPE) + + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.healthy( + f"Azure provider healthy - Token fetched successfully in {response_time_ms}, " + f"Region: {self._config.region}", + response_time_ms, + ) + except Exception as exc: + self._logger.warning("Azure health check failed: %s", exc, exc_info=True) + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.unhealthy( + f"Health check error: {exc!s}", + {"error": str(exc), "response_time_ms": response_time_ms}, + ) + + async def check_health_async(self) -> ProviderHealthStatus: + """Perform an async health check with short-lived async Azure credentials.""" + start_time = time.time() + + try: + from orb.infrastructure.mocking.dry_run_context import is_dry_run_active + + if is_dry_run_active(): + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.healthy( + f"Azure provider healthy (DRY-RUN) - Region: {self._config.region}", + response_time_ms, + ) + + token_provider = AsyncDefaultAzureAccessTokenProvider( + client_id=self._config.client_id, + logger=self._logger, + ) + await token_provider.get_access_token(self._MANAGEMENT_SCOPE) + + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.healthy( + f"Azure provider healthy - Token fetched successfully in {response_time_ms}, " + f"Region: {self._config.region}", + response_time_ms, + ) + except Exception as exc: + self._logger.warning("Azure health check failed: %s", exc, exc_info=True) + response_time_ms = (time.time() - start_time) * 1000 + return ProviderHealthStatus.unhealthy( + f"Health check error: {exc!s}", + {"error": str(exc), "response_time_ms": response_time_ms}, + ) diff --git a/src/orb/providers/azure/services/inventory_service.py b/src/orb/providers/azure/services/inventory_service.py new file mode 100644 index 000000000..e4142841a --- /dev/null +++ b/src/orb/providers/azure/services/inventory_service.py @@ -0,0 +1,785 @@ +"""Azure read/query orchestration for status and resource discovery. + +Owns the read-side vocabulary (``AzureReadOperationContext``, +``AzureStatusResult``, the normalization helpers) and the +``AzureInventoryService`` orchestrator. Generic ``ProviderOperation`` +parsing lives in ``operation_parsing`` because it is also used by the +write paths (``termination_service``, ``provisioning_service``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Optional, Protocol, TypedDict, runtime_checkable + +from orb.domain.base.ports import LoggingPort +from orb.domain.request.aggregate import Request +from orb.domain.request.value_objects import RequestType +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError +from orb.providers.azure.infrastructure.cyclecloud_session import CycleCloudRequestContext +from orb.providers.azure.infrastructure.error_codes import ProviderErrorEntry +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + RAISE_ON_STATUS_ERROR_METADATA_KEY, + AzureHandler, + AzureHandlerStatusResult, +) +from orb.providers.azure.infrastructure.vmss_cleanup import VmssCleanupCoordinator +from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager +from orb.providers.azure.services.operation_parsing import ( + group_instance_ids_by_resource, + operation_request_id, + resolve_operation_provider_api, + resolve_operation_resource_group, +) +from orb.providers.azure.services.resource_metadata_service import ( + AzureDeploymentStatusServiceProtocol, + AzureResourceMetadataService, +) +from orb.providers.base.strategy import ProviderOperation, ProviderResult + +if TYPE_CHECKING: + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + + +# --------------------------------------------------------------------------- +# Read-side vocabulary +# --------------------------------------------------------------------------- + + +@dataclass +class AzureReadOperationContext: + """Provider-owned runtime context for Azure status and describe operations.""" + + operation_name: str + request_id: str | None + template_id: str + request_metadata: dict[str, Any] + cyclecloud_request_context: CycleCloudRequestContext + provider_api: Optional[AzureProviderApi] + provider_api_key: str | None + resource_group: str | None + instance_ids: list[str] = field(default_factory=list) + resource_ids: list[str] = field(default_factory=list) + grouped_resource_mapping: dict[str, list[str]] = field(default_factory=dict) + direct_resource_id: str | None = None + raise_on_status_error: bool = False + + +class AzureStatusProviderData(TypedDict, total=False): + """Provider-owned identity fields surfaced by Azure handlers.""" + + vm_id: str + vmss_instance_id: str + node_id: str + node_name: str + vm_name: str + + +class AzureStatusResult(TypedDict, total=False): + """Normalized Azure status result used for cross-handler identity matching.""" + + instance_id: str + name: str + provider_data: AzureStatusProviderData + + +CYCLECLOUD_METADATA_KEYS = ( + "cluster_name", + "node_array", + "node_ids", + "operation_id", + "operation_location", + "cyclecloud_url", + "cyclecloud_credential_path", + "cyclecloud_verify_ssl", + "cyclecloud_auth_mode", + "cyclecloud_aad_scope", +) + + +def normalize_status_result(result: AzureHandlerStatusResult) -> AzureStatusResult: + """Build an AzureStatusResult from a generic handler status dict.""" + normalized: AzureStatusResult = {} + + instance_id = result.get("instance_id") + if isinstance(instance_id, str) and instance_id: + normalized["instance_id"] = instance_id + + name = result.get("name") + if isinstance(name, str) and name: + normalized["name"] = name + + raw_provider_data = result.get("provider_data") + if raw_provider_data: + provider_data: AzureStatusProviderData = {} + for key in ("vm_id", "vmss_instance_id", "node_id", "node_name", "vm_name"): + value = raw_provider_data.get(key) + if isinstance(value, str) and value: + provider_data[key] = value + if provider_data: + normalized["provider_data"] = provider_data + + return normalized + + +def status_candidate_ids(result: AzureHandlerStatusResult) -> set[str]: + """Return all plausible instance identifiers from a status result.""" + identity = normalize_status_result(result) + provider_data = identity.get("provider_data") or {} + candidate_ids = { + identity.get("instance_id", ""), + identity.get("name", ""), + provider_data.get("vm_id", ""), + provider_data.get("vmss_instance_id", ""), + provider_data.get("node_id", ""), + provider_data.get("node_name", ""), + provider_data.get("vm_name", ""), + } + candidate_ids.discard("") + return candidate_ids + + +def observed_status_ids(instance_details: list[AzureHandlerStatusResult]) -> set[str]: + """Collect all candidate identifiers across a list of instance detail dicts.""" + observed_ids: set[str] = set() + for instance in instance_details: + observed_ids.update(status_candidate_ids(instance)) + return observed_ids + + +def filter_status_results( + results: list[AzureHandlerStatusResult], + requested_ids: list[str], +) -> list[AzureHandlerStatusResult]: + """Return only the results whose candidate IDs overlap with the requested set.""" + requested = {str(item) for item in requested_ids} + filtered: list[AzureHandlerStatusResult] = [] + for result in results: + if status_candidate_ids(result) & requested: + filtered.append(result) + return filtered + + +def collect_status_resource_ids( + grouped_resource_mapping: dict[str, list[str]], + direct_resource_id: object, +) -> list[str]: + """Collect all Azure resource ids relevant to a status query.""" + resource_ids: list[str] = [] + for resource_id, mapped_ids in grouped_resource_mapping.items(): + if resource_id and mapped_ids and resource_id not in resource_ids: + resource_ids.append(resource_id) + + if direct_resource_id not in (None, "") and str(direct_resource_id) not in resource_ids: + resource_ids.append(str(direct_resource_id)) + return resource_ids + + +def status_resource_ids( + operation: ProviderOperation, + instance_ids: list[str], +) -> list[str]: + """Resolve the Azure resource IDs relevant to the given instance IDs.""" + raw_resource_mapping = operation.parameters.get("resource_mapping", {}) or {} + grouped_resource_mapping = group_instance_ids_by_resource( + instance_ids, + raw_resource_mapping, + ) + return collect_status_resource_ids( + grouped_resource_mapping, + operation.parameters.get("resource_id"), + ) + + +def build_cyclecloud_request_metadata( + *, + operation: ProviderOperation, + resource_group: Optional[str], +) -> dict[str, Any]: + """Build the metadata dict required for CycleCloud handler calls.""" + metadata: dict[str, Any] = {"resource_group": resource_group} + metadata_from_request = operation.parameters.get("request_metadata") or {} + for key in CYCLECLOUD_METADATA_KEYS: + value = metadata_from_request.get(key) + if value not in (None, ""): + metadata[key] = value + return metadata + + +def build_read_operation_context( + *, + operation: ProviderOperation, + operation_name: str, + default_resource_group: Optional[str], +) -> AzureReadOperationContext: + """Build the provider-owned runtime context for Azure read operations.""" + metadata: dict[str, Any] = dict(operation.parameters.get("request_metadata") or {}) + cyclecloud_request_context = CycleCloudRequestContext.from_mapping(metadata) + provider_api = resolve_operation_provider_api(operation) + provider_api_key = provider_api.value if provider_api is not None else None + + resource_group = resolve_operation_resource_group(operation, default_resource_group) + request_id = operation_request_id(operation) + template_id = str(operation.parameters.get("template_id", "unknown")) + + if operation_name == "get_instance_status": + instance_ids = list(operation.parameters.get("instance_ids", []) or []) + if not instance_ids: + raise AzureValidationError( + "Instance IDs are required for status query", + error_code="MISSING_INSTANCE_IDS", + ) + if not resource_group: + raise AzureValidationError( + "resource_group is required for status query", + error_code="MISSING_RESOURCE_GROUP", + ) + + grouped_resource_mapping = group_instance_ids_by_resource( + instance_ids, + operation.parameters.get("resource_mapping", {}) or {}, + ) + direct_resource_id = operation.parameters.get("resource_id") + if ( + direct_resource_id in (None, "") + and provider_api == AzureProviderApi.CYCLECLOUD + and cyclecloud_request_context.cluster_name not in (None, "") + ): + direct_resource_id = cyclecloud_request_context.cluster_name + + return AzureReadOperationContext( + operation_name=operation_name, + request_id=str(request_id) if request_id not in (None, "") else None, + template_id=template_id, + request_metadata=metadata, + cyclecloud_request_context=cyclecloud_request_context, + provider_api=provider_api, + provider_api_key=provider_api_key, + resource_group=resource_group, + instance_ids=instance_ids, + resource_ids=collect_status_resource_ids(grouped_resource_mapping, direct_resource_id), + grouped_resource_mapping=grouped_resource_mapping, + direct_resource_id=str(direct_resource_id) + if direct_resource_id not in (None, "") + else None, + ) + + resource_ids = list(operation.parameters.get("resource_ids", []) or []) + if not resource_ids: + raise AzureValidationError( + "Resource IDs are required for instance discovery", + error_code="MISSING_RESOURCE_IDS", + ) + if provider_api in (None, ""): + raise AzureValidationError( + "provider_api is required for Azure resource discovery", + error_code="MISSING_PROVIDER_API", + ) + + return AzureReadOperationContext( + operation_name=operation_name, + request_id=str(request_id) if request_id not in (None, "") else None, + template_id=template_id, + request_metadata=metadata, + cyclecloud_request_context=cyclecloud_request_context, + provider_api=provider_api, + provider_api_key=provider_api_key, + resource_group=resource_group, + resource_ids=resource_ids, + direct_resource_id=resource_ids[0] if len(resource_ids) == 1 else None, + ) + + +def build_read_handler_request( + *, + read_context: AzureReadOperationContext, + provider_name: str, + resource_ids: list[str], + additional_metadata: Optional[dict[str, Any]] = None, +) -> Request: + """Build the Request object used for Azure read/query handler calls.""" + metadata: dict[str, Any] = {"resource_group": read_context.resource_group} + metadata.update(read_context.cyclecloud_request_context.to_metadata()) + if additional_metadata: + metadata.update(additional_metadata) + + request = Request.create_new_request( + request_type=RequestType.ACQUIRE, + template_id=read_context.template_id, + machine_count=1, + provider_type="azure", + provider_name=provider_name, + request_id=read_context.request_id, + metadata=metadata, + ) + request.resource_ids = resource_ids + if read_context.provider_api_key: + request.provider_api = read_context.provider_api_key + return request + + +# --------------------------------------------------------------------------- +# Read-side orchestrator +# --------------------------------------------------------------------------- + + +@runtime_checkable +class VmssResourceErrorReader(Protocol): + """Capability interface for VMSS-level fleet error inspection.""" + + async def get_vmss_resource_errors_async( + self, + resource_group: str, + resource_id: str, + ) -> list[ProviderErrorEntry]: + """Return VMSS resource-level errors for one scale set.""" + ... + + +class AzureInventoryService: + """Own Azure read/query orchestration separate from strategy lifecycle concerns.""" + + def __init__( + self, + *, + logger: LoggingPort, + provider_instance_name: str, + resource_metadata_service: AzureResourceMetadataService, + handler_provider: AzureProviderStrategy, + vmss_cleanup_coordinator: VmssCleanupCoordinator, + ) -> None: + self._logger = logger + self._provider_instance_name = provider_instance_name + self._resource_metadata_service = resource_metadata_service + self._handler_provider = handler_provider + self._vmss_cleanup_coordinator = vmss_cleanup_coordinator + + async def get_instance_status_async( + self, + read_context: AzureReadOperationContext, + ) -> ProviderResult: + """Resolve instance status through the async Azure handler contract.""" + self._vmss_cleanup_coordinator.restore_from_request_metadata(read_context.request_metadata) + handler_machines = await self._get_instance_status_via_handlers_async(read_context) + if handler_machines is not None: + return await self._build_instance_status_result( + read_context=read_context, + handler_machines=handler_machines, + ) + + raise AzureValidationError( + "Azure get_instance_status requires provider_api-backed handler resolution.", + error_code="MISSING_PROVIDER_API", + ) + + async def describe_resource_instances_async( + self, + *, + read_context: AzureReadOperationContext, + resource_manager: Optional[AzureResourceManager], + deployment_service: AzureDeploymentStatusServiceProtocol | None, + ) -> ProviderResult: + """Describe Azure resource-backed instances through the async handler contract.""" + resource_group = read_context.resource_group + is_vmss = self._prepare_describe_context(read_context) + result = await self._describe_resource_instances_via_handler_async( + read_context=read_context, + resource_manager=resource_manager, + deployment_service=deployment_service, + ) + + if result.success and is_vmss and resource_group: + instance_details = result.data.get("instances", []) if result.data else [] + await self._vmss_cleanup_coordinator.reconcile( + resource_group=resource_group, + resource_ids=read_context.resource_ids, + observed_ids=observed_status_ids(instance_details), + ) + return result.model_copy( + update={ + "metadata": { + **(result.metadata or {}), + **self._vmss_cleanup_coordinator.status_metadata( + resource_group=resource_group, + resource_ids=read_context.resource_ids, + ), + } + } + ) + + return result + + async def _get_instance_status_via_handlers_async( + self, + read_context: AzureReadOperationContext, + ) -> Optional[list[AzureHandlerStatusResult]]: + provider_api = read_context.provider_api + grouped_resource_mapping = read_context.grouped_resource_mapping + + if not provider_api: + return None + + handler = self._resolve_status_handler(provider_api) + if not handler and not grouped_resource_mapping: + return None + + if provider_api == AzureProviderApi.SINGLE_VM and handler: + request = build_read_handler_request( + read_context=read_context, + provider_name=self._provider_instance_name, + resource_ids=read_context.instance_ids, + ) + return await handler.check_hosts_status_async(request) + + if grouped_resource_mapping: + all_results: list[AzureHandlerStatusResult] = [] + seen_instance_ids: set[str] = set() + for resource_id, mapped_ids in grouped_resource_mapping.items(): + group_handler = handler or self._resolve_status_handler(provider_api) + if not group_handler: + continue + + request = build_read_handler_request( + read_context=read_context, + provider_name=self._provider_instance_name, + resource_ids=[resource_id], + additional_metadata=self._handler_status_metadata( + provider_api=provider_api, + instance_ids=mapped_ids, + ), + ) + machines = await group_handler.check_hosts_status_async(request) + self._append_unique_status_results( + destination=all_results, + seen_instance_ids=seen_instance_ids, + machines=filter_status_results(machines, mapped_ids), + ) + + if all_results: + return all_results + + resource_id = read_context.direct_resource_id + if not handler or not resource_id: + return None + + request = build_read_handler_request( + read_context=read_context, + provider_name=self._provider_instance_name, + resource_ids=( + read_context.instance_ids + if provider_api == AzureProviderApi.SINGLE_VM + else [resource_id] + ), + additional_metadata=self._handler_status_metadata( + provider_api=provider_api, + instance_ids=read_context.instance_ids, + ), + ) + machines = await handler.check_hosts_status_async(request) + if provider_api == AzureProviderApi.SINGLE_VM: + return machines + return filter_status_results(machines, read_context.instance_ids) + + async def _describe_resource_instances_via_handler_async( + self, + *, + read_context: AzureReadOperationContext, + resource_manager: Optional[AzureResourceManager], + deployment_service: AzureDeploymentStatusServiceProtocol | None, + ) -> ProviderResult: + resolved = self._resolve_describe_handler(read_context) + if isinstance(resolved, ProviderResult): + return resolved + provider_api, handler = resolved + + request = self._build_describe_handler_request(read_context=read_context) + instance_details = await handler.check_hosts_status_async(request) + return await self._build_describe_instances_result( + read_context=read_context, + handler=handler, + instance_details=instance_details, + resource_manager=resource_manager, + deployment_service=deployment_service, + include_shortfall_metadata=provider_api != AzureProviderApi.CYCLECLOUD, + ) + + async def _build_instance_status_result( + self, + *, + read_context: AzureReadOperationContext, + handler_machines: list[AzureHandlerStatusResult], + ) -> ProviderResult: + """Build the normalized handler-backed instance-status result.""" + is_vmss = read_context.provider_api in ( + AzureProviderApi.VMSS, + AzureProviderApi.VMSS_UNIFORM, + ) + metadata: dict[str, Any] = { + "operation": "get_instance_status", + "instance_ids": read_context.instance_ids, + "method": "handler", + } + status_result = self._resource_metadata_service.attach_provider_fulfilment( + metadata, + instances=handler_machines, + target_units=len(read_context.instance_ids), + ) + result = ProviderResult.success_result( + { + "instances": status_result.instances, + "queried_count": len(read_context.instance_ids), + }, + metadata, + ) + if is_vmss and read_context.resource_group and read_context.resource_ids: + await self._vmss_cleanup_coordinator.reconcile( + resource_group=read_context.resource_group, + resource_ids=read_context.resource_ids, + observed_ids=observed_status_ids(handler_machines), + ) + return result.model_copy( + update={ + "metadata": { + **(result.metadata or {}), + **self._vmss_cleanup_coordinator.status_metadata( + resource_group=read_context.resource_group, + resource_ids=read_context.resource_ids, + ), + } + } + ) + return result + + def _prepare_describe_context( + self, + read_context: AzureReadOperationContext, + ) -> bool: + """Restore VMSS cleanup state and set the describe status failure policy.""" + is_vmss = read_context.provider_api in ( + AzureProviderApi.VMSS, + AzureProviderApi.VMSS_UNIFORM, + ) + self._vmss_cleanup_coordinator.restore_from_request_metadata(read_context.request_metadata) + read_context.raise_on_status_error = is_vmss and self._vmss_cleanup_coordinator.has_pending( + resource_group=read_context.resource_group, + resource_ids=read_context.resource_ids, + ) + return is_vmss + + def _build_describe_handler_request( + self, + *, + read_context: AzureReadOperationContext, + ) -> Request: + """Build the handler request used for Azure resource-instance discovery.""" + extra_metadata: dict[str, Any] = {} + if read_context.provider_api == AzureProviderApi.SINGLE_VM: + deployment_name = read_context.request_metadata.get("deployment_name") + if deployment_name not in (None, ""): + extra_metadata["deployment_name"] = str(deployment_name) + extra_metadata[RAISE_ON_STATUS_ERROR_METADATA_KEY] = read_context.raise_on_status_error + return build_read_handler_request( + read_context=read_context, + provider_name=self._provider_instance_name, + resource_ids=read_context.resource_ids, + additional_metadata=extra_metadata or None, + ) + + @staticmethod + def _append_unique_status_results( + *, + destination: list[AzureHandlerStatusResult], + seen_instance_ids: set[str], + machines: list[AzureHandlerStatusResult], + ) -> None: + """Append status results while preserving first-seen instance identities.""" + for machine in machines: + identity = normalize_status_result(machine) + provider_data = identity.get("provider_data") or {} + machine_id = next( + ( + value + for value in ( + identity.get("instance_id"), + identity.get("name"), + provider_data.get("vm_id"), + provider_data.get("vmss_instance_id"), + provider_data.get("node_id"), + provider_data.get("node_name"), + provider_data.get("vm_name"), + ) + if value + ), + "", + ) + if machine_id and machine_id in seen_instance_ids: + continue + destination.append(machine) + if machine_id: + seen_instance_ids.add(machine_id) + + @staticmethod + def _handler_status_metadata( + *, + provider_api: AzureProviderApi, + instance_ids: list[str], + ) -> dict[str, Any] | None: + """Return provider-specific handler metadata for status requests.""" + if provider_api == AzureProviderApi.CYCLECLOUD: + return {"node_ids": instance_ids} + return None + + def _resolve_status_handler( + self, + provider_api: AzureProviderApi, + ) -> Optional[AzureHandler]: + """Resolve a handler for status operations with VMSS uniform fallback enabled.""" + return self._handler_provider.resolve_handler( + provider_api, + allow_vmss_uniform_fallback=True, + ) + + def _resolve_describe_handler( + self, + read_context: AzureReadOperationContext, + ) -> ProviderResult | tuple[AzureProviderApi, AzureHandler]: + """Resolve the concrete handler for describe operations or return a failure result.""" + provider_api = read_context.provider_api + provider_api_key = read_context.provider_api_key or "" + if provider_api is None: + return ProviderResult.error_result( + "provider_api is required for Azure resource discovery", + "MISSING_PROVIDER_API", + ) + + handler = self._handler_provider.resolve_handler(provider_api) + if not handler: + return ProviderResult.error_result( + f"No handler available for provider_api: {provider_api_key}", + "HANDLER_NOT_FOUND", + ) + return provider_api, handler + + @staticmethod + def _collect_instance_fleet_errors( + instance_details: list[AzureHandlerStatusResult], + ) -> list[dict[str, Any]]: + """Collect distinct fleet errors embedded in handler provider data.""" + fleet_errors: list[dict[str, Any]] = [] + for inst in instance_details: + provider_data = inst.get("provider_data") or {} + for error in provider_data.get("fleet_errors") or []: + if error not in fleet_errors: + fleet_errors.append(error) + return fleet_errors + + @staticmethod + async def _get_optional_vmss_resource_errors( + handler: AzureHandler, + logger: LoggingPort, + *, + resource_group: str | None, + resource_ids: list[str], + ) -> list[ProviderErrorEntry]: + """Read VMSS resource-level errors when the concrete handler exposes them.""" + vmss_errors: list[ProviderErrorEntry] = [] + if not resource_group: + return vmss_errors + if not isinstance(handler, VmssResourceErrorReader): + logger.warning( + "VMSS resource error lookup requested from handler '%s' without VMSS error support", + type(handler).__name__, + ) + return vmss_errors + for resource_id in resource_ids: + raw_errors = await handler.get_vmss_resource_errors_async( + resource_group, + resource_id, + ) + for error in raw_errors: + if error not in vmss_errors: + vmss_errors.append(error) + return vmss_errors + + async def _build_describe_instances_result( + self, + *, + read_context: AzureReadOperationContext, + handler: AzureHandler, + instance_details: list[AzureHandlerStatusResult], + resource_manager: Optional[AzureResourceManager], + deployment_service: AzureDeploymentStatusServiceProtocol | None, + include_shortfall_metadata: bool, + ) -> ProviderResult: + """Build the normalized describe-resource-instances result and metadata.""" + provider_api = read_context.provider_api + provider_api_key = read_context.provider_api_key or "" + resource_ids = read_context.resource_ids + resource_group = read_context.resource_group + metadata: dict[str, Any] = { + "operation": "describe_resource_instances", + "resource_ids": resource_ids, + "provider_api": provider_api_key, + "handler_used": provider_api_key, + "instance_count": len(instance_details), + } + + if not instance_details: + if provider_api in (AzureProviderApi.VMSS, AzureProviderApi.VMSS_UNIFORM): + vmss_errors = await self._get_optional_vmss_resource_errors( + handler, + self._logger, + resource_group=resource_group, + resource_ids=resource_ids, + ) + if vmss_errors: + metadata["fleet_errors"] = vmss_errors + await self._resource_metadata_service.augment_vmss_capacity_metadata_async( + metadata, + resource_ids, + resource_manager=resource_manager, + resource_group=resource_group, + ) + elif provider_api == AzureProviderApi.SINGLE_VM: + await self._resource_metadata_service.augment_single_vm_deployment_metadata_async( + metadata, + read_context.request_metadata, + resource_group=resource_group, + deployment_service=deployment_service, + ) + status_result = self._resource_metadata_service.attach_provider_fulfilment( + metadata, + instances=[], + target_units=( + len(resource_ids) if provider_api == AzureProviderApi.SINGLE_VM else None + ), + ) + return ProviderResult.success_result( + {"instances": status_result.instances}, + metadata, + ) + + fleet_errors = self._collect_instance_fleet_errors(instance_details) + if fleet_errors: + metadata["fleet_errors"] = fleet_errors + if provider_api in (AzureProviderApi.VMSS, AzureProviderApi.VMSS_UNIFORM): + await self._resource_metadata_service.augment_vmss_capacity_metadata_async( + metadata, + resource_ids, + resource_manager=resource_manager, + resource_group=resource_group, + ) + if include_shortfall_metadata: + self._resource_metadata_service.augment_shortfall_metadata(metadata) + status_result = self._resource_metadata_service.attach_provider_fulfilment( + metadata, + instances=instance_details, + target_units=( + len(resource_ids) if provider_api == AzureProviderApi.SINGLE_VM else None + ), + ) + return ProviderResult.success_result( + data={"instances": status_result.instances}, + metadata=metadata, + ) diff --git a/src/orb/providers/azure/services/operation_parsing.py b/src/orb/providers/azure/services/operation_parsing.py new file mode 100644 index 000000000..a94cdbd17 --- /dev/null +++ b/src/orb/providers/azure/services/operation_parsing.py @@ -0,0 +1,74 @@ +"""Generic helpers for parsing Azure ``ProviderOperation`` parameters. + +Used by every Azure service that consumes a ``ProviderOperation`` — read +paths (``inventory_service``) and write paths (``termination_service``, +``provisioning_service``) alike. No read- or write-specific logic. +""" + +from __future__ import annotations + +from typing import Optional + +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError +from orb.providers.base.strategy import ProviderOperation + + +def operation_request_id(operation: ProviderOperation) -> str | None: + """Extract the request id from operation parameters or context.""" + return operation.parameters.get("request_id") or ( + operation.context.get("request_id") if operation.context else None + ) + + +def resolve_operation_resource_group( + operation: ProviderOperation, + default_resource_group: Optional[str], +) -> Optional[str]: + """Return the resource group from request metadata, falling back to the default.""" + metadata = operation.parameters.get("request_metadata") or {} + request_resource_group = metadata.get("resource_group") + if request_resource_group not in (None, ""): + return str(request_resource_group) + return default_resource_group + + +def resolve_operation_provider_api( + operation: ProviderOperation, +) -> Optional[AzureProviderApi]: + """Resolve the Azure provider API carried by an operation. + + The ``provider_api`` parameter is a string (the enum's ``.value``); enum + constructors are idempotent so an enum instance would also coerce, but + every construction site in the repo passes a string. + """ + raw = operation.parameters.get("provider_api") + if not raw: + return None + try: + return AzureProviderApi(raw) + except ValueError as exc: + raise AzureValidationError( + f"Invalid Azure provider_api: {raw!r}", + error_code="INVALID_PROVIDER_API", + ) from exc + + +def group_instance_ids_by_resource( + instance_ids: list[str], + resource_mapping: dict[str, tuple[Optional[str], int]], +) -> dict[str, list[str]]: + """Group requested instance IDs by their owning Azure resource ID. + + The mapping shape comes from the deprovisioning orchestrator — + ``{instance_id: (resource_id, desired_capacity)}``. Entries without a + ``resource_id`` and instances not in ``instance_ids`` are skipped. + """ + grouped: dict[str, list[str]] = {} + for instance_id, (resource_id, _capacity) in resource_mapping.items(): + if not resource_id or instance_id not in instance_ids: + continue + bucket = grouped.setdefault(resource_id, []) + if instance_id not in bucket: + bucket.append(instance_id) + return grouped diff --git a/src/orb/providers/azure/services/provisioning_service.py b/src/orb/providers/azure/services/provisioning_service.py new file mode 100644 index 000000000..4ca4e1e27 --- /dev/null +++ b/src/orb/providers/azure/services/provisioning_service.py @@ -0,0 +1,225 @@ +"""Azure create-instance orchestration helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from orb.domain.request.aggregate import Request +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions import AzureValidationError +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + AzureAcquireHostsResult, + AzureHandler, +) +from orb.providers.base.strategy import ProviderOperation, ProviderResult + + +@dataclass +class CreateOperationContext: + """Context object encapsulating all necessary information for handling a create operation.""" + + template_config: dict[str, Any] + count: int + provider_api: AzureProviderApi + provider_api_key: str + handler: AzureHandler + azure_template: AzureTemplate + + +def get_create_template_config(operation: ProviderOperation) -> dict[str, Any]: + """Extract the template configuration from the provider operation.""" + return dict(operation.parameters.get("template_config") or {}) + + +def get_create_count(operation: ProviderOperation) -> int: + """Extract the instance count from the provider operation.""" + return operation.parameters.get("count", 1) + + +def validate_create_template_config( + template_config: dict[str, Any], +) -> None: + """Validate that the template configuration is present.""" + if template_config: + return + raise AzureValidationError( + "Template configuration is required for instance creation", + error_code="MISSING_TEMPLATE_CONFIG", + ) + + +def provider_api_key(provider_api: AzureProviderApi) -> str: + """Get the string key for the provider API value.""" + return provider_api.value + + +def resolve_create_provider_api( + template_config: dict[str, Any], +) -> AzureProviderApi: + """Resolve the provider API from the template config.""" + provider_api = template_config.get("provider_api", AzureProviderApi.VMSS) + if isinstance(provider_api, AzureProviderApi): + return provider_api + if isinstance(provider_api, str): + try: + return AzureProviderApi(provider_api) + except ValueError as exc: + raise AzureValidationError( + f"Invalid Azure provider_api: {provider_api!r}", + error_code="INVALID_PROVIDER_API", + ) from exc + raise AzureValidationError( + f"Invalid Azure provider_api: {provider_api!r}", + error_code="INVALID_PROVIDER_API", + ) + + +def create_instances_dry_run_result( + create_context: CreateOperationContext, +) -> ProviderResult: + """Return a dry-run result for create instances operation.""" + return ProviderResult.success_result( + { + "resource_ids": ["dry-run-resource-id"], + "instances": [], + "provider_api": create_context.provider_api_key, + "count": create_context.count, + "template_id": create_context.azure_template.template_id, + }, + { + "operation": "create_instances", + "template_config": create_context.template_config, + "handler_used": create_context.provider_api_key, + "method": "dry_run", + "provider_data": {"dry_run": True}, + }, + ) + + +class AzureProvisioningService: + """Own Azure create-operation orchestration and handler result shaping.""" + + def build_create_operation_context( + self, + *, + operation: ProviderOperation, + resolve_handler: Callable[[AzureProviderApi], Optional[AzureHandler]], + build_template: Callable[[dict[str, Any]], AzureTemplate], + ) -> CreateOperationContext: + """Build and validate the context required for a create operation.""" + template_config = get_create_template_config(operation) + count = get_create_count(operation) + validate_create_template_config(template_config) + + provider_api = resolve_create_provider_api(template_config) + provider_api_value = provider_api_key(provider_api) + handler = resolve_handler(provider_api) + if handler is None: + raise AzureValidationError( + f"No handler available for provider_api: {provider_api_value}", + error_code="HANDLER_NOT_FOUND", + ) + + return CreateOperationContext( + template_config=template_config, + count=count, + provider_api=provider_api, + provider_api_key=provider_api_value, + handler=handler, + azure_template=build_template(template_config), + ) + + def build_create_request( + self, + *, + operation: ProviderOperation, + azure_template: AzureTemplate, + count: int, + provider_api: AzureProviderApi, + provider_instance_name: str, + ) -> Request: + """Build a request object for the create operation handler.""" + from orb.domain.request.value_objects import RequestType + + request_metadata = dict(operation.parameters.get("request_metadata", {}) or {}) + request_id = operation.parameters.get("request_id") or ( + operation.context.get("request_id") if operation.context else None + ) + request = Request.create_new_request( + request_type=RequestType.ACQUIRE, + template_id=azure_template.template_id, + machine_count=count, + provider_type="azure", + provider_name=provider_instance_name, + metadata=request_metadata, + request_id=request_id, + ) + request.provider_api = provider_api_key(provider_api) + return request + + def normalize_handler_create_result( + self, + handler_result: AzureAcquireHostsResult, + *, + template_config: dict[str, Any], + provider_api: AzureProviderApi, + count: int, + template_id: str, + ) -> ProviderResult: + """Normalize the result from the handler into a ProviderResult.""" + resource_ids = handler_result["resource_ids"] + instances = handler_result["instances"] + success = handler_result["success"] + error_message = handler_result.get("error_message") + provider_data = handler_result.get("provider_data", {}) + + if not success: + return ProviderResult.error_result( + f"Provisioning failed: {error_message}", + "PROVISIONING_ADAPTER_ERROR", + { + "operation": "create_instances", + "template_config": template_config, + "handler_used": provider_api_key(provider_api), + "method": "handler", + "provider_data": provider_data, + }, + ) + + return ProviderResult.success_result( + { + "resource_ids": resource_ids, + "instances": instances, + "provider_api": provider_api_key(provider_api), + "count": count, + "template_id": template_id, + }, + { + "operation": "create_instances", + "template_config": template_config, + "handler_used": provider_api_key(provider_api), + "method": "handler", + "provider_data": provider_data, + }, + ) + + async def execute_create_handler_async( + self, + *, + create_context: CreateOperationContext, + request: Request, + ) -> ProviderResult: + """Execute the handler acquire call through the async Azure handler contract.""" + handler_result = await create_context.handler.acquire_hosts_async( + request, + create_context.azure_template, + ) + return self.normalize_handler_create_result( + handler_result, + template_config=create_context.template_config, + provider_api=create_context.provider_api, + count=create_context.count, + template_id=create_context.azure_template.template_id, + ) diff --git a/src/orb/providers/azure/services/resource_metadata_service.py b/src/orb/providers/azure/services/resource_metadata_service.py new file mode 100644 index 000000000..ce9100a19 --- /dev/null +++ b/src/orb/providers/azure/services/resource_metadata_service.py @@ -0,0 +1,398 @@ +"""Azure resource metadata enrichment.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Optional, Protocol, TypedDict + +from orb.domain.base.ports import LoggingPort +from orb.domain.base.provider_fulfilment import ( + CheckHostsStatusResult, + FulfilmentState, + ProviderFulfilment, +) + +_RUNNING_STATES = frozenset({"running"}) +_PENDING_STATES = frozenset({"pending", "creating", "starting", "updating", "unknown"}) +_FAILED_STATES = frozenset({"failed", "terminated", "stopped", "deallocated"}) + + +@dataclass +class VmssCapacitySnapshot: + """Normalized VMSS capacity details for one scale set.""" + + target_capacity_units: int + fulfilled_capacity_units: int + provisioned_instance_count: int + state: Optional[str] + + def as_metadata(self) -> dict[str, Any]: + """Return the snapshot as a plain dict suitable for metadata output.""" + return { + "target_capacity_units": self.target_capacity_units, + "fulfilled_capacity_units": self.fulfilled_capacity_units, + "provisioned_instance_count": self.provisioned_instance_count, + "state": self.state, + } + + +class VmssCapacityInfo(TypedDict): + """VMSS capacity payload returned by Azure resource management.""" + + vmss_name: str + resource_group: str + capacity: int + vm_size: str | None + provisioning_state: str | None + provisioned_instance_count: int + + +class AzureResourceManagerProtocol(Protocol): + """Structural subset of AzureResourceManager used for metadata enrichment.""" + + async def get_vmss_capacity_async( + self, resource_group: str, vmss_name: str + ) -> VmssCapacityInfo: + """Return VMSS capacity details for one scale set via the async SDK.""" + ... + + +class AzureDeploymentStatusServiceProtocol(Protocol): + """Structural subset of AzureDeploymentService used for deployment status enrichment.""" + + async def get_deployment_status_async( + self, + *, + resource_group: str, + deployment_name: str, + ) -> Optional[dict[str, object]]: + """Return deployment provisioning/error state for one ARM deployment via the async SDK.""" + ... + + +class AzureResourceMetadataService: + """Own Azure-specific metadata enrichment for discovery/status flows.""" + + def __init__(self, default_resource_group: Optional[str], logger: LoggingPort) -> None: + self._default_resource_group = default_resource_group + self._logger = logger + + @staticmethod + def _dedupe_resource_ids(resource_ids: list[str]) -> list[str]: + deduped: list[str] = [] + for resource_id in resource_ids: + vmss_name = str(resource_id) + if vmss_name and vmss_name not in deduped: + deduped.append(vmss_name) + return deduped + + async def augment_vmss_capacity_metadata_async( + self, + metadata: dict[str, Any], + resource_ids: list[str], + *, + resource_manager: AzureResourceManagerProtocol | None, + resource_group: Optional[str] = None, + ) -> None: + """Async enrich metadata with aggregate VMSS capacity fulfilment.""" + if not resource_ids or resource_manager is None: + return + + resolved_resource_group = resource_group or self._default_resource_group + if not resolved_resource_group: + return + + per_resource_capacity = await self._collect_vmss_capacity_async( + resource_group=resolved_resource_group, + resource_ids=resource_ids, + resource_manager=resource_manager, + ) + if not per_resource_capacity: + return + + aggregate_snapshot = self._aggregate_vmss_capacity(per_resource_capacity) + metadata["fleet_capacity_fulfilment"] = aggregate_snapshot.as_metadata() + if len(per_resource_capacity) > 1: + metadata["fleet_capacity_fulfilment_by_resource"] = { + vmss_name: snapshot.as_metadata() + for vmss_name, snapshot in per_resource_capacity.items() + } + + async def _collect_vmss_capacity_async( + self, + *, + resource_group: str, + resource_ids: list[str], + resource_manager: AzureResourceManagerProtocol, + ) -> dict[str, VmssCapacitySnapshot]: + vmss_names = self._dedupe_resource_ids(resource_ids) + snapshots = await asyncio.gather( + *[ + self._get_vmss_capacity_snapshot_async( + resource_group=resource_group, + vmss_name=vmss_name, + resource_manager=resource_manager, + ) + for vmss_name in vmss_names + ] + ) + + per_resource_capacity: dict[str, VmssCapacitySnapshot] = {} + for vmss_name, snapshot in zip(vmss_names, snapshots, strict=True): + if snapshot is not None: + per_resource_capacity[vmss_name] = snapshot + return per_resource_capacity + + async def _get_vmss_capacity_snapshot_async( + self, + *, + resource_group: str, + vmss_name: str, + resource_manager: AzureResourceManagerProtocol, + ) -> Optional[VmssCapacitySnapshot]: + try: + capacity_info = await resource_manager.get_vmss_capacity_async( + resource_group, + vmss_name, + ) + except Exception as exc: + self._logger.warning( + "Could not fetch VMSS capacity for %s: %s", + vmss_name, + exc, + exc_info=True, + ) + return None + + provisioned_instance_count = capacity_info["provisioned_instance_count"] + target_capacity = capacity_info["capacity"] + provisioning_state = capacity_info.get("provisioning_state") + return VmssCapacitySnapshot( + target_capacity_units=target_capacity, + fulfilled_capacity_units=provisioned_instance_count, + provisioned_instance_count=provisioned_instance_count, + state=str(provisioning_state) if provisioning_state not in (None, "") else None, + ) + + @staticmethod + def _aggregate_vmss_capacity( + per_resource_capacity: dict[str, VmssCapacitySnapshot], + ) -> VmssCapacitySnapshot: + states = [snapshot.state for snapshot in per_resource_capacity.values() if snapshot.state] + aggregate_state = None + if len(per_resource_capacity) == 1: + aggregate_state = next(iter(per_resource_capacity.values())).state + elif states: + aggregate_state = states[0] if len(set(states)) == 1 else "multiple" + + target_capacity = sum( + snapshot.target_capacity_units for snapshot in per_resource_capacity.values() + ) + fulfilled_capacity = sum( + snapshot.fulfilled_capacity_units for snapshot in per_resource_capacity.values() + ) + return VmssCapacitySnapshot( + target_capacity_units=target_capacity, + fulfilled_capacity_units=fulfilled_capacity, + provisioned_instance_count=fulfilled_capacity, + state=aggregate_state, + ) + + async def augment_single_vm_deployment_metadata_async( + self, + metadata: dict[str, Any], + request_metadata: dict[str, Any], + *, + resource_group: Optional[str], + deployment_service: AzureDeploymentStatusServiceProtocol | None, + ) -> None: + """Async enrich metadata with ARM deployment status for single-VM resources.""" + deployment_name = request_metadata.get("deployment_name") + if deployment_name in (None, "") or not resource_group or deployment_service is None: + return + + try: + deployment_status = await deployment_service.get_deployment_status_async( + resource_group=str(resource_group), + deployment_name=str(deployment_name), + ) + except Exception as exc: + self._logger.warning( + "Could not fetch SingleVM deployment status for %s: %s", + deployment_name, + exc, + ) + return + + if not deployment_status: + return + + metadata["deployment_name"] = str(deployment_name) + provisioning_state = deployment_status.get("provisioning_state") + if provisioning_state not in (None, ""): + metadata["deployment_provisioning_state"] = provisioning_state + + error_code = deployment_status.get("error_code") + error_message = deployment_status.get("error_message") + if str(provisioning_state).lower() == "failed" and error_code in (None, ""): + error_code = "DeploymentFailed" + if error_code not in (None, "") or error_message not in (None, ""): + metadata["fleet_errors"] = [ + { + "error_code": error_code or "DeploymentFailed", + "error_message": error_message or f"ARM deployment '{deployment_name}' failed", + "resource_group": str(resource_group), + "instance_id": str(deployment_name), + } + ] + + @staticmethod + def augment_shortfall_metadata(metadata: dict[str, Any]) -> None: + """Add a capacity_shortfall summary when fulfilled capacity is below the target.""" + capacity = metadata.get("fleet_capacity_fulfilment") or {} + target = capacity.get("target_capacity_units") + fulfilled = capacity.get("fulfilled_capacity_units") + fleet_errors = metadata.get("fleet_errors") or [] + + if target is None or fulfilled is None: + return + if fulfilled >= target and not fleet_errors: + return + + missing_capacity = max(int(target) - int(fulfilled), 0) + likely_causes: list[str] = [] + seen_causes: set[str] = set() + + for error in fleet_errors: + error_code = str((error or {}).get("error_code") or "") + cause = error_code or "Unknown" + if cause not in seen_causes: + likely_causes.append(cause) + seen_causes.add(cause) + + metadata["capacity_shortfall"] = { + "missing_capacity_units": missing_capacity, + "likely_causes": likely_causes, + "summary": ( + f"Shortfall {fulfilled}/{target}" + + (f"; causes={', '.join(likely_causes)}" if likely_causes else "") + ), + } + + def attach_provider_fulfilment( + self, + metadata: dict[str, Any], + *, + instances: Sequence[Mapping[str, Any]], + target_units: int | None, + ) -> CheckHostsStatusResult: + """Attach and return the canonical provider status contract.""" + status_result = CheckHostsStatusResult( + instances=[dict(instance) for instance in instances], + fulfilment=_build_provider_fulfilment( + metadata=metadata, + instances=instances, + target_units=target_units, + ), + ) + metadata["provider_fulfilment"] = status_result.fulfilment + return status_result + + +def _build_provider_fulfilment( + *, + metadata: dict[str, Any], + instances: Sequence[Mapping[str, Any]], + target_units: int | None, +) -> ProviderFulfilment: + fleet_errors = metadata.get("fleet_errors") or [] + # ProviderResult.metadata is the core provider result bag typed as dict[str, Any]. + # Azure owns this key, but reading it back from the core metadata boundary + # requires validating the shape before trusting the capacity fields. + capacity = metadata.get("fleet_capacity_fulfilment") + if isinstance(capacity, dict): + target = capacity.get("target_capacity_units") + fulfilled = capacity.get("fulfilled_capacity_units") + if isinstance(target, int) and isinstance(fulfilled, int): + capacity_state = str(capacity.get("state") or "").lower() + if fulfilled >= target and target > 0 and not fleet_errors: + return ProviderFulfilment( + state="fulfilled", + message=f"Azure capacity fulfilled: {fulfilled}/{target}", + target_units=target, + fulfilled_units=fulfilled, + ) + + if capacity_state == "failed" or fleet_errors: + capacity_fulfilment_state: FulfilmentState = ( + "partial" if fulfilled > 0 else "failed" + ) + return ProviderFulfilment( + state=capacity_fulfilment_state, + message=f"Azure capacity shortfall: {fulfilled}/{target}", + target_units=target, + fulfilled_units=fulfilled, + ) + + return ProviderFulfilment( + state="in_progress", + message=f"Azure capacity provisioning: {fulfilled}/{target}", + target_units=target, + fulfilled_units=fulfilled, + ) + + # Handler-backed status, SingleVM, CycleCloud, and dry-run results do not + # have VMSS capacity metadata; derive their verdict from observed statuses. + target = target_units + running_count = _count_statuses(instances, _RUNNING_STATES) + pending_count = _count_statuses(instances, _PENDING_STATES) + failed_count = _count_statuses(instances, _FAILED_STATES) + + if fleet_errors and running_count == 0: + return ProviderFulfilment( + state="failed", + message="Azure provisioning failed before capacity became available", + target_units=target, + fulfilled_units=running_count, + ) + + if target is not None and target > 0 and running_count >= target and failed_count == 0: + return ProviderFulfilment( + state="fulfilled", + message=f"Azure instances running: {running_count}/{target}", + target_units=target, + fulfilled_units=running_count, + ) + + if failed_count > 0 and pending_count == 0: + state: FulfilmentState = "partial" if running_count > 0 else "failed" + return ProviderFulfilment( + state=state, + message=( + f"Azure instance shortfall: {running_count}/{target}" + if target is not None + else f"Azure instance failure: {running_count} running" + ), + target_units=target, + fulfilled_units=running_count, + ) + + return ProviderFulfilment( + state="in_progress", + message=( + f"Azure instances provisioning: {running_count}/{target}" + if target is not None + else f"Azure instances provisioning: {running_count} running" + ), + target_units=target, + fulfilled_units=running_count, + ) + + +def _count_statuses( + instances: Sequence[Mapping[str, Any]], + statuses: frozenset[str], +) -> int: + return sum(1 for instance in instances if str(instance.get("status") or "").lower() in statuses) diff --git a/src/orb/providers/azure/services/runtime_dependencies.py b/src/orb/providers/azure/services/runtime_dependencies.py new file mode 100644 index 000000000..a02e0a25f --- /dev/null +++ b/src/orb/providers/azure/services/runtime_dependencies.py @@ -0,0 +1,172 @@ +"""Azure runtime-owned lazy dependency resolution.""" + +from __future__ import annotations + +from threading import RLock +from typing import TYPE_CHECKING, Callable, Optional + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.configuration.config import AzureProviderConfig + +if TYPE_CHECKING: + from orb.providers.azure.infrastructure.azure_client import AzureClient + from orb.providers.azure.infrastructure.azure_handler_factory import AzureHandlerFactory + from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, + ) + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + +class AzureRuntimeDependencies: + """Own Azure lazy runtime dependency resolution and cache lifecycle.""" + + def __init__( + self, + *, + config: AzureProviderConfig, + logger: LoggingPort, + azure_client_resolver: Optional[Callable[[], AzureClient]] = None, + azure_handler_factory_resolver: Optional[Callable[[], AzureHandlerFactory]] = None, + azure_resource_manager_resolver: Optional[Callable[[], AzureResourceManager | None]] = None, + azure_deployment_service_resolver: Optional[ + Callable[[], AzureDeploymentService | None] + ] = None, + azure_native_spec_service: Optional[AzureNativeSpecService] = None, + ) -> None: + """Initialize the Azure runtime dependency holder and its lazy resolvers.""" + self._config = config + self._logger = logger + self._azure_client_resolver = azure_client_resolver + self._azure_handler_factory_resolver = azure_handler_factory_resolver + self._azure_resource_manager_resolver = azure_resource_manager_resolver + self._azure_deployment_service_resolver = azure_deployment_service_resolver + self._azure_native_spec_service = azure_native_spec_service + self._lock = RLock() + self._client: Optional[AzureClient] = None + self._resource_manager: Optional[AzureResourceManager] = None + self._deployment_service: Optional[AzureDeploymentService] = None + self._handler_factory: Optional[AzureHandlerFactory] = None + + @property + def azure_client(self) -> Optional[AzureClient]: + """Resolve and cache AzureClient on first access.""" + with self._lock: + if self._client is None: + self._logger.debug("Creating Azure client on first access") + if self._azure_client_resolver: + try: + self._client = self._azure_client_resolver() + except Exception as exc: + self._logger.warning( + "Failed to resolve AzureClient lazily: %s", + exc, + exc_info=True, + ) + self._client = None + else: + self._logger.warning("AzureClient resolver not provided") + return self._client + + @property + def resource_manager(self) -> Optional[AzureResourceManager]: + """Resolve and cache AzureResourceManager on first access.""" + with self._lock: + azure_client = self.azure_client + if self._resource_manager is None and self._azure_resource_manager_resolver is not None: + try: + self._resource_manager = self._azure_resource_manager_resolver() + except Exception as exc: + self._logger.warning( + "Failed to resolve AzureResourceManager lazily: %s", + exc, + exc_info=True, + ) + self._resource_manager = None + if self._resource_manager is None and azure_client: + self._logger.debug("Creating Azure resource manager on first access") + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + self._resource_manager = AzureResourceManager( + azure_client=azure_client, + config=self._config, + logger=self._logger, + ) + return self._resource_manager + + @property + def deployment_service(self) -> Optional[AzureDeploymentService]: + """Resolve and cache AzureDeploymentService on first access.""" + with self._lock: + azure_client = self.azure_client + if ( + self._deployment_service is None + and self._azure_deployment_service_resolver is not None + ): + try: + self._deployment_service = self._azure_deployment_service_resolver() + except Exception as exc: + self._logger.warning( + "Failed to resolve AzureDeploymentService lazily: %s", + exc, + exc_info=True, + ) + self._deployment_service = None + if self._deployment_service is None and azure_client: + from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, + ) + + self._deployment_service = AzureDeploymentService( + azure_client=azure_client, + logger=self._logger, + ) + return self._deployment_service + + @property + def handler_factory(self) -> Optional[AzureHandlerFactory]: + """Resolve and cache AzureHandlerFactory on first access.""" + with self._lock: + if self._handler_factory is not None: + return self._handler_factory + + if self._azure_handler_factory_resolver is not None: + try: + self._handler_factory = self._azure_handler_factory_resolver() + return self._handler_factory + except Exception as exc: + self._logger.warning( + "Failed to resolve AzureHandlerFactory lazily: %s", + exc, + exc_info=True, + ) + self._handler_factory = None + return None + + azure_client = self.azure_client + if azure_client is None: + return None + + from orb.providers.azure.infrastructure.azure_handler_factory import ( + AzureHandlerFactory, + ) + + self._handler_factory = AzureHandlerFactory( + azure_client=azure_client, + logger=self._logger, + azure_native_spec_service=self._azure_native_spec_service, + azure_resource_manager=self.resource_manager, + ) + return self._handler_factory + + def clear_cached_runtime(self) -> Optional[AzureClient]: + """Clear all cached runtime dependencies and return the cached client.""" + with self._lock: + client = self._client + self._client = None + self._resource_manager = None + self._deployment_service = None + self._handler_factory = None + return client diff --git a/src/orb/providers/azure/services/spot_launch_service.py b/src/orb/providers/azure/services/spot_launch_service.py new file mode 100644 index 000000000..1064e4f59 --- /dev/null +++ b/src/orb/providers/azure/services/spot_launch_service.py @@ -0,0 +1,435 @@ +"""Azure spot-placement planning and execution helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any, Awaitable, Callable, Mapping, Optional, Protocol + +from orb.domain.base.ports import LoggingPort +from orb.domain.request.aggregate import Request +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import ( + AzureLocationName, + AzureProviderApi, +) +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.infrastructure.handlers.azure_handler import AzureHandler +from orb.providers.azure.infrastructure.services.spot_placement_score_adapter import ( + AzureSpotPlacementScoreAdapter, +) +from orb.providers.azure.services.spot_placement_execution import ( + SpotPlacementExecutionSummary, + build_planned_execution_metadata, + create_acquire_request, +) +from orb.providers.azure.services.spot_placement_planner import ( + PlacementCandidate, + PlacementPlanEntry, + PlacementScore, + SpotPlacementPlanner, +) +from orb.providers.base.strategy import ProviderOperation, ProviderResult + + +class SpotPlacementExecutionPort(Protocol): + """Structural subset of SpotPlacementExecutionService used by Azure spot launches.""" + + async def execute_plan_async( + self, + plan: list[PlacementPlanEntry], + total_count: int, + build_child_template: Callable[[PlacementPlanEntry], AzureTemplate], + build_child_request: Callable[[int, int], Request], + launch_child: Callable[[Request, AzureTemplate], Awaitable[Mapping[str, Any]]], + is_capacity_like_failure: Callable[[dict[str, Any]], bool], + ) -> SpotPlacementExecutionSummary: + """Execute a placement plan asynchronously and return the aggregated summary.""" + ... + + +@dataclass +class AzureSpotPlacementTemplateView: + """Template view matching the adapter's structural spot-placement input.""" + + vm_size: str + location: AzureLocationName + placement_regions: list[str] + placement_zones: list[str] + zones: list[str] + candidate_vm_sizes: list[str] + + +class AzureSpotLaunchService: + """Own Azure spot-placement planning and planned launch execution.""" + + def __init__( + self, + *, + config: AzureProviderConfig, + logger: LoggingPort, + planner: SpotPlacementPlanner, + execution_service: SpotPlacementExecutionPort, + ) -> None: + self._config = config + self._logger = logger + self._planner = planner + self._execution_service = execution_service + + @staticmethod + def should_use_spot_placement(template: AzureTemplate) -> bool: + """Return whether the template opts into Spot Placement Score planning.""" + return template.spot_placement_score_enabled + + def build_spot_placement_plan( + self, + *, + azure_template: AzureTemplate, + count: int, + azure_client: AzureClient | None, + ) -> list[PlacementPlanEntry]: + """Score candidate regions/zones and build a placement plan for spot launches.""" + template_view = self._template_view(azure_template) + if azure_client is None: + return self._fallback_plan_for_missing_client(template_view, count) + + adapter = AzureSpotPlacementScoreAdapter( + azure_client=azure_client, + logger=self._logger, + subscription_id=azure_template.subscription_id or self._config.subscription_id, + base_location=azure_template.location.value or self._config.region, + ) + scores = adapter.score_candidates(requested_count=count, template=template_view) + return self._plan_from_scores(azure_template=azure_template, count=count, scores=scores) + + async def build_spot_placement_plan_async( + self, + *, + azure_template: AzureTemplate, + count: int, + azure_client: AzureClient | None, + ) -> list[PlacementPlanEntry]: + """Score candidate regions/zones without blocking the async create flow.""" + template_view = self._template_view(azure_template) + if azure_client is None: + return self._fallback_plan_for_missing_client(template_view, count) + + adapter = AzureSpotPlacementScoreAdapter( + azure_client=azure_client, + logger=self._logger, + subscription_id=azure_template.subscription_id or self._config.subscription_id, + base_location=azure_template.location.value or self._config.region, + ) + scores = await adapter.score_candidates_async(requested_count=count, template=template_view) + return self._plan_from_scores(azure_template=azure_template, count=count, scores=scores) + + @staticmethod + def _template_view(azure_template: AzureTemplate) -> AzureSpotPlacementTemplateView: + """Build the structural template view used by the scoring adapter.""" + return AzureSpotPlacementTemplateView( + vm_size=azure_template.vm_size, + location=azure_template.location, + placement_regions=list(azure_template.placement_regions or []), + placement_zones=list(azure_template.placement_zones or []), + zones=list(azure_template.zones or []), + candidate_vm_sizes=list(azure_template.candidate_vm_sizes), + ) + + def _fallback_plan_for_missing_client( + self, + template_view: AzureSpotPlacementTemplateView, + count: int, + ) -> list[PlacementPlanEntry]: + """Build a deterministic fallback plan when live scoring cannot run.""" + self._logger.warning("Azure client not available; falling back to template candidate order") + return self.build_fallback_spot_placement_plan( + self._build_approximate_template_scores(template_view), + count, + ) + + def _plan_from_scores( + self, + *, + azure_template: AzureTemplate, + count: int, + scores: list[PlacementScore], + ) -> list[PlacementPlanEntry]: + """Create a placement plan from scores with a deterministic fallback.""" + plan = self._planner.create_plan( + requested_count=count, + scores=scores, + split_strategy=azure_template.placement_split_strategy, + primary_share_percent=azure_template.placement_primary_share_percent, + ) + if plan: + return plan + if scores: + self._logger.warning( + "Azure spot placement scores returned no viable candidates; " + "falling back to template candidate order" + ) + return self.build_fallback_spot_placement_plan(scores, count) + return [] + + @staticmethod + def _build_approximate_template_scores( + template: AzureSpotPlacementTemplateView, + ) -> list[PlacementScore]: + """Build approximate scores in template order when live scoring cannot run.""" + regions = template.placement_regions or [template.location.value] + zones = template.placement_zones or template.zones or [None] + + return [ + PlacementScore( + candidate=PlacementCandidate( + candidate_id=f"azure:{region}:{zone or 'regional'}:{vm_size}", + instance_type=vm_size, + region=region, + zone=zone, + ), + raw_score="DataNotFoundOrStale", + normalized_score=0.0, + approximate=True, + metadata={ + "fallback_reason": "azure_client_unavailable", + "raw_entry": {"score": "DataNotFoundOrStale"}, + }, + ) + for region in regions + for vm_size in template.candidate_vm_sizes + for zone in zones + ] + + @staticmethod + def build_fallback_spot_placement_plan( + scores: list[PlacementScore], + requested_count: int, + ) -> list[PlacementPlanEntry]: + """Build a single-candidate fallback plan when the planner returns no viable entries.""" + if requested_count <= 0 or not scores: + return [] + + fallback_scores = [ + replace( + score, + approximate=True, + metadata={ + **score.metadata, + "fallback_reason": "no_viable_provider_scores", + }, + ) + for score in scores + ] + return [ + PlacementPlanEntry(score=fallback_scores[0], planned_count=requested_count), + *[PlacementPlanEntry(score=score, planned_count=0) for score in fallback_scores[1:]], + ] + + @staticmethod + def is_capacity_like_failure(child_result: dict[str, Any]) -> bool: + """Return whether the child result error codes indicate a capacity-related failure.""" + error_codes = set(child_result.get("error_codes", [])) + return bool( + error_codes + & { + "AllocationFailed", + "ZonalAllocationFailed", + "SkuNotAvailable", + "OverconstrainedAllocationRequest", + } + ) + + @staticmethod + def clone_template_for_plan_entry( + azure_template: AzureTemplate, + plan_entry: PlacementPlanEntry, + ) -> AzureTemplate: + """Clone a template with the VM size, location, and zone from a plan entry.""" + cloned_data = azure_template.model_dump(mode="json", exclude_none=True) + selected_vm_size = plan_entry.score.candidate.instance_type + cloned_data["vm_size"] = selected_vm_size + cloned_data["vm_sizes"] = [] + cloned_data["vm_size_preferences"] = [] + cloned_data.pop("allocation_strategy", None) + cloned_data["spot_placement_score_enabled"] = False + cloned_data["location"] = plan_entry.score.candidate.region or azure_template.location.value + cloned_data["zones"] = ( + [plan_entry.score.candidate.zone] if plan_entry.score.candidate.zone else [] + ) + cloned_data["placement_regions"] = [] + cloned_data["placement_zones"] = [] + return AzureTemplate.model_validate(cloned_data) + + @staticmethod + def _planned_child_result_with_fulfillment( + *, + provider_api_key: str, + requested_count: int, + raw_result: Mapping[str, Any], + ) -> dict[str, Any]: + result = dict(raw_result) + if "fulfilled_count" in result: + return result + + if not result.get("success", True): + result["fulfilled_count"] = 0 + return result + + if provider_api_key == "CycleCloud": + provider_data = result.get("provider_data") + added_count = ( + provider_data.get("added_count") if isinstance(provider_data, Mapping) else None + ) + result["fulfilled_count"] = int(added_count or 0) + return result + + result["fulfilled_count"] = requested_count + return result + + async def execute_planned_spot_launches_async( + self, + *, + azure_template: AzureTemplate, + provider_api: AzureProviderApi, + provider_api_key: str, + count: int, + template_config: dict[str, Any], + operation: ProviderOperation, + provider_instance_name: str, + handler: Optional[AzureHandler], + azure_client: AzureClient | None, + plan_override: Optional[list[PlacementPlanEntry]] = None, + capacity_like_failure_checker: Optional[Callable[[dict[str, Any]], bool]] = None, + ) -> ProviderResult: + """Async variant of planned spot launches using the async handler contract.""" + if not handler: + return ProviderResult.error_result( + f"No handler available for provider_api: {provider_api_key}", + "HANDLER_NOT_FOUND", + ) + + plan = ( + plan_override + if plan_override is not None + else await self.build_spot_placement_plan_async( + azure_template=azure_template, + count=count, + azure_client=azure_client, + ) + ) + if not plan: + return ProviderResult.error_result( + "No viable spot placement candidates returned scores", + "NO_PLACEMENT_CANDIDATES", + ) + + request_metadata = dict(operation.parameters.get("request_metadata", {}) or {}) + base_request_id = operation.parameters.get("request_id") or ( + operation.context.get("request_id") if operation.context else None + ) + + summary = await self._execution_service.execute_plan_async( + plan=plan, + total_count=count, + build_child_template=lambda plan_entry: self.clone_template_for_plan_entry( + azure_template, plan_entry + ), + build_child_request=lambda requested_for_entry, idx: create_acquire_request( + template_id=azure_template.template_id, + count=requested_for_entry, + provider_type="azure", + provider_name=provider_instance_name, + provider_api=provider_api_key, + request_metadata=request_metadata, + parent_request_id=base_request_id, + plan_entry_index=idx, + ), + launch_child=lambda child_request, child_template: self._launch_planned_child_async( + handler=handler, + provider_api_key=provider_api_key, + child_request=child_request, + child_template=child_template, + ), + is_capacity_like_failure=capacity_like_failure_checker or self.is_capacity_like_failure, + ) + return self._planned_execution_result( + plan=plan, + summary=summary, + count=count, + template_config=template_config, + provider_api_key=provider_api_key, + template_id=azure_template.template_id, + ) + + async def _launch_planned_child_async( + self, + *, + handler: AzureHandler, + provider_api_key: str, + child_request: Request, + child_template: AzureTemplate, + ) -> dict[str, Any]: + """Launch one planned child request through the async Azure handler contract.""" + return self._planned_child_result_with_fulfillment( + provider_api_key=provider_api_key, + requested_count=child_request.requested_count, + raw_result=await handler.acquire_hosts_async(child_request, child_template), + ) + + def _planned_execution_result( + self, + *, + plan: list[PlacementPlanEntry], + summary: SpotPlacementExecutionSummary, + count: int, + template_config: dict[str, Any], + provider_api_key: str, + template_id: str, + ) -> ProviderResult: + """Build the final ProviderResult for planned spot execution.""" + provider_data = build_planned_execution_metadata(plan, summary) + provider_data["fulfillment_final"] = ( + bool(summary.resource_ids) + and summary.unfulfilled_count == 0 + and not summary.terminal_error_message + and not summary.terminated_early + ) + + if ( + not summary.resource_ids + and not summary.instances + and (summary.terminal_error_message or summary.unfulfilled_count > 0) + ): + error_message = ( + f"Provisioning failed: {summary.terminal_error_message}" + if summary.terminal_error_message + else "Spot placement plan could not provision any instances" + ) + return ProviderResult.error_result( + error_message, + "PROVISIONING_ADAPTER_ERROR", + { + "operation": "create_instances", + "template_config": template_config, + "handler_used": provider_api_key, + "method": "planned_handler", + "provider_data": provider_data, + }, + ) + + return ProviderResult.success_result( + { + "resource_ids": summary.resource_ids, + "instances": summary.instances, + "provider_api": provider_api_key, + "count": count, + "template_id": template_id, + }, + { + "operation": "create_instances", + "template_config": template_config, + "handler_used": provider_api_key, + "method": "planned_handler", + "provider_data": provider_data, + }, + ) diff --git a/src/orb/providers/azure/services/spot_placement_execution.py b/src/orb/providers/azure/services/spot_placement_execution.py new file mode 100644 index 000000000..605967605 --- /dev/null +++ b/src/orb/providers/azure/services/spot_placement_execution.py @@ -0,0 +1,249 @@ +"""Azure execution helpers for spot placement plans.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from typing import Any + +from orb.domain.base.exceptions import DomainException +from orb.providers.azure.services.spot_placement_planner import PlacementPlanEntry + + +@dataclass(frozen=True) +class SpotPlacementExecutionSummary: + """Aggregated execution outcome for a placement plan.""" + + resource_ids: list[str] + instances: list[dict[str, Any]] + child_results: list[dict[str, Any]] + failed_subplans: list[dict[str, Any]] + unfulfilled_count: int + terminated_early: bool = False + terminal_error_message: str | None = None + + @property + def provider_data(self) -> dict[str, Any]: + """Return JSON-friendly execution metadata for request/provider data.""" + return { + "child_results": self.child_results, + "failed_subplans": self.failed_subplans, + "unfulfilled_count": self.unfulfilled_count, + "terminated_early": self.terminated_early, + "terminal_error_message": self.terminal_error_message, + } + + +def serialize_placement_plan(plan: list[PlacementPlanEntry]) -> list[dict[str, Any]]: + """Convert plan entries into JSON-friendly metadata.""" + return [ + { + "candidate_id": entry.score.candidate.candidate_id, + "instance_type": entry.score.candidate.instance_type, + "region": entry.score.candidate.region, + "zone": entry.score.candidate.zone, + "raw_score": entry.score.raw_score, + "normalized_score": entry.score.normalized_score, + "approximate": entry.score.approximate, + "planned_count": entry.planned_count, + "metadata": entry.score.metadata, + } + for entry in plan + ] + + +def create_acquire_request( + template_id: str, + count: int, + provider_type: str, + provider_name: str | None, + provider_api: str, + request_metadata: dict[str, Any], + request_id: str | None = None, + parent_request_id: str | None = None, + plan_entry_index: int | None = None, +) -> Any: + """Create a standard acquire request for a child placement launch.""" + from orb.domain.request.aggregate import Request + from orb.domain.request.request_types import RequestType + + child_metadata = dict(request_metadata) + if parent_request_id: + child_metadata["parent_request_id"] = parent_request_id + if plan_entry_index is not None: + child_metadata["spot_placement_plan_entry_index"] = plan_entry_index + + request = Request.create_new_request( + request_type=RequestType.ACQUIRE, + template_id=template_id, + machine_count=count, + provider_type=provider_type, + provider_name=provider_name, + metadata=child_metadata, + request_id=request_id, + ) + request.provider_api = provider_api + return request + + +def build_planned_execution_metadata( + plan: list[PlacementPlanEntry], + summary: SpotPlacementExecutionSummary, +) -> dict[str, Any]: + """Build metadata payload for planned execution results.""" + return { + "placement_plan": serialize_placement_plan(plan), + **summary.provider_data, + } + + +class SpotPlacementExecutionService: + """Execute a placement plan via Azure callbacks.""" + + async def execute_plan_async( + self, + plan: list[PlacementPlanEntry], + total_count: int, + build_child_template: Callable[[PlacementPlanEntry], Any], + build_child_request: Callable[[int, int], Any], + launch_child: Callable[[Any, Any], Awaitable[Mapping[str, Any]]], + is_capacity_like_failure: Callable[[dict[str, Any]], bool], + ) -> SpotPlacementExecutionSummary: + resource_ids: list[str] = [] + instances: list[dict[str, Any]] = [] + child_results: list[dict[str, Any]] = [] + failed_subplans: list[dict[str, Any]] = [] + carryover = 0 + fulfilled = 0 + terminated_early = False + terminal_error_message: str | None = None + + for idx, plan_entry in enumerate(plan): + requested_for_entry = plan_entry.planned_count + carryover + if requested_for_entry <= 0: + continue + + child_template = build_child_template(plan_entry) + child_request = build_child_request(requested_for_entry, idx) + try: + raw_result = await launch_child(child_request, child_template) + except DomainException as exc: + raw_result = self._domain_exception_child_result(exc) + + child_result = self._normalize_child_result( + plan_entry=plan_entry, + requested_count=requested_for_entry, + raw_result=raw_result, + ) + child_results.append(child_result) + + if child_result["success"]: + resource_ids.extend(child_result["resource_ids"]) + instances.extend(child_result["instances"]) + fulfilled_for_child = child_result["fulfilled_count"] + fulfilled += fulfilled_for_child + carryover = max(requested_for_entry - fulfilled_for_child, 0) + if carryover == 0: + continue + + failed_subplans.append(child_result) + if is_capacity_like_failure(child_result): + continue + + terminated_early = True + terminal_error_message = child_result["error_message"] + break + + failed_subplans.append(child_result) + if is_capacity_like_failure(child_result): + carryover = requested_for_entry + continue + + if resource_ids or instances: + terminated_early = True + terminal_error_message = child_result["error_message"] + break + + return SpotPlacementExecutionSummary( + resource_ids=[], + instances=[], + child_results=child_results, + failed_subplans=failed_subplans, + unfulfilled_count=total_count, + terminated_early=True, + terminal_error_message=child_result["error_message"], + ) + + return SpotPlacementExecutionSummary( + resource_ids=resource_ids, + instances=instances, + child_results=child_results, + failed_subplans=failed_subplans, + unfulfilled_count=max(total_count - fulfilled, carryover), + terminated_early=terminated_early, + terminal_error_message=terminal_error_message, + ) + + @staticmethod + def _domain_exception_child_result(exc: DomainException) -> dict[str, Any]: + """Normalize a domain exception into the child-launch result contract.""" + return { + "success": False, + "resource_ids": [], + "instances": [], + "error_message": SpotPlacementExecutionService._format_launch_error(exc), + "provider_data": { + "error_codes": [exc.error_code] if exc.error_code else [], + }, + } + + @staticmethod + def _format_launch_error(exc: DomainException) -> str: + error_code = exc.error_code + message = str(exc) + if error_code and not message.startswith(f"{error_code}:"): + return f"{error_code}: {message}" + return message + + @staticmethod + def _normalize_child_result( + plan_entry: PlacementPlanEntry, + requested_count: int, + raw_result: Mapping[str, Any], + ) -> dict[str, Any]: + result = dict(raw_result) + success = result.get("success", True) + error_message = result.get("error_message") + resource_ids = result.get("resource_ids", []) + instances = result.get("instances", []) + provider_data = result.get("provider_data") or {} + raw_fulfilled_count = result.get("fulfilled_count") + if success: + fulfilled_count = ( + requested_count if raw_fulfilled_count is None else int(raw_fulfilled_count) + ) + else: + fulfilled_count = 0 if raw_fulfilled_count is None else int(raw_fulfilled_count) + fulfilled_count = min(max(fulfilled_count, 0), requested_count) + + return { + "candidate_id": plan_entry.score.candidate.candidate_id, + "requested_count": requested_count, + "fulfilled_count": fulfilled_count, + "success": success, + "resource_ids": resource_ids, + "instances": instances, + "error_message": error_message, + "error_codes": SpotPlacementExecutionService._extract_error_codes(provider_data), + "provider_data": provider_data, + } + + @staticmethod + def _extract_error_codes(provider_data: dict[str, Any]) -> list[str]: + error_codes: list[str] = [] + direct_error_codes = provider_data.get("error_codes", []) + if isinstance(direct_error_codes, (list, tuple, set)): + error_codes.extend(str(error_code) for error_code in direct_error_codes if error_code) + elif direct_error_codes: + error_codes.append(str(direct_error_codes)) + return list(dict.fromkeys(error_codes)) diff --git a/src/orb/providers/azure/services/spot_placement_planner.py b/src/orb/providers/azure/services/spot_placement_planner.py new file mode 100644 index 000000000..00e11e5cb --- /dev/null +++ b/src/orb/providers/azure/services/spot_placement_planner.py @@ -0,0 +1,173 @@ +"""Azure spot placement score planning service.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from math import ceil +from typing import Any, Protocol + + +class PlacementSplitStrategy(str, Enum): + """How to split requested capacity across spot placement candidates.""" + + GREEDY = "greedy" + HYBRID = "hybrid" + + +@dataclass(frozen=True) +class PlacementCandidate: + """Azure spot placement candidate.""" + + candidate_id: str + instance_type: str + region: str | None = None + zone: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PlacementScore: + """Placement score for a candidate.""" + + candidate: PlacementCandidate + raw_score: Any + normalized_score: float + approximate: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PlacementPlanEntry: + """Concrete count allocation for a placement candidate.""" + + score: PlacementScore + planned_count: int + + +class SpotPlacementScoreAdapter(Protocol): + """Provider-specific scoring adapter contract.""" + + def score_candidates(self, requested_count: int, template: Any) -> list[PlacementScore]: + """Return normalized placement scores for the provider's candidates.""" + ... + + +class SpotPlacementPlanner: + """Build a placement plan from normalized candidate scores.""" + + def create_plan( + self, + requested_count: int, + scores: list[PlacementScore], + split_strategy: PlacementSplitStrategy, + primary_share_percent: int, + ) -> list[PlacementPlanEntry]: + """Create an ordered placement plan.""" + if requested_count <= 0: + return [] + + ranked_scores = self._rank_scores(scores) + if not ranked_scores: + return [] + + if split_strategy == PlacementSplitStrategy.GREEDY or len(ranked_scores) == 1: + return self._create_greedy_plan(requested_count, ranked_scores) + + return self._create_hybrid_plan( + requested_count=requested_count, + scores=ranked_scores, + primary_share_percent=primary_share_percent, + ) + + def _rank_scores(self, scores: list[PlacementScore]) -> list[PlacementScore]: + ranked = [score for score in scores if score.normalized_score > 0] + ranked.sort( + key=lambda score: ( + score.normalized_score, + self._raw_score_sort_value(score.raw_score), + score.candidate.candidate_id, + ), + reverse=True, + ) + return ranked + + @staticmethod + def _raw_score_sort_value(raw_score: Any) -> float: + if isinstance(raw_score, (int, float)): + return float(raw_score) + if isinstance(raw_score, str): + mapping = {"low": 1.0, "medium": 2.0, "high": 3.0} + return mapping.get(raw_score.lower(), 0.0) + return 0.0 + + @staticmethod + def _create_greedy_plan( + requested_count: int, + scores: list[PlacementScore], + ) -> list[PlacementPlanEntry]: + return [ + PlacementPlanEntry(score=scores[0], planned_count=requested_count), + *[PlacementPlanEntry(score=score, planned_count=0) for score in scores[1:]], + ] + + def _create_hybrid_plan( + self, + requested_count: int, + scores: list[PlacementScore], + primary_share_percent: int, + ) -> list[PlacementPlanEntry]: + top_score = scores[0] + remainder_scores = scores[1:] + + top_count = ceil(requested_count * primary_share_percent / 100) + top_count = min(top_count, requested_count) + remainder_count = requested_count - top_count + + plan_entries = [PlacementPlanEntry(score=top_score, planned_count=top_count)] + if remainder_count <= 0 or not remainder_scores: + return plan_entries + [ + PlacementPlanEntry(score=score, planned_count=0) for score in remainder_scores + ] + + total_weight = sum(score.normalized_score for score in remainder_scores) + if total_weight <= 0: + return self._create_greedy_plan(requested_count, scores) + + raw_allocations: list[tuple[PlacementScore, int, float]] = [] + assigned = 0 + for score in remainder_scores: + exact_count = remainder_count * (score.normalized_score / total_weight) + allocated = int(exact_count) + raw_allocations.append((score, allocated, exact_count - allocated)) + assigned += allocated + + leftovers = remainder_count - assigned + if leftovers > 0: + raw_allocations.sort( + key=lambda item: ( + item[2], + item[0].normalized_score, + self._raw_score_sort_value(item[0].raw_score), + item[0].candidate.candidate_id, + ), + reverse=True, + ) + for idx in range(leftovers): + score, allocated, fraction = raw_allocations[idx] + raw_allocations[idx] = (score, allocated + 1, fraction) + + raw_allocations.sort( + key=lambda item: ( + item[0].normalized_score, + self._raw_score_sort_value(item[0].raw_score), + item[0].candidate.candidate_id, + ), + reverse=True, + ) + + plan_entries.extend( + PlacementPlanEntry(score=score, planned_count=allocated) + for score, allocated, _ in raw_allocations + ) + return plan_entries diff --git a/src/orb/providers/azure/services/template_catalog_service.py b/src/orb/providers/azure/services/template_catalog_service.py new file mode 100644 index 000000000..9b1e62121 --- /dev/null +++ b/src/orb/providers/azure/services/template_catalog_service.py @@ -0,0 +1,109 @@ +"""Azure template catalog loading.""" + +from __future__ import annotations + +from typing import Any, Protocol, cast + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.domain.template.value_objects import AzureProviderApi + + +class SchedulerTemplateStrategy(Protocol): + """Structural subset used from the active scheduler strategy.""" + + def get_template_paths(self) -> list[str]: + """Return configured template search paths.""" + ... + + def load_templates_from_path(self, path: str) -> list[dict[str, Any]]: + """Load templates from one scheduler-managed path.""" + ... + + +class AzureTemplateCatalogService: + """Load Azure templates from the active scheduler or a local fallback catalog.""" + + def __init__(self, logger: LoggingPort) -> None: + self._logger = logger + + def get_available_templates(self) -> list[dict[str, Any]]: + """Load templates from the active scheduler, falling back to built-in defaults.""" + try: + from orb.infrastructure.scheduler.registry import get_scheduler_registry + + scheduler_registry = get_scheduler_registry() + # TODO: SchedulerRegistry.get_active_strategy() does not exist; this + # call raises AttributeError at runtime, which the outer try/except + # below swallows so we always fall back to hardcoded templates. + # Either implement get_active_strategy() on SchedulerRegistry (the + # registry advertises SINGLE_CHOICE mode but exposes no accessor for + # the active type's strategy instance) or delete this code path and + # call get_fallback_templates() directly. Same bug at + # aws/services/template_validation_service.py:34. See + # provider-quality-comparison.md defect #54. + scheduler_strategy = cast( + SchedulerTemplateStrategy | None, + scheduler_registry.get_active_strategy(), # type: ignore[attr-defined] + ) + + if scheduler_strategy: + templates: list[dict[str, Any]] = [] + for path in scheduler_strategy.get_template_paths(): + try: + templates.extend(scheduler_strategy.load_templates_from_path(path)) + except Exception as exc: + self._logger.warning( + "Failed to load templates from %s: %s", path, exc, exc_info=True + ) + return templates + + self._logger.warning("No scheduler strategy available, using fallback templates") + return self.get_fallback_templates() + except Exception as exc: + self._logger.error("Failed to load templates via scheduler strategy: %s", exc) + return self.get_fallback_templates() + + @staticmethod + def get_fallback_templates() -> list[dict[str, Any]]: + """Return hard-coded sample templates used when no scheduler is available.""" + return [ + { + "template_id": "azure-vmss-linux-basic", + "name": "Azure VMSS Linux Basic", + "description": "VMSS with Ubuntu 22.04 LTS on Standard_D4s_v5", + "provider_type": "azure", + "provider_api": AzureProviderApi.VMSS.value, + "vm_size": "Standard_D4s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "ssh_key_name": "my-azure-ssh-key", + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "max_instances": 2, + }, + { + "template_id": "azure-vmss-spot", + "name": "Azure VMSS Spot Instances", + "description": "VMSS with Spot VMs for cost-effective workloads", + "provider_type": "azure", + "provider_api": AzureProviderApi.VMSS.value, + "vm_size": "Standard_D4s_v5", + "resource_group": "my-resource-group", + "location": "eastus2", + "ssh_key_name": "my-azure-ssh-key", + "priority": "Spot", + "eviction_policy": "Deallocate", + "billing_profile_max_price": -1.0, + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "max_instances": 5, + }, + ] diff --git a/src/orb/providers/azure/services/termination_service.py b/src/orb/providers/azure/services/termination_service.py new file mode 100644 index 000000000..d715f0acd --- /dev/null +++ b/src/orb/providers/azure/services/termination_service.py @@ -0,0 +1,319 @@ +"""Azure terminate-instance orchestration.""" + +from __future__ import annotations + +import asyncio +import builtins +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Optional + +from orb.domain.base.ports import LoggingPort +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions import AzureValidationError +from orb.providers.azure.infrastructure.cyclecloud_session import CycleCloudRequestContext +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + AzureHandler, + AzureReleaseContext, + AzureReleaseHostsResult, + AzureReleaseProviderData, +) +from orb.providers.azure.services.operation_parsing import ( + group_instance_ids_by_resource, + resolve_operation_provider_api, + resolve_operation_resource_group, +) +from orb.providers.base.strategy import ProviderOperation, ProviderResult + +if TYPE_CHECKING: + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + + +@dataclass +class _TerminationOperationContext: + """Resolved parameters needed to execute a termination dispatch.""" + + instance_ids: list[str] + grouped_resource_mapping: dict[str, list[str]] + release_context: AzureReleaseContext + handler: AzureHandler + default_resource_id: str + + +@dataclass(frozen=True) +class _TerminationDispatchFailure: + """One failed resource-group termination dispatch.""" + + resource_id: str + instance_ids: list[str] + error_message: str + error_type: str + + +@dataclass(frozen=True) +class _TerminationDispatchResult: + """Outcome of dispatching termination across one or more Azure resources.""" + + provider_data: list[AzureReleaseProviderData] + successful_instance_ids: list[str] + failures: list[_TerminationDispatchFailure] + + +async def _dispatch_release_groups_async( + *, + handler: AzureHandler, + instance_ids: list[str], + grouped_resource_mapping: dict[str, list[str]], + default_resource_id: str, + context: AzureReleaseContext, + logger: LoggingPort, + record_pending_cleanup: Callable[[AzureReleaseHostsResult | None], None], +) -> _TerminationDispatchResult: + """Fan out async release_hosts calls per resource and collect provider data. + + Module-level so the existing fan-out tests can exercise this behavior in + isolation without constructing the full ``AzureTerminationService``. + """ + termination_provider_data: list[AzureReleaseProviderData] = [] + successful_instance_ids: list[str] = [] + dispatch_failures: list[Exception] = [] + failure_details: list[_TerminationDispatchFailure] = [] + + dispatch_groups = grouped_resource_mapping or {default_resource_id: instance_ids} + dispatch_group_items = list(dispatch_groups.items()) + handler_results = await asyncio.gather( + *[ + handler.release_hosts_async( + machine_ids=mapped_instance_ids, + resource_id=resource_id, + context=context, + ) + for resource_id, mapped_instance_ids in dispatch_group_items + ], + return_exceptions=True, + ) + for (resource_id, mapped_instance_ids), handler_result in zip( + dispatch_group_items, + handler_results, + strict=True, + ): + if isinstance(handler_result, BaseException): + # Cancellation and system-level exceptions (CancelledError, + # KeyboardInterrupt, SystemExit) must propagate, not be swallowed + # as a dispatch failure. + if not isinstance(handler_result, Exception): + raise handler_result + dispatch_failures.append(handler_result) + failure_details.append( + _TerminationDispatchFailure( + resource_id=resource_id, + instance_ids=list(mapped_instance_ids), + error_message=str(handler_result), + error_type=type(handler_result).__name__, + ) + ) + logger.warning( + "Azure termination dispatch group failed: %s", + handler_result, + exc_info=True, + ) + continue + record_pending_cleanup(handler_result) + if handler_result is None: + continue + + provider_data = handler_result.get("provider_data") + if provider_data is not None: + termination_provider_data.append(provider_data) + successful_instance_ids.extend(mapped_instance_ids) + + if dispatch_failures and not successful_instance_ids: + if len(dispatch_failures) > 1: + raise builtins.ExceptionGroup( + "All Azure termination dispatch groups failed", + dispatch_failures, + ) + raise dispatch_failures[0] + return _TerminationDispatchResult( + provider_data=termination_provider_data, + successful_instance_ids=successful_instance_ids, + failures=failure_details, + ) + + +class AzureTerminationService: + """Own Azure termination from validation through handler dispatch to result shaping.""" + + def __init__( + self, + *, + logger: LoggingPort, + handler_provider: AzureProviderStrategy, + record_pending_cleanup: Callable[[AzureReleaseHostsResult | None], None], + default_resource_group: Optional[str], + ) -> None: + self._logger = logger + self._handler_provider = handler_provider + self._record_pending_cleanup = record_pending_cleanup + self._default_resource_group = default_resource_group + + async def terminate_instances_async( + self, + operation: ProviderOperation, + *, + is_dry_run: bool, + ) -> ProviderResult: + """Validate, dispatch, and shape the result for an Azure terminate-instances operation.""" + context = self._build_context(operation, is_dry_run=is_dry_run) + if is_dry_run: + return _dry_run_result(context) + + dispatch_result = await _dispatch_release_groups_async( + handler=context.handler, + instance_ids=context.instance_ids, + grouped_resource_mapping=context.grouped_resource_mapping, + default_resource_id=context.default_resource_id, + context=context.release_context, + logger=self._logger, + record_pending_cleanup=self._record_pending_cleanup, + ) + return _termination_result(context.instance_ids, dispatch_result) + + def _build_context( + self, + operation: ProviderOperation, + *, + is_dry_run: bool, + ) -> _TerminationOperationContext: + """Validate and resolve a termination operation into a dispatch context.""" + instance_ids = operation.parameters.get("instance_ids", []) + if not instance_ids: + raise AzureValidationError( + "Instance IDs are required for termination", + error_code="MISSING_INSTANCE_IDS", + ) + + provider_api = resolve_operation_provider_api(operation) + if provider_api is None: + raise AzureValidationError( + "provider_api is required for Azure termination", + error_code="MISSING_PROVIDER_API", + ) + + provider_api_value = provider_api.value + handler = self._handler_provider.resolve_handler(provider_api) + if handler is None: + raise AzureValidationError( + f"No handler available for provider_api: {provider_api_value}", + error_code="HANDLER_NOT_FOUND", + ) + + raw_resource_mapping = operation.parameters.get("resource_mapping", {}) + grouped_resource_mapping = group_instance_ids_by_resource( + instance_ids, raw_resource_mapping + ) + request_metadata = operation.parameters.get("request_metadata") or {} + + default_resource_id = operation.parameters.get("resource_id") + if not default_resource_id and grouped_resource_mapping: + default_resource_id = next(iter(grouped_resource_mapping.keys())) + if not default_resource_id and provider_api_value == AzureProviderApi.CYCLECLOUD.value: + cyclecloud_cluster_name = request_metadata.get("cluster_name") + if cyclecloud_cluster_name not in (None, ""): + default_resource_id = str(cyclecloud_cluster_name) + if not default_resource_id and not is_dry_run: + raise AzureValidationError( + "resource_id or resource_mapping is required for Azure termination", + error_code="MISSING_RESOURCE_ID", + ) + + resolved_resource_group = resolve_operation_resource_group( + operation, self._default_resource_group + ) + cyclecloud_request_context = CycleCloudRequestContext.from_mapping(request_metadata) + release_context = AzureReleaseContext( + resource_group=resolved_resource_group, + resource_id=(default_resource_id or None), + cyclecloud_request_context=cyclecloud_request_context, + ) + + return _TerminationOperationContext( + instance_ids=instance_ids, + grouped_resource_mapping=grouped_resource_mapping, + release_context=release_context, + handler=handler, + default_resource_id=default_resource_id or "", + ) + + +def _dry_run_result(context: _TerminationOperationContext) -> ProviderResult: + """Return a success result describing what a real termination would do.""" + return ProviderResult.success_result( + { + "success": True, + "terminated_count": len(context.instance_ids), + }, + { + "operation": "terminate_instances", + "instance_ids": context.instance_ids, + "method": "dry_run", + "provider_data": {"dry_run": True}, + }, + ) + + +def _termination_result( + instance_ids: list[str], + dispatch_result: _TerminationDispatchResult, +) -> ProviderResult: + """Build the final termination result from handler responses.""" + provider_data = ( + { + "termination_requests": dispatch_result.provider_data, + } + if dispatch_result.provider_data + else {} + ) + + if dispatch_result.failures: + failed_instances = [ + instance_id + for failure in dispatch_result.failures + for instance_id in failure.instance_ids + ] + failure_details = [ + { + "resource_id": failure.resource_id, + "instance_ids": failure.instance_ids, + "error_message": failure.error_message, + "error_type": failure.error_type, + } + for failure in dispatch_result.failures + ] + return ProviderResult.error_result( + "Azure termination partially failed", + "AZURE_TERMINATION_PARTIAL_FAILURE", + { + "operation": "terminate_instances", + "instance_ids": instance_ids, + "method": "handler", + "provider_data": { + **provider_data, + "dispatch_failures": failure_details, + "successful_instance_ids": dispatch_result.successful_instance_ids, + "failed_instance_ids": failed_instances, + }, + }, + ) + + return ProviderResult.success_result( + { + "success": True, + "terminated_count": len(instance_ids), + }, + { + "operation": "terminate_instances", + "instance_ids": instance_ids, + "method": "handler", + "provider_data": provider_data, + }, + ) diff --git a/src/orb/providers/azure/strategy/__init__.py b/src/orb/providers/azure/strategy/__init__.py new file mode 100644 index 000000000..89722f0ad --- /dev/null +++ b/src/orb/providers/azure/strategy/__init__.py @@ -0,0 +1,5 @@ +"""Azure provider strategy.""" + +from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + +__all__: list[str] = ["AzureProviderStrategy"] diff --git a/src/orb/providers/azure/strategy/azure_provider_strategy.py b/src/orb/providers/azure/strategy/azure_provider_strategy.py new file mode 100644 index 000000000..bf54e7b1f --- /dev/null +++ b/src/orb/providers/azure/strategy/azure_provider_strategy.py @@ -0,0 +1,1003 @@ +"""Azure Provider Strategy. + +Implements ``ProviderStrategy`` for Azure, routing all seven operation types +to the appropriate handlers via the VMSS / SingleVM infrastructure layer. +""" + +from __future__ import annotations + +import argparse +import asyncio +import inspect +import time +from collections.abc import Mapping +from threading import Condition, RLock +from typing import TYPE_CHECKING, Any, Callable, Optional + +from orb.domain.base.ports import LoggingPort +from orb.infrastructure.di.injectable import injectable +from orb.providers.azure.capabilities import ( + get_supported_api_capabilities, + get_supported_apis, +) +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.configuration.validator import validate_azure_template +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzurePriority, AzureProviderApi +from orb.providers.azure.exceptions import AzureError, AzureValidationError +from orb.providers.azure.infrastructure.error_utils import ( + canonical_azure_error_code, + extract_azure_error_details, +) +from orb.providers.azure.infrastructure.vmss_cleanup import VmssCleanupCoordinator +from orb.providers.azure.services.cyclecloud_request_context_service import ( + resolve_cyclecloud_request_metadata, +) +from orb.providers.azure.services.health_check_service import AzureHealthCheckService +from orb.providers.azure.services.inventory_service import ( + AzureInventoryService, + AzureReadOperationContext, + build_read_operation_context, +) +from orb.providers.azure.services.operation_parsing import resolve_operation_provider_api +from orb.providers.azure.services.provisioning_service import ( + AzureProvisioningService, + create_instances_dry_run_result, +) +from orb.providers.azure.services.resource_metadata_service import ( + AzureResourceMetadataService, +) +from orb.providers.azure.services.runtime_dependencies import AzureRuntimeDependencies +from orb.providers.azure.services.spot_launch_service import AzureSpotLaunchService +from orb.providers.azure.services.spot_placement_execution import ( + SpotPlacementExecutionService, +) +from orb.providers.azure.services.spot_placement_planner import ( + PlacementPlanEntry, + SpotPlacementPlanner, +) +from orb.providers.azure.services.template_catalog_service import AzureTemplateCatalogService +from orb.providers.azure.services.termination_service import AzureTerminationService +from orb.providers.base.strategy import ( + ProviderCapabilities, + ProviderHealthStatus, + ProviderOperation, + ProviderOperationType, + ProviderResult, + ProviderStrategy, +) + +if TYPE_CHECKING: + from orb.providers.azure.infrastructure.azure_client import AzureClient + from orb.providers.azure.infrastructure.azure_handler_factory import AzureHandlerFactory + from orb.providers.azure.infrastructure.handlers.azure_handler import AzureHandler + from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, + ) + from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, + ) + from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + +def _provider_config_mapping(template_config: dict[str, Any]) -> Mapping[str, Any]: + raw_provider_config = template_config.get("provider_config") + if raw_provider_config is None: + return {} + if not isinstance(raw_provider_config, Mapping): + raise AzureValidationError("Azure provider_config must be a mapping") + return raw_provider_config + + +def _has_provider_config_value( + provider_config: Mapping[str, Any], + key: str, +) -> bool: + return provider_config.get(key) not in (None, "") + + +@injectable +class AzureProviderStrategy(ProviderStrategy): + """Azure implementation of ``ProviderStrategy``. + + Adapts the Azure infrastructure layer (VMSS / SingleVM handlers, + AzureClient, AzureResourceManager) to the strategy pattern, enabling + runtime provider switching and composition with the advanced strategy + wrappers (fallback, composite, load-balancing). + """ + + def __init__( + self, + config: AzureProviderConfig, + logger: LoggingPort, + provider_instance_name: str, + azure_client_resolver: Optional[Callable[[], AzureClient]] = None, + azure_handler_factory_resolver: Optional[Callable[[], AzureHandlerFactory]] = None, + azure_resource_manager_resolver: Optional[Callable[[], AzureResourceManager | None]] = None, + azure_deployment_service_resolver: Optional[ + Callable[[], AzureDeploymentService | None] + ] = None, + azure_native_spec_service: Optional[AzureNativeSpecService] = None, + vmss_cleanup_coordinator: Optional[VmssCleanupCoordinator] = None, + ) -> None: + """Initialise the Azure strategy with config, logger, and optional client resolver.""" + if not isinstance(config, AzureProviderConfig): + raise ValueError("AzureProviderStrategy requires AzureProviderConfig") + + super().__init__(config) + self._logger = logger + self._azure_config = config + self._provider_instance_name = provider_instance_name + self._azure_native_spec_service = azure_native_spec_service + self._handlers: dict[str, AzureHandler] = {} + self._spot_placement_planner = SpotPlacementPlanner() + self._spot_placement_execution = SpotPlacementExecutionService() + self._health_check_service = AzureHealthCheckService(config=config, logger=logger) + self._provisioning_service = AzureProvisioningService() + self._resource_metadata_service = AzureResourceMetadataService( + default_resource_group=config.resource_group, + logger=logger, + ) + self._spot_launch_service = AzureSpotLaunchService( + config=config, + logger=logger, + planner=self._spot_placement_planner, + execution_service=self._spot_placement_execution, + ) + self._template_catalog_service = AzureTemplateCatalogService(logger=logger) + self._runtime = AzureRuntimeDependencies( + config=config, + logger=logger, + azure_client_resolver=azure_client_resolver, + azure_handler_factory_resolver=azure_handler_factory_resolver, + azure_resource_manager_resolver=azure_resource_manager_resolver, + azure_deployment_service_resolver=azure_deployment_service_resolver, + azure_native_spec_service=azure_native_spec_service, + ) + self._lifecycle_lock = RLock() + self._lifecycle_condition = Condition(self._lifecycle_lock) + self._active_operations = 0 + self._cleanup_requested = False + self._cleanup_wait_timeout_seconds = 30.0 + self._vmss_cleanup_coordinator = vmss_cleanup_coordinator or VmssCleanupCoordinator( + logger=self._logger, + get_vmss_member_count=self._current_vmss_member_count_async, + vmss_exists=self._vmss_exists_async, + begin_delete_vmss=self._begin_delete_vmss_async, + ) + self._termination_service = AzureTerminationService( + logger=logger, + handler_provider=self, + record_pending_cleanup=self._vmss_cleanup_coordinator.record, + default_resource_group=config.resource_group, + ) + self._inventory_service = AzureInventoryService( + logger=logger, + provider_instance_name=provider_instance_name, + resource_metadata_service=self._resource_metadata_service, + handler_provider=self, + vmss_cleanup_coordinator=self._vmss_cleanup_coordinator, + ) + + # ------------------------------------------------------------------ + # Lazy-initialised properties + # ------------------------------------------------------------------ + + @property + def provider_type(self) -> str: + """Return the provider type identifier.""" + return "azure" + + @property + def provider_instance_name(self) -> str: + """Return the configured name for this provider instance.""" + return self._provider_instance_name + + @property + def azure_client(self) -> Optional[AzureClient]: + """Get the Azure client with lazy initialisation.""" + if self._cleanup_requested: + return None + return self._runtime.azure_client + + @property + def resource_manager(self) -> Optional[AzureResourceManager]: + """Get the Azure resource manager with lazy initialisation.""" + return self._runtime.resource_manager + + @property + def deployment_service(self) -> Optional[AzureDeploymentService]: + """Get the ARM deployment service with lazy initialisation.""" + return self._runtime.deployment_service + + @property + def handlers(self) -> dict[str, AzureHandler]: + """Get handler mapping, including any explicit test overrides.""" + handler_factory = self._get_handler_factory() + handlers = handler_factory.get_all_handlers() if handler_factory is not None else {} + handlers.update(self._handlers) + return handlers + + def _get_handler_factory(self) -> Optional[AzureHandlerFactory]: + """Resolve the Azure handler factory lazily using the strategy-owned client.""" + return self._runtime.handler_factory + + def resolve_handler( + self, + provider_api: AzureProviderApi, + *, + allow_vmss_uniform_fallback: bool = False, + ) -> Optional[AzureHandler]: + """Resolve one handler from explicit overrides or the canonical factory.""" + provider_api_value = provider_api.value + handler = self._handlers.get(provider_api_value) + if handler is not None: + return handler + + handler_factory = self._get_handler_factory() + if handler_factory is None: + return None + + try: + return handler_factory.create_handler(provider_api) + except AzureValidationError: + if allow_vmss_uniform_fallback and provider_api == AzureProviderApi.VMSS_UNIFORM: + override_handler = self._handlers.get(AzureProviderApi.VMSS.value) + if override_handler is not None: + return override_handler + return handler_factory.create_handler(AzureProviderApi.VMSS) + return None + + def _build_azure_template_config(self, template_config: dict[str, Any]) -> dict[str, Any]: + """Coalesce provider-owned and Azure-default fields before AzureTemplate validation.""" + enhanced_config = dict(template_config) + provider_config = _provider_config_mapping(enhanced_config) + + raw_subnet_id = enhanced_config.get("subnet_id") + if raw_subnet_id and raw_subnet_id != "default-subnet": + enhanced_config["subnet_ids"] = [raw_subnet_id] + elif enhanced_config.get("subnet_ids") == ["default-subnet"]: + enhanced_config.pop("subnet_ids", None) + + if enhanced_config.get("priority") in ( + None, + "", + ) and not _has_provider_config_value(provider_config, "priority"): + enhanced_config["priority"] = AzurePriority.REGULAR + if enhanced_config.get("admin_username") in ( + None, + "", + ) and not _has_provider_config_value(provider_config, "admin_username"): + enhanced_config["admin_username"] = "azureuser" + if enhanced_config.get("node_attributes") in ( + None, + "", + ) and not _has_provider_config_value(provider_config, "node_attributes"): + enhanced_config["node_attributes"] = {} + + if ( + enhanced_config.get("resource_group") in (None, "") + and not _has_provider_config_value(provider_config, "resource_group") + and self._azure_config.resource_group + ): + enhanced_config["resource_group"] = self._azure_config.resource_group + if ( + enhanced_config.get("location") in (None, "") + and not _has_provider_config_value(provider_config, "location") + and self._azure_config.region + ): + # Provider config follows the shared ``region`` interface, but the + # Azure template model is Azure-native and owns ``location``. + enhanced_config["location"] = self._azure_config.region + if ( + enhanced_config.get("subscription_id") in (None, "") + and not _has_provider_config_value(provider_config, "subscription_id") + and self._azure_config.subscription_id + ): + enhanced_config["subscription_id"] = self._azure_config.subscription_id + + enhanced_config.setdefault("provider_type", "azure") + enhanced_config.setdefault("provider_name", self.provider_instance_name) + return enhanced_config + + def _build_spot_placement_plan( + self, + azure_template: AzureTemplate, + count: int, + ) -> list[PlacementPlanEntry]: + """Compatibility wrapper for tests and callers that patch this seam.""" + return self._spot_launch_service.build_spot_placement_plan( + azure_template=azure_template, + count=count, + azure_client=self.azure_client, + ) + + async def _build_spot_placement_plan_async( + self, + azure_template: AzureTemplate, + count: int, + ) -> list[PlacementPlanEntry]: + """Build the spot placement plan without blocking the async create flow.""" + patched_sync_builder = self.__dict__.get("_build_spot_placement_plan") + if patched_sync_builder is not None: + return patched_sync_builder(azure_template, count) + return await self._spot_launch_service.build_spot_placement_plan_async( + azure_template=azure_template, + count=count, + azure_client=self.azure_client, + ) + + def _is_capacity_like_failure(self, child_result: dict[str, Any]) -> bool: + """Compatibility wrapper for tests and callers that patch this seam.""" + return self._spot_launch_service.is_capacity_like_failure(child_result) + + # ------------------------------------------------------------------ + # ProviderStrategy contract + # ------------------------------------------------------------------ + + def initialize(self) -> bool: + """Mark the strategy as initialised and ready to execute operations.""" + with self._lifecycle_condition: + self._cleanup_requested = False + self._initialized = True + self._logger.info( + "Azure provider strategy ready for region: %s", + self._azure_config.region, + ) + return True + + def _begin_operation(self) -> bool: + """Reserve an execution slot unless cleanup has already started.""" + with self._lifecycle_condition: + if self._cleanup_requested: + return False + self._active_operations += 1 + return True + + def _end_operation(self) -> None: + """Release an execution slot and wake any waiting cleanup path.""" + with self._lifecycle_condition: + self._active_operations -= 1 + if self._active_operations == 0: + self._lifecycle_condition.notify_all() + + async def execute_operation(self, operation: ProviderOperation) -> ProviderResult: + """Compatibility entrypoint that delegates to the native async override.""" + return await self.execute_operation_async(operation) + + async def execute_operation_async(self, operation: ProviderOperation) -> ProviderResult: + """Execute an Azure provider operation via the native async strategy path.""" + self._logger.debug( + "azure_provider_strategy execute_operation [%s, %s, %s]", + operation.operation_type, + operation.parameters, + operation.context, + ) + # Azure's primary SDK surface is already async-capable, so the strategy + # stays natively async here and only pushes specific sync-only helpers + # behind narrow service-level bridges where needed. + if not self._initialized: + return ProviderResult.error_result( + "Azure provider strategy not initialized", "NOT_INITIALIZED" + ) + if not self._begin_operation(): + return ProviderResult.error_result( + "Azure provider strategy is shutting down", + "STRATEGY_SHUTTING_DOWN", + ) + + start_time = time.time() + is_dry_run = bool(operation.context and operation.context.get("dry_run", False)) + + try: + from orb.providers.azure.infrastructure.dry_run_adapter import azure_dry_run_context + + if is_dry_run: + # Activates the global dry-run flag checked by is_dry_run_active(). + # Individual operations short-circuit via early returns; the context + # manager is the integration point for future SDK-level mocking + # (analogous to the AWS moto adapter). + with azure_dry_run_context(): + result = await self._execute_operation_internal(operation) + else: + result = await self._execute_operation_internal(operation) + + execution_time_ms = int((time.time() - start_time) * 1000) + return result.model_copy( + update={ + "routing_info": { + "execution_time_ms": execution_time_ms, + "provider": "azure", + }, + "metadata": { + **(result.metadata or {}), + "dry_run": is_dry_run, + "execution_time_ms": execution_time_ms, + "provider": "azure", + }, + } + ) + except asyncio.CancelledError: + raise + except Exception as exc: + execution_time_ms = int((time.time() - start_time) * 1000) + self._logger.error("Azure operation failed: %s", exc, exc_info=True) + return ProviderResult.error_result( + f"Azure operation failed: {exc!s}", + "OPERATION_FAILED", + {"dry_run": is_dry_run}, + ).model_copy( + update={ + "routing_info": { + "execution_time_ms": execution_time_ms, + "provider": "azure", + } + } + ) + finally: + self._end_operation() + + def get_capabilities(self) -> ProviderCapabilities: + """Get Azure provider capabilities.""" + return ProviderCapabilities( + provider_type="azure", + supported_operations=[ + ProviderOperationType.CREATE_INSTANCES, + ProviderOperationType.TERMINATE_INSTANCES, + ProviderOperationType.GET_INSTANCE_STATUS, + ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + ProviderOperationType.VALIDATE_TEMPLATE, + ProviderOperationType.GET_AVAILABLE_TEMPLATES, + ProviderOperationType.HEALTH_CHECK, + ], + supported_apis=get_supported_apis(), + features={ + "api_capabilities": get_supported_api_capabilities(), + "instance_management": True, + "spot_instances": True, + "fleet_management": True, + "auto_scaling": True, + "load_balancing": True, + "vpc_support": True, + "security_groups": True, + "key_pairs": True, + "tags_support": True, + "monitoring": True, + }, + ) + + def check_health(self) -> ProviderHealthStatus: + """Check Azure connectivity and return the current health status.""" + return self._health_check_service.check_health() + + def get_available_credential_sources(self) -> list[dict]: + """Return Azure credential source options.""" + return [{"name": None, "description": "DefaultAzureCredential / managed identity"}] + + def get_credential_requirements(self) -> dict: + """Azure auth uses ambient credentials and does not require pre-auth prompts.""" + return {} + + def get_operational_requirements(self) -> dict: + """Return the Azure values init must collect to build a working provider config.""" + return { + "subscription_id": {"required": True, "description": "Azure subscription ID"}, + "resource_group": {"required": True, "description": "Azure resource group"}, + "region": {"required": True, "description": "Azure location"}, + "client_id": { + "required": False, + "prompt": True, + "description": "Managed identity client ID (optional)", + }, + } + + def get_default_region(self) -> str: + """Return the default Azure location for CLI prompts.""" + return "eastus2" + + def get_cli_provider_config(self, args: argparse.Namespace) -> dict[str, Any]: + """Extract Azure provider config from init CLI args.""" + args_dict = vars(args) + provider_config: dict[str, Any] = { + "region": args_dict.get("azure_location") or self.get_default_region(), + } + + field_map = { + "subscription_id": "azure_subscription_id", + "resource_group": "azure_resource_group", + "client_id": "azure_client_id", + } + for config_key, arg_name in field_map.items(): + value = args_dict.get(arg_name) + if value not in (None, ""): + provider_config[config_key] = value + + cyclecloud: dict[str, Any] = {} + cyclecloud_field_map = { + "url": "azure_cyclecloud_url", + "credential_path": "azure_cyclecloud_credential_path", + "auth_mode": "azure_cyclecloud_auth_mode", + "aad_scope": "azure_cyclecloud_aad_scope", + } + for config_key, arg_name in cyclecloud_field_map.items(): + value = args_dict.get(arg_name) + if value not in (None, ""): + cyclecloud[config_key] = value + + if args_dict.get("azure_cyclecloud_verify_ssl"): + cyclecloud["verify_ssl"] = True + elif args_dict.get("azure_cyclecloud_no_verify_ssl"): + cyclecloud["verify_ssl"] = False + + if cyclecloud: + provider_config["cyclecloud"] = cyclecloud + + return provider_config + + def test_credentials(self, credential_source: Optional[str] = None, **kwargs) -> dict: + """Validate Azure credentials by performing a health check.""" + del credential_source + health = self._health_check_service.check_health() + if health.is_healthy: + return {"success": True} + return { + "success": False, + "error": health.status_message, + "details": health.error_details or {}, + } + + def generate_provider_name(self, config: dict[str, Any]) -> str: + """Generate Azure provider name.""" + subscription_id = config.get("subscription_id", "default") + region = config.get("region", "eastus") + return f"azure_{subscription_id}_{region}" + + def parse_provider_name(self, provider_name: str) -> dict[str, str]: + """Parse Azure provider name.""" + parts = provider_name.split("_") + return { + "type": parts[0] if len(parts) > 0 else "azure", + "subscription_id": parts[1] if len(parts) > 1 else "default", + "region": parts[2] if len(parts) > 2 else "eastus2", + } + + def get_provider_name_pattern(self) -> str: + """Return the regex pattern used to validate Azure provider names.""" + return "{type}_{subscription_id}_{region}" + + def get_supported_apis(self) -> list[str]: + """Get supported Azure provider APIs.""" + return get_supported_apis() + + def cleanup(self) -> None: + """Release Azure client resources and reset all lazily initialised state.""" + client = self._prepare_cleanup() + if client is not None: + client.close() + self._logger.debug("Azure provider cleaned up") + + async def cleanup_async(self) -> None: + """Release Azure client resources through the native async close path.""" + client = await asyncio.to_thread(self._prepare_cleanup) + if client is not None: + await client.aclose() + self._logger.debug("Azure provider cleaned up") + + def _prepare_cleanup(self) -> Optional[AzureClient]: + """Wait for active work and clear strategy-owned runtime state.""" + with self._lifecycle_condition: + self._cleanup_requested = True + deadline = time.monotonic() + self._cleanup_wait_timeout_seconds + while self._active_operations > 0: + remaining = deadline - time.monotonic() + if remaining <= 0: + self._logger.warning( + "Azure provider cleanup timed out with %d active operation(s); " + "leaving strategy in shutdown mode until cleanup is retried", + self._active_operations, + ) + return None + self._lifecycle_condition.wait(timeout=remaining) + + client = self._runtime.clear_cached_runtime() + self._handlers = {} + self._vmss_cleanup_coordinator.clear() + self._initialized = False + return client + + @staticmethod + def _error_result( + message: str, + default_code: str, + exc: Exception, + metadata: Optional[dict[str, Any]] = None, + ) -> ProviderResult: + """Build an error result, preserving Azure-specific error codes and details.""" + # getattr: exc may be any Exception subclass, not just AzureError. + error_code = getattr(exc, "error_code", None) or default_code + merged = dict(metadata or {}) + merged["error_class"] = type(exc).__name__ + if isinstance(exc, AzureError): + merged["provider_error"] = exc.to_dict() + return ProviderResult.error_result(message, error_code, merged) + + # ------------------------------------------------------------------ + # Internal operation dispatch + # ------------------------------------------------------------------ + + async def _execute_operation_internal(self, operation: ProviderOperation) -> ProviderResult: + if operation.operation_type == ProviderOperationType.CREATE_INSTANCES: + return await self._handle_create_instances(operation) + elif operation.operation_type == ProviderOperationType.TERMINATE_INSTANCES: + return await self._handle_terminate_instances(operation) + elif operation.operation_type == ProviderOperationType.GET_INSTANCE_STATUS: + return await self._handle_get_instance_status(operation) + elif operation.operation_type == ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES: + return await self._handle_describe_resource_instances(operation) + elif operation.operation_type == ProviderOperationType.VALIDATE_TEMPLATE: + return await self._handle_validate_template(operation) + elif operation.operation_type == ProviderOperationType.GET_AVAILABLE_TEMPLATES: + return await self._handle_get_available_templates() + elif operation.operation_type == ProviderOperationType.HEALTH_CHECK: + return await self._handle_health_check(operation) + else: + return ProviderResult.error_result( + f"Unsupported operation: {operation.operation_type}", + "UNSUPPORTED_OPERATION", + ) + + # ------------------------------------------------------------------ + # CREATE_INSTANCES + # ------------------------------------------------------------------ + + async def _handle_create_instances(self, operation: ProviderOperation) -> ProviderResult: + template_config: dict[str, Any] = {} + provider_api_key: Optional[str] = None + try: + create_context = self._provisioning_service.build_create_operation_context( + operation=operation, + resolve_handler=self.resolve_handler, + build_template=lambda tc: AzureTemplate.model_validate( + self._build_azure_template_config(tc) + ), + ) + + template_config = create_context.template_config + provider_api_key = create_context.provider_api_key + + if bool(operation.context and operation.context.get("dry_run", False)): + return create_instances_dry_run_result(create_context) + + if self._spot_launch_service.should_use_spot_placement(create_context.azure_template): + plan_result = self._build_spot_placement_plan_async( + create_context.azure_template, + create_context.count, + ) + plan_override = ( + await plan_result if inspect.isawaitable(plan_result) else plan_result + ) + return await self._spot_launch_service.execute_planned_spot_launches_async( + azure_template=create_context.azure_template, + provider_api=create_context.provider_api, + provider_api_key=create_context.provider_api.value, + count=create_context.count, + template_config=template_config, + operation=operation, + provider_instance_name=self.provider_instance_name, + handler=self.resolve_handler(create_context.provider_api), + azure_client=self.azure_client, + plan_override=plan_override, + capacity_like_failure_checker=self._is_capacity_like_failure, + ) + + request = self._provisioning_service.build_create_request( + operation=operation, + azure_template=create_context.azure_template, + count=create_context.count, + provider_api=create_context.provider_api, + provider_instance_name=self.provider_instance_name, + ) + return await self._provisioning_service.execute_create_handler_async( + create_context=create_context, + request=request, + ) + + except asyncio.CancelledError: + raise + except AzureValidationError as exc: + return self._error_result( + str(exc), + "CREATE_INSTANCES_ERROR", + exc, + ) + except Exception as exc: + error_details = extract_azure_error_details(exc) + fleet_error = { + "error_code": canonical_azure_error_code(exc), + "error_message": error_details["message"], + "status_code": error_details["status_code"], + "raw_error_code": error_details["raw_error_code"], + "details": error_details["details"], + } + return self._error_result( + f"Failed to create instances: {exc!s}", + "CREATE_INSTANCES_ERROR", + exc, + metadata={ + "operation": "create_instances", + "template_config": template_config, + "handler_used": provider_api_key, + "provider_data": {"fleet_errors": [fleet_error]}, + }, + ) + + # ------------------------------------------------------------------ + # TERMINATE_INSTANCES + # ------------------------------------------------------------------ + + async def _handle_terminate_instances(self, operation: ProviderOperation) -> ProviderResult: + self._logger.debug("_handle_terminate_instances") + try: + operation = self._resolve_cyclecloud_operation(operation) + is_dry_run = bool(operation.context and operation.context.get("dry_run", False)) + return await self._termination_service.terminate_instances_async( + operation, is_dry_run=is_dry_run + ) + + except asyncio.CancelledError: + raise + except AzureValidationError as exc: + return self._error_result( + str(exc), + "TERMINATE_INSTANCES_ERROR", + exc, + ) + except Exception as exc: + return self._error_result( + f"Failed to terminate instances: {exc!s}", + "TERMINATE_INSTANCES_ERROR", + exc, + ) + + def _resolve_cyclecloud_operation(self, operation: ProviderOperation) -> ProviderOperation: + """Merge durable CycleCloud follow-up context into an operation on demand.""" + provider_api = resolve_operation_provider_api(operation) + if provider_api != AzureProviderApi.CYCLECLOUD: + return operation + + resolved_request_metadata = resolve_cyclecloud_request_metadata( + operation=operation, + ) + return ProviderOperation( + operation_type=operation.operation_type, + parameters={**operation.parameters, "request_metadata": resolved_request_metadata}, + context=operation.context, + ) + + # ------------------------------------------------------------------ + # GET_INSTANCE_STATUS + # ------------------------------------------------------------------ + + async def _handle_get_instance_status(self, operation: ProviderOperation) -> ProviderResult: + try: + operation = self._resolve_cyclecloud_operation(operation) + read_context = build_read_operation_context( + operation=operation, + operation_name="get_instance_status", + default_resource_group=self._azure_config.resource_group, + ) + + if read_context.resource_group is None: + raise RuntimeError( + "build_read_operation_context must provide resource_group for status queries" + ) + instance_ids = read_context.instance_ids + + if bool(operation.context and operation.context.get("dry_run", False)): + return self._dry_run_instance_status_result(instance_ids) + + return await self._inventory_service.get_instance_status_async(read_context) + + except asyncio.CancelledError: + raise + except Exception as exc: + return self._error_result( + f"Failed to get instance status: {exc!s}", + "GET_INSTANCE_STATUS_ERROR", + exc, + ) + + async def _current_vmss_member_count_async( + self, + *, + resource_group: str, + vmss_name: str, + ) -> Optional[int]: + if not self.resource_manager: + return None + + return await self.resource_manager.get_vmss_member_count_async( + resource_group=resource_group, + vmss_name=vmss_name, + ) + + async def _vmss_exists_async( + self, + *, + resource_group: str, + vmss_name: str, + ) -> Optional[bool]: + if not self.resource_manager: + return None + + return await self.resource_manager.vmss_exists_async( + resource_group=resource_group, + vmss_name=vmss_name, + ) + + async def _begin_delete_vmss_async( + self, + *, + resource_group: str, + vmss_name: str, + ) -> None: + azure_client = self.azure_client + if not azure_client: + raise RuntimeError("Azure client not available for VMSS cleanup delete") + compute = await azure_client.get_async_compute_client() + await compute.virtual_machine_scale_sets.begin_delete( + resource_group_name=resource_group, + vm_scale_set_name=vmss_name, + ) + + # ------------------------------------------------------------------ + # DESCRIBE_RESOURCE_INSTANCES + # ------------------------------------------------------------------ + + async def _handle_describe_resource_instances( + self, operation: ProviderOperation + ) -> ProviderResult: + try: + operation = self._resolve_cyclecloud_operation(operation) + read_context = build_read_operation_context( + operation=operation, + operation_name="describe_resource_instances", + default_resource_group=self._azure_config.resource_group, + ) + if bool(operation.context and operation.context.get("dry_run", False)): + return self._dry_run_describe_instances_result(read_context) + return await self._inventory_service.describe_resource_instances_async( + read_context=read_context, + resource_manager=self.resource_manager, + deployment_service=self.deployment_service, + ) + + except asyncio.CancelledError: + raise + except Exception as exc: + return self._error_result( + f"Failed to describe resource instances: {exc!s}", + "DESCRIBE_RESOURCE_INSTANCES_ERROR", + exc, + ) + + def _dry_run_instance_status_result(self, instance_ids: list[str]) -> ProviderResult: + """Build a dry-run status result with the same fulfilment contract as real status.""" + dry_run_instances = [ + { + "instance_id": instance_id, + "status": "unknown", + "provider_type": "azure", + "provider_data": {"dry_run": True}, + } + for instance_id in instance_ids + ] + metadata: dict[str, Any] = { + "operation": "get_instance_status", + "instance_ids": instance_ids, + "method": "dry_run", + "provider_data": {"dry_run": True}, + } + status_result = self._resource_metadata_service.attach_provider_fulfilment( + metadata, + instances=dry_run_instances, + target_units=len(instance_ids), + ) + return ProviderResult.success_result( + { + "instances": status_result.instances, + "queried_count": len(instance_ids), + }, + metadata, + ) + + def _dry_run_describe_instances_result( + self, + read_context: AzureReadOperationContext, + ) -> ProviderResult: + """Build a dry-run resource discovery result with an explicit unknown target.""" + metadata: dict[str, Any] = { + "operation": "describe_resource_instances", + "resource_ids": read_context.resource_ids, + "provider_api": read_context.provider_api_key, + "method": "dry_run", + "provider_data": {"dry_run": True}, + } + status_result = self._resource_metadata_service.attach_provider_fulfilment( + metadata, + instances=[], + target_units=None, + ) + return ProviderResult.success_result( + {"instances": status_result.instances}, + metadata, + ) + + # ------------------------------------------------------------------ + # VALIDATE_TEMPLATE + # ------------------------------------------------------------------ + + async def _handle_validate_template(self, operation: ProviderOperation) -> ProviderResult: + try: + template_config = operation.parameters.get("template_config", {}) + if not template_config: + return ProviderResult.error_result( + "Template configuration is required for validation", + "MISSING_TEMPLATE_CONFIG", + ) + + validation_result = validate_azure_template(template_config) + + return ProviderResult.success_result( + validation_result, + {"operation": "validate_template", "template_config": template_config}, + ) + except Exception as exc: + return self._error_result( + f"Failed to validate template: {exc!s}", + "VALIDATE_TEMPLATE_ERROR", + exc, + ) + + # ------------------------------------------------------------------ + # GET_AVAILABLE_TEMPLATES + # ------------------------------------------------------------------ + + async def _handle_get_available_templates(self) -> ProviderResult: + try: + templates = self._template_catalog_service.get_available_templates() + return ProviderResult.success_result( + {"templates": templates, "count": len(templates)}, + {"operation": "get_available_templates"}, + ) + except Exception as exc: + return ProviderResult.error_result( + f"Failed to get available templates: {exc!s}", + "GET_TEMPLATES_ERROR", + ) + + # ------------------------------------------------------------------ + # HEALTH_CHECK + # ------------------------------------------------------------------ + + async def _handle_health_check(self, operation: ProviderOperation) -> ProviderResult: + _ = operation + health_status = await self._health_check_service.check_health_async() + return ProviderResult.success_result( + { + "is_healthy": health_status.is_healthy, + "status_message": health_status.status_message, + "response_time_ms": health_status.response_time_ms, + }, + {"operation": "health_check"}, + ) + + # ------------------------------------------------------------------ + # String representations + # ------------------------------------------------------------------ + + def __str__(self) -> str: + return ( + f"AzureProviderStrategy(region={self._azure_config.region}, " + f"initialized={self._initialized})" + ) diff --git a/src/orb/providers/registration.py b/src/orb/providers/registration.py index 23e27189e..96a7c4833 100644 --- a/src/orb/providers/registration.py +++ b/src/orb/providers/registration.py @@ -28,7 +28,7 @@ # --------------------------------------------------------------------------- # Central provider list – the single line to edit when adding a new provider. # --------------------------------------------------------------------------- -_REGISTERED_PROVIDERS: list[str] = ["aws", "k8s"] +_REGISTERED_PROVIDERS: list[str] = ["aws", "azure", "k8s"] # Entry-point group used by third-party plugins to register provider # extensions. See ``discover_provider_plugins`` and diff --git a/tests/providers/azure/__init__.py b/tests/providers/azure/__init__.py new file mode 100644 index 000000000..e938a42f8 --- /dev/null +++ b/tests/providers/azure/__init__.py @@ -0,0 +1 @@ +"""Azure provider test package.""" diff --git a/tests/providers/azure/conftest.py b/tests/providers/azure/conftest.py new file mode 100644 index 000000000..7ae4603b1 --- /dev/null +++ b/tests/providers/azure/conftest.py @@ -0,0 +1,32 @@ +"""Shared Azure strategy test fixtures.""" + +from unittest.mock import MagicMock + +import pytest + +from orb.providers.azure.configuration.config import AzureProviderConfig +from tests.providers.azure.strategy_test_support import build_strategy_harness + + +@pytest.fixture +def azure_config(): + return AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="test-rg", + region="eastus2", + ) + + +@pytest.fixture +def logger(): + return MagicMock() + + +@pytest.fixture +def strategy_harness(azure_config, logger): + return build_strategy_harness(config=azure_config, logger=logger) + + +@pytest.fixture +def strategy(strategy_harness): + return strategy_harness.strategy diff --git a/tests/providers/azure/strategy_test_support.py b/tests/providers/azure/strategy_test_support.py new file mode 100644 index 000000000..ca6d3c17d --- /dev/null +++ b/tests/providers/azure/strategy_test_support.py @@ -0,0 +1,201 @@ +"""Shared helpers for Azure strategy tests.""" + +import asyncio +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError +from orb.providers.azure.services.provisioning_service import provider_api_key +from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + +_UNSET = object() + + +@dataclass +class AzureStrategyHarness: + """Mutable test harness that feeds explicit dependencies into the strategy.""" + + strategy: AzureProviderStrategy | None = None + handlers: dict[str, Any] = field(default_factory=dict) + azure_client: Any | None = None + resource_manager: Any | None = None + deployment_service: Any | None = None + + +class _TestAzureHandlerFactory: + """Minimal handler factory that resolves from the harness-owned handler map.""" + + def __init__(self, harness: AzureStrategyHarness) -> None: + self._harness = harness + + def create_handler(self, handler_type: object) -> Any: + handler_key = provider_api_key(handler_type) + handler = self._harness.handlers.get(handler_key) + if handler is None: + raise AzureValidationError(f"No handler class registered for type: {handler_key}") + return handler + + def get_all_handlers(self) -> dict[str, Any]: + return dict(self._harness.handlers) + + +class AsyncPager: + """Minimal async iterator helper for Azure SDK pager-shaped tests.""" + + def __init__(self, items: list[Any]) -> None: + self._items = list(items) + + def __aiter__(self) -> "AsyncPager": + self._iterator = iter(self._items) + return self + + async def __anext__(self) -> Any: + try: + return next(self._iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + +def make_azure_template( + *, + template_id: str, + provider_api: str, + vm_size: str = "Standard_D4s_v5", + resource_group: str = "test-rg", + location: str = "eastus2", + network_config: dict[str, Any] | None | object = _UNSET, + ssh_public_keys: list[str] | None | object = _UNSET, + image: dict[str, Any] | None | object = _UNSET, + **overrides: Any, +) -> AzureTemplate: + """Build a canonical Azure test template with overridable defaults.""" + config: dict[str, Any] = { + "template_id": template_id, + "provider_api": provider_api, + "vm_size": vm_size, + "resource_group": resource_group, + "location": location, + "network_config": ( + {"subnet_id": "/subscriptions/.../subnets/default"} + if network_config is _UNSET + else network_config + ), + "ssh_public_keys": ( + ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"] + if ssh_public_keys is _UNSET + else ssh_public_keys + ), + "image": ( + { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + } + if image is _UNSET + else image + ), + } + config.update(overrides) + return AzureTemplate(**config) + + +def make_single_vm_azure_client() -> MagicMock: + """Build the SingleVM handler's async-over-sync Azure client double.""" + azure_client = MagicMock() + async_compute = MagicMock() + async_resource = MagicMock() + + async_compute.virtual_machines.get = AsyncMock( + side_effect=azure_client.compute_client.virtual_machines.get + ) + async_compute.virtual_machines.begin_delete = AsyncMock( + side_effect=azure_client.compute_client.virtual_machines.begin_delete + ) + async_compute.virtual_machines.list = MagicMock( + side_effect=lambda *args, **kwargs: AsyncPager( + azure_client.compute_client.virtual_machines.list(*args, **kwargs) + ) + ) + async_compute.ssh_public_keys.get = AsyncMock( + side_effect=azure_client.compute_client.ssh_public_keys.get + ) + async_resource.resources.begin_create_or_update = AsyncMock( + side_effect=azure_client.resource_client.resources.begin_create_or_update + ) + async_resource.resources.get = AsyncMock(side_effect=azure_client.resource_client.resources.get) + + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource) + return azure_client + + +def make_vmss_azure_client() -> MagicMock: + """Build the VMSS handler's async-over-sync Azure client double.""" + azure_client = MagicMock() + async_compute = MagicMock() + + async_compute.virtual_machine_scale_sets.begin_create_or_update = AsyncMock( + side_effect=azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update + ) + async_compute.virtual_machine_scale_sets.get = AsyncMock( + side_effect=azure_client.compute_client.virtual_machine_scale_sets.get + ) + async_compute.virtual_machine_scale_sets.begin_delete_instances = AsyncMock( + side_effect=azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances + ) + async_compute.virtual_machine_scale_sets.begin_delete = AsyncMock( + side_effect=azure_client.compute_client.virtual_machine_scale_sets.begin_delete + ) + async_compute.virtual_machine_scale_set_vms.list = MagicMock( + side_effect=lambda *args, **kwargs: AsyncPager( + azure_client.compute_client.virtual_machine_scale_set_vms.list(*args, **kwargs) + ) + ) + async_compute.virtual_machines.list = MagicMock( + side_effect=lambda *args, **kwargs: AsyncPager( + azure_client.compute_client.virtual_machines.list(*args, **kwargs) + ) + ) + async_compute.virtual_machines.begin_delete = AsyncMock( + side_effect=azure_client.compute_client.virtual_machines.begin_delete + ) + async_compute.ssh_public_keys.get = AsyncMock( + side_effect=azure_client.compute_client.ssh_public_keys.get + ) + + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + return azure_client + + +def run_operation(coro): + """Run a coroutine in a fresh event loop for synchronous tests.""" + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + +def build_strategy_harness( + *, + config, + logger, + provider_instance_name: str = "azure-default", +) -> AzureStrategyHarness: + """Build a strategy plus mutable dependency holders for focused tests.""" + harness = AzureStrategyHarness() + handler_factory = _TestAzureHandlerFactory(harness) + harness.strategy = AzureProviderStrategy( + config=config, + logger=logger, + provider_instance_name=provider_instance_name, + azure_client_resolver=lambda: harness.azure_client, + azure_handler_factory_resolver=lambda: handler_factory, + azure_resource_manager_resolver=lambda: harness.resource_manager, + azure_deployment_service_resolver=lambda: harness.deployment_service, + ) + harness.strategy.initialize() + return harness diff --git a/tests/providers/azure/test_async_azure_services.py b/tests/providers/azure/test_async_azure_services.py new file mode 100644 index 000000000..bbd2afd59 --- /dev/null +++ b/tests/providers/azure/test_async_azure_services.py @@ -0,0 +1,665 @@ +"""Direct async coverage for Azure async service entry points.""" + +import asyncio +from types import SimpleNamespace +from typing import TYPE_CHECKING, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError +from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, +) +from orb.providers.azure.services.health_check_service import AzureHealthCheckService +from orb.providers.azure.services.inventory_service import ( + AzureInventoryService, + AzureReadOperationContext, +) +from orb.providers.azure.services.provisioning_service import ( + AzureProvisioningService, + CreateOperationContext, +) +from orb.providers.azure.services.termination_service import ( + AzureTerminationService, +) +from orb.providers.base.strategy import ProviderOperation, ProviderOperationType +from tests.providers.azure.strategy_test_support import make_azure_template + +if TYPE_CHECKING: + from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy + + +def _make_template(**overrides): + return make_azure_template( + template_id="azure-service-test", + provider_api="SingleVM", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABg service@test"], + **overrides, + ) + + +def _handler_provider(handler) -> "AzureProviderStrategy": + """Build a strategy stand-in that returns ``handler`` from any ``resolve_handler`` call. + + Cast to ``AzureProviderStrategy`` so the test fixture satisfies the inventory + service's annotated dependency without spinning up a real strategy. Duck-typed + at runtime via ``SimpleNamespace.resolve_handler``. + """ + return cast( + "AzureProviderStrategy", + SimpleNamespace(resolve_handler=lambda *_args, **_kwargs: handler), + ) + + +def _termination_operation( + *, + instance_ids: list[str], + resource_mapping: dict[str, tuple[str, int]] | None = None, + resource_id: str | None = None, +) -> ProviderOperation: + parameters = { + "instance_ids": instance_ids, + "provider_api": "VMSS", + } + if resource_mapping is not None: + parameters["resource_mapping"] = resource_mapping + if resource_id is not None: + parameters["resource_id"] = resource_id + + return ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters=parameters, + ) + + +async def _terminate_with_handler( + handler, + *, + instance_ids: list[str], + resource_mapping: dict[str, tuple[str, int]] | None = None, + resource_id: str | None = None, + logger: MagicMock | None = None, + cleanup_recorder: MagicMock | None = None, +): + service = AzureTerminationService( + logger=logger or MagicMock(), + handler_provider=_handler_provider(handler), + record_pending_cleanup=cleanup_recorder or MagicMock(), + default_resource_group="test-rg", + ) + return await service.terminate_instances_async( + _termination_operation( + instance_ids=instance_ids, + resource_mapping=resource_mapping, + resource_id=resource_id, + ), + is_dry_run=False, + ) + + +@pytest.fixture +def resource_metadata_service() -> MagicMock: + service = MagicMock() + service.augment_vmss_capacity_metadata_async = AsyncMock() + service.augment_single_vm_deployment_metadata_async = AsyncMock() + service.attach_provider_fulfilment.side_effect = ( + lambda _metadata, *, instances, target_units=None: SimpleNamespace(instances=instances) + ) + return service + + +@pytest.mark.asyncio +async def test_dispatch_collects_provider_data_and_records_cleanup(): + cleanup_recorder = MagicMock() + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + return_value={"provider_data": {"operation_status": "submitted"}} + ) + + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1"], + resource_mapping={"vm-1": ("vmss-1", 1)}, + cleanup_recorder=cleanup_recorder, + ) + + assert result.success is True + assert result.metadata["provider_data"]["termination_requests"] == [ + {"operation_status": "submitted"} + ] + cleanup_recorder.assert_called_once() + + +@pytest.mark.asyncio +async def test_dispatch_fans_out_across_multiple_resource_groups(): + cleanup_recorder = MagicMock() + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=[ + {"provider_data": {"resource_id": "vmss-1"}}, + {"provider_data": {"resource_id": "vmss-2"}}, + ] + ) + + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1", "vm-2"], + resource_mapping={"vm-1": ("vmss-1", 1), "vm-2": ("vmss-2", 1)}, + cleanup_recorder=cleanup_recorder, + ) + + assert result.success is True + assert result.metadata["provider_data"]["termination_requests"] == [ + {"resource_id": "vmss-1"}, + {"resource_id": "vmss-2"}, + ] + assert handler.release_hosts_async.await_count == 2 + cleanup_recorder.assert_called() + + +@pytest.mark.asyncio +async def test_dispatch_preserves_group_order_when_release_calls_complete_out_of_order(): + cleanup_recorder = MagicMock() + handler = MagicMock() + + async def release_hosts_async(*, machine_ids, resource_id, context): + _ = machine_ids, context + if resource_id == "vmss-1": + await asyncio.sleep(0.02) + return {"provider_data": {"resource_id": resource_id}} + + handler.release_hosts_async = AsyncMock(side_effect=release_hosts_async) + + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1", "vm-2"], + resource_mapping={"vm-1": ("vmss-1", 1), "vm-2": ("vmss-2", 1)}, + cleanup_recorder=cleanup_recorder, + ) + + assert result.success is True + assert result.metadata["provider_data"]["termination_requests"] == [ + {"resource_id": "vmss-1"}, + {"resource_id": "vmss-2"}, + ] + assert cleanup_recorder.call_args_list[0].args[0]["provider_data"]["resource_id"] == "vmss-1" + assert cleanup_recorder.call_args_list[1].args[0]["provider_data"]["resource_id"] == "vmss-2" + + +@pytest.mark.asyncio +async def test_dispatch_reports_partial_failure_when_one_group_fails(): + cleanup_recorder = MagicMock() + logger = MagicMock() + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=[ + {"provider_data": {"resource_id": "vmss-1"}}, + RuntimeError("delete failed"), + ] + ) + + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1", "vm-2"], + resource_mapping={"vm-1": ("vmss-1", 1), "vm-2": ("vmss-2", 1)}, + logger=logger, + cleanup_recorder=cleanup_recorder, + ) + + assert result.success is False + assert result.error_code == "AZURE_TERMINATION_PARTIAL_FAILURE" + provider_data = result.metadata["provider_data"] + assert provider_data["termination_requests"] == [{"resource_id": "vmss-1"}] + assert provider_data["successful_instance_ids"] == ["vm-1"] + assert provider_data["dispatch_failures"] == [ + { + "resource_id": "vmss-2", + "instance_ids": ["vm-2"], + "error_message": "delete failed", + "error_type": "RuntimeError", + } + ] + cleanup_recorder.assert_called_once() + logger.warning.assert_called() + + +@pytest.mark.asyncio +async def test_dispatch_treats_success_without_provider_data_as_successful_group(): + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=[ + {}, + RuntimeError("delete failed"), + ] + ) + + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1", "vm-2"], + resource_mapping={"vm-1": ("vmss-1", 1), "vm-2": ("vmss-2", 1)}, + ) + + assert result.success is False + provider_data = result.metadata["provider_data"] + assert "termination_requests" not in provider_data + assert provider_data["successful_instance_ids"] == ["vm-1"] + assert provider_data["failed_instance_ids"] == ["vm-2"] + + +@pytest.mark.asyncio +async def test_terminate_instances_async_reports_partial_dispatch_failure(): + cleanup_recorder = MagicMock() + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=[ + {"provider_data": {"resource_id": "vmss-1"}}, + RuntimeError("delete failed"), + ] + ) + result = await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1", "vm-2"], + resource_mapping={"vm-1": ("vmss-1", 1), "vm-2": ("vmss-2", 1)}, + cleanup_recorder=cleanup_recorder, + ) + + assert result.success is False + assert result.error_code == "AZURE_TERMINATION_PARTIAL_FAILURE" + provider_data = result.metadata["provider_data"] + assert provider_data["termination_requests"] == [{"resource_id": "vmss-1"}] + assert provider_data["successful_instance_ids"] == ["vm-1"] + assert provider_data["failed_instance_ids"] == ["vm-2"] + assert provider_data["dispatch_failures"] == [ + { + "resource_id": "vmss-2", + "instance_ids": ["vm-2"], + "error_message": "delete failed", + "error_type": "RuntimeError", + } + ] + + +@pytest.mark.asyncio +async def test_dispatch_raises_when_all_groups_fail(): + handler = MagicMock() + handler.release_hosts_async = AsyncMock(side_effect=RuntimeError("delete failed")) + + with pytest.raises(RuntimeError, match="delete failed"): + await _terminate_with_handler( + handler=handler, + instance_ids=["vm-1"], + resource_mapping={"vm-1": ("vmss-1", 1)}, + ) + + +@pytest.mark.asyncio +async def test_execute_create_handler_async_normalizes_handler_result(): + service = AzureProvisioningService() + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["vm-1"], + "instances": [], + "error_message": None, + "provider_data": {"operation_status": "submitted"}, + } + ) + create_context = CreateOperationContext( + template_config={"provider_api": "SingleVM"}, + count=1, + provider_api=AzureProviderApi.SINGLE_VM, + provider_api_key="SingleVM", + handler=handler, + azure_template=_make_template(), + ) + + result = await service.execute_create_handler_async( + create_context=create_context, + request=MagicMock(), + ) + + assert result.success is True + assert result.data["resource_ids"] == ["vm-1"] + assert result.metadata["handler_used"] == "SingleVM" + + +@pytest.mark.asyncio +async def test_execute_create_handler_async_returns_provider_error_for_failed_handler_result(): + service = AzureProvisioningService() + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": False, + "resource_ids": [], + "instances": [], + "error_message": "quota exhausted", + "provider_data": {"error_codes": ["AllocationFailed"]}, + } + ) + create_context = CreateOperationContext( + template_config={"provider_api": "SingleVM"}, + count=1, + provider_api=AzureProviderApi.SINGLE_VM, + provider_api_key="SingleVM", + handler=handler, + azure_template=_make_template(), + ) + + result = await service.execute_create_handler_async( + create_context=create_context, + request=MagicMock(), + ) + + assert result.success is False + assert result.error_code == "PROVISIONING_ADAPTER_ERROR" + assert result.error_message == "Provisioning failed: quota exhausted" + assert result.metadata["provider_data"]["error_codes"] == ["AllocationFailed"] + + +@pytest.mark.asyncio +async def test_get_instance_status_async_uses_async_handler_dispatch(resource_metadata_service): + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[{"instance_id": "vm-1", "provider_data": {"vm_name": "vm-1"}}] + ) + service = AzureInventoryService( + logger=MagicMock(), + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(handler), + vmss_cleanup_coordinator=MagicMock(), + ) + read_context = AzureReadOperationContext( + operation_name="get_instance_status", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={}, + cyclecloud_request_context=MagicMock(), + provider_api=AzureProviderApi.SINGLE_VM, + provider_api_key="SingleVM", + resource_group="test-rg", + instance_ids=["vm-1"], + resource_ids=["vm-1"], + direct_resource_id="vm-1", + ) + + result = await service.get_instance_status_async(read_context) + + assert result.success is True + assert result.data["instances"][0]["instance_id"] == "vm-1" + handler.check_hosts_status_async.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_instance_status_async_rejects_queries_without_provider_api( + resource_metadata_service, +): + service = AzureInventoryService( + logger=MagicMock(), + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(None), + vmss_cleanup_coordinator=MagicMock(), + ) + read_context = AzureReadOperationContext( + operation_name="get_instance_status", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={}, + cyclecloud_request_context=MagicMock(), + provider_api=None, + provider_api_key=None, + resource_group="test-rg", + instance_ids=["vm-1"], + ) + with pytest.raises(AzureValidationError, match="provider_api-backed handler resolution"): + await service.get_instance_status_async(read_context) + + +@pytest.mark.asyncio +async def test_describe_resource_instances_async_builds_result_from_async_handler( + resource_metadata_service, +): + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[{"instance_id": "vm-1", "provider_data": {"vm_name": "vm-1"}}] + ) + service = AzureInventoryService( + logger=MagicMock(), + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(handler), + vmss_cleanup_coordinator=MagicMock(), + ) + read_context = AzureReadOperationContext( + operation_name="describe_resource_instances", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={"deployment_name": "dep-1"}, + cyclecloud_request_context=MagicMock(), + provider_api=AzureProviderApi.SINGLE_VM, + provider_api_key="SingleVM", + resource_group="test-rg", + resource_ids=["vm-1"], + ) + + result = await service.describe_resource_instances_async( + read_context=read_context, + resource_manager=None, + deployment_service=None, + ) + + assert result.success is True + assert result.data["instances"][0]["instance_id"] == "vm-1" + resource_metadata_service.augment_shortfall_metadata.assert_called_once() + resource_metadata_service.augment_single_vm_deployment_metadata_async.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_submit_template_deployment_async_uses_async_resources_api(): + azure_client = MagicMock() + async_resource_client = MagicMock() + async_resource_client.resources.begin_create_or_update = AsyncMock() + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource_client) + service = AzureDeploymentService(azure_client=azure_client, logger=MagicMock()) + + deployment_name = await service.submit_template_deployment_async( + resource_group="test-rg", + deployment_name="dep-1", + template={"resources": []}, + parameters={"vmName": {"value": "vm-1"}}, + ) + + assert deployment_name == "dep-1" + async_resource_client.resources.begin_create_or_update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_deployment_status_async_extracts_provisioning_and_error_fields(): + azure_client = MagicMock() + async_resource_client = MagicMock() + async_resource_client.resources.get = AsyncMock( + return_value=SimpleNamespace( + properties={ + "provisioningState": "Failed", + "error": { + "code": "DeploymentFailed", + "message": "template validation failed", + }, + } + ) + ) + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource_client) + service = AzureDeploymentService(azure_client=azure_client, logger=MagicMock()) + + status = await service.get_deployment_status_async( + resource_group="test-rg", + deployment_name="dep-1", + ) + + assert status == { + "provisioning_state": "Failed", + "error_code": "DeploymentFailed", + "error_message": "template validation failed", + } + + +@pytest.mark.asyncio +async def test_describe_resource_instances_async_awaits_vmss_async_metadata_and_cleanup( + resource_metadata_service, +): + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock( + return_value=[{"error_code": "ProvisioningStateFailed"}] + ) + cleanup = MagicMock() + cleanup.reconcile = AsyncMock() + service = AzureInventoryService( + logger=MagicMock(), + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(handler), + vmss_cleanup_coordinator=cleanup, + ) + read_context = AzureReadOperationContext( + operation_name="describe_resource_instances", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={}, + cyclecloud_request_context=MagicMock(), + provider_api=AzureProviderApi.VMSS, + provider_api_key="VMSS", + resource_group="test-rg", + resource_ids=["vmss-1"], + ) + + result = await service.describe_resource_instances_async( + read_context=read_context, + resource_manager=MagicMock(), + deployment_service=None, + ) + + assert result.success is True + assert result.metadata["fleet_errors"] == [{"error_code": "ProvisioningStateFailed"}] + handler.get_vmss_resource_errors_async.assert_awaited_once_with("test-rg", "vmss-1") + resource_metadata_service.augment_vmss_capacity_metadata_async.assert_awaited_once() + cleanup.reconcile.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_describe_resource_instances_async_uses_empty_vmss_error_list_without_override( + resource_metadata_service, +): + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + service = AzureInventoryService( + logger=MagicMock(), + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(handler), + vmss_cleanup_coordinator=MagicMock(reconcile=AsyncMock()), + ) + read_context = AzureReadOperationContext( + operation_name="describe_resource_instances", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={}, + cyclecloud_request_context=MagicMock(), + provider_api=AzureProviderApi.VMSS, + provider_api_key="VMSS", + resource_group="test-rg", + resource_ids=["vmss-1"], + ) + + result = await service.describe_resource_instances_async( + read_context=read_context, + resource_manager=MagicMock(), + deployment_service=None, + ) + + assert result.success is True + assert "fleet_errors" not in result.metadata + handler.get_vmss_resource_errors_async.assert_awaited_once_with("test-rg", "vmss-1") + + +@pytest.mark.asyncio +async def test_describe_resource_instances_async_warns_when_vmss_handler_lacks_error_reader( + resource_metadata_service, +): + logger = MagicMock() + handler = MagicMock(spec=["check_hosts_status_async"]) + handler.check_hosts_status_async = AsyncMock(return_value=[]) + service = AzureInventoryService( + logger=logger, + provider_instance_name="azure-default", + resource_metadata_service=resource_metadata_service, + handler_provider=_handler_provider(handler), + vmss_cleanup_coordinator=MagicMock(reconcile=AsyncMock()), + ) + read_context = AzureReadOperationContext( + operation_name="describe_resource_instances", + request_id="req-12345678-1234-1234-1234-123456789012", + template_id="tpl-1", + request_metadata={}, + cyclecloud_request_context=MagicMock(), + provider_api=AzureProviderApi.VMSS, + provider_api_key="VMSS", + resource_group="test-rg", + resource_ids=["vmss-1"], + ) + + result = await service.describe_resource_instances_async( + read_context=read_context, + resource_manager=MagicMock(), + deployment_service=None, + ) + + assert result.success is True + assert "fleet_errors" not in result.metadata + logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_health_check_async_uses_short_lived_async_token_provider(monkeypatch): + config = MagicMock() + config.region = "eastus2" + config.client_id = "client-id" + service = AzureHealthCheckService(config=config, logger=MagicMock()) + provider = MagicMock() + provider.get_access_token = AsyncMock(return_value="token") + provider_cls = MagicMock(return_value=provider) + monkeypatch.setattr( + "orb.providers.azure.services.health_check_service.AsyncDefaultAzureAccessTokenProvider", + provider_cls, + ) + + result = await service.check_health_async() + + assert result.is_healthy is True + provider_cls.assert_called_once_with(client_id="client-id", logger=service._logger) + provider.get_access_token.assert_awaited_once_with("https://management.azure.com/.default") + + +def test_health_check_sync_uses_short_lived_sync_token_provider(monkeypatch): + config = MagicMock() + config.region = "eastus2" + config.client_id = "client-id" + service = AzureHealthCheckService(config=config, logger=MagicMock()) + provider = MagicMock() + provider.get_access_token.return_value = "token" + provider_cls = MagicMock(return_value=provider) + monkeypatch.setattr( + "orb.providers.azure.services.health_check_service.DefaultAzureAccessTokenProvider", + provider_cls, + ) + + result = service.check_health() + + assert result.is_healthy is True + provider_cls.assert_called_once_with(client_id="client-id", logger=service._logger) + provider.get_access_token.assert_called_once_with("https://management.azure.com/.default") diff --git a/tests/providers/azure/test_azure_client.py b/tests/providers/azure/test_azure_client.py new file mode 100644 index 000000000..b72960b0a --- /dev/null +++ b/tests/providers/azure/test_azure_client.py @@ -0,0 +1,631 @@ +"""Focused tests for Azure client and auth behavior.""" + +import asyncio +import sys +import threading +import types +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from orb.config import PerformanceConfig +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.exceptions.azure_exceptions import ( + AuthenticationError, + AzureConfigurationError, +) +from orb.providers.azure.infrastructure.azure_client import ( + AzureClient, + AzureClientRuntimeConfig, +) +from orb.providers.azure.infrastructure.services.arm_resource_id_parser import ( + ArmResourceIdParser, +) +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentityResolver, +) + + +class TestAzureAuthStrategy: + @staticmethod + def _build_auth_context(): + from orb.infrastructure.adapters.ports.auth import AuthContext + + return AuthContext( + method="GET", + path="/providers/azure/health", + headers={}, + query_params={}, + ) + + @pytest.mark.asyncio + async def test_auth_strategy_returns_failed_result_for_expected_azure_auth_error(self): + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + + token_provider = MagicMock() + strategy = AzureAuthStrategy(logger=MagicMock(), token_provider=token_provider) + token_provider.get_auth_error_types.return_value = (RuntimeError,) + token_provider.get_access_token.side_effect = RuntimeError("credential unavailable") + + result = await strategy.authenticate(self._build_auth_context()) + + assert result.status.name == "FAILED" + assert "credential unavailable" in result.error_message + + @pytest.mark.asyncio + async def test_auth_strategy_propagates_unexpected_errors(self): + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + + token_provider = MagicMock() + strategy = AzureAuthStrategy(logger=MagicMock(), token_provider=token_provider) + token_provider.get_auth_error_types.return_value = (ValueError,) + token_provider.get_access_token.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + await strategy.authenticate(self._build_auth_context()) + + @pytest.mark.asyncio + async def test_auth_strategy_delegates_to_token_provider(self): + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + + token_provider = MagicMock() + token_provider.get_auth_error_types.return_value = (RuntimeError,) + token_provider.get_access_token.return_value = "access-token" + strategy = AzureAuthStrategy(logger=MagicMock(), token_provider=token_provider) + + result = await strategy.authenticate(self._build_auth_context()) + + assert result.status.name == "SUCCESS" + assert result.token == "access-token" + assert result.permissions == [] + token_provider.get_access_token.assert_called_once_with( + "https://management.azure.com/.default" + ) + + @pytest.mark.asyncio + async def test_auth_strategy_awaits_async_token_provider(self): + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + + async_token_provider = MagicMock() + async_token_provider.get_auth_error_types.return_value = (RuntimeError,) + async_token_provider.get_access_token = AsyncMock(return_value="access-token") + strategy = AzureAuthStrategy( + logger=MagicMock(), + async_token_provider=async_token_provider, + ) + + result = await strategy.authenticate(self._build_auth_context()) + + assert result.status.name == "SUCCESS" + assert result.token == "access-token" + async_token_provider.get_access_token.assert_awaited_once_with( + "https://management.azure.com/.default" + ) + + +class TestAzureClientOperationalBehavior: + @staticmethod + def _build_runtime_config( + azure_config: AzureProviderConfig | None = None, + perf_config: PerformanceConfig | None = None, + ) -> AzureClientRuntimeConfig: + return AzureClientRuntimeConfig( + azure_config=azure_config + or AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + ), + performance_config=perf_config or PerformanceConfig(), + ) + + @classmethod + def _build_client(cls, logger: MagicMock | None = None) -> AzureClient: + return AzureClient( + runtime_config=cls._build_runtime_config(), + logger=logger or MagicMock(), + ) + + @staticmethod + def _build_partial_client() -> AzureClient: + client = object.__new__(AzureClient) + client._logger = MagicMock() + client._lazy_init_lock = threading.RLock() + client._closed = False + client._credentials_validated = False + client._async_credential = None + client._async_compute_client = None + client._async_network_client = None + client._async_resource_client = None + client._async_subscription_client = None + client._pending_async_close_task = None + client._arm_resource_id_parser = ArmResourceIdParser() + client._network_identity_resolver = AzureNetworkIdentityResolver( + async_network_client_getter=client.get_async_network_client, + logger=client._logger, + arm_resource_id_parser=client._arm_resource_id_parser, + network_lookup_error_types=client._network_lookup_error_types, + ) + return client + + def test_azure_client_uses_explicit_runtime_config(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + ) + client = AzureClient( + runtime_config=self._build_runtime_config(azure_config=azure_config), + logger=MagicMock(), + ) + + assert client.subscription_id == azure_config.subscription_id + assert client.resource_group == "rg-explicit" + assert client.region_name == "westeurope" + + @pytest.mark.asyncio + async def test_azure_client_passes_managed_identity_client_id_when_configured(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + client_id="managed-identity-client-id", + ) + client = AzureClient( + runtime_config=self._build_runtime_config(azure_config=azure_config), + logger=MagicMock(), + ) + + fake_ctor = MagicMock(return_value=MagicMock()) + + with patch( + "orb.providers.azure.infrastructure.azure_client.create_default_azure_credential_async", + fake_ctor, + ): + await client.get_async_credential() + + fake_ctor.assert_called_once_with( + client_id="managed-identity-client-id", + logger=client._logger, + ) + + @pytest.mark.asyncio + async def test_azure_client_omits_managed_identity_client_id_when_unset(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + ) + client = AzureClient( + runtime_config=self._build_runtime_config(azure_config=azure_config), + logger=MagicMock(), + ) + + fake_ctor = MagicMock(return_value=MagicMock()) + + with patch( + "orb.providers.azure.infrastructure.azure_client.create_default_azure_credential_async", + fake_ctor, + ): + await client.get_async_credential() + + fake_ctor.assert_called_once_with( + client_id=None, + logger=client._logger, + ) + + @pytest.mark.asyncio + async def test_get_async_credential_preserves_nested_import_error_details(self): + client = self._build_client() + + with patch( + "orb.providers.azure.infrastructure.azure_client.create_default_azure_credential_async", + side_effect=ImportError("aiohttp package is not installed"), + ): + with pytest.raises( + AuthenticationError, + match="azure-identity dependency error: aiohttp package is not installed", + ): + await client.get_async_credential() + + @pytest.mark.asyncio + async def test_azure_client_passes_retry_and_timeout_kwargs_to_compute_client(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + max_retries=7, + connect_timeout=11, + read_timeout=22, + ) + client = AzureClient( + runtime_config=self._build_runtime_config(azure_config=azure_config), + logger=MagicMock(), + ) + fake_credential = MagicMock() + client.get_async_credential = AsyncMock(return_value=fake_credential) + + fake_compute_module = types.ModuleType("azure.mgmt.compute.aio") + fake_ctor = MagicMock(return_value=MagicMock()) + fake_compute_module.ComputeManagementClient = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.mgmt": types.ModuleType("azure.mgmt"), + "azure.mgmt.compute.aio": fake_compute_module, + }, + ): + await client.get_async_compute_client() + + fake_ctor.assert_called_once_with( + credential=fake_credential, + subscription_id="12345678-1234-1234-1234-123456789012", + retry_total=7, + connection_timeout=11, + read_timeout=22, + ) + + @pytest.mark.asyncio + async def test_azure_client_passes_retry_and_timeout_kwargs_to_subscription_client(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + max_retries=4, + connect_timeout=9, + read_timeout=19, + ) + client = AzureClient( + runtime_config=self._build_runtime_config(azure_config=azure_config), + logger=MagicMock(), + ) + fake_credential = MagicMock() + client.get_async_credential = AsyncMock(return_value=fake_credential) + + fake_subscription_module = types.ModuleType("azure.mgmt.resource.subscriptions.aio") + fake_ctor = MagicMock(return_value=MagicMock()) + fake_subscription_module.SubscriptionClient = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.mgmt": types.ModuleType("azure.mgmt"), + "azure.mgmt.resource": types.ModuleType("azure.mgmt.resource"), + "azure.mgmt.resource.subscriptions.aio": fake_subscription_module, + }, + ): + await client.get_async_subscription_client() + + fake_ctor.assert_called_once_with( + credential=fake_credential, + retry_total=4, + connection_timeout=9, + read_timeout=19, + ) + + def test_collect_error_types_deduplicates_base_and_imported_types(self): + error_types = AzureClient._collect_error_types( + AuthenticationError, + AuthenticationError, + optional_error_loaders=(), + ) + + assert error_types == (AuthenticationError,) + + def test_azure_client_rejects_non_integer_timeout_values(self): + with pytest.raises( + AzureConfigurationError, + match=r"connect_timeout.*must be an integer", + ): + AzureClient( + runtime_config=AzureClientRuntimeConfig( + azure_config=types.SimpleNamespace( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + max_retries=3, + connect_timeout="fast", + read_timeout=22, + ), + performance_config=PerformanceConfig(), + ), + logger=MagicMock(), + ) + + def test_azure_client_rejects_timeout_values_below_minimum(self): + with pytest.raises( + AzureConfigurationError, + match=r"connect_timeout.*must be >= 1", + ): + AzureClient( + runtime_config=AzureClientRuntimeConfig( + azure_config=types.SimpleNamespace( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + max_retries=3, + connect_timeout=0, + read_timeout=22, + ), + performance_config=PerformanceConfig(), + ), + logger=MagicMock(), + ) + + def test_azure_client_maps_typed_performance_config(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + ) + perf_config = PerformanceConfig( + enable_batching=False, + enable_parallel=False, + max_workers=3, + caching={"request_status": {"enabled": False, "ttl_seconds": 42}}, + ) + client = AzureClient( + runtime_config=self._build_runtime_config( + azure_config=azure_config, + perf_config=perf_config, + ), + logger=MagicMock(), + ) + + assert client.perf_config["enable_batching"] is False + assert client.perf_config["enable_parallel"] is False + assert client.perf_config["max_workers"] == 3 + assert client.perf_config["enable_caching"] is False + assert client.perf_config["cache_ttl"] == 42 + assert client.perf_config["batch_sizes"] == { + "deallocate_vms": 25, + "create_tags": 20, + "describe_vms": 25, + } + + def test_azure_client_uses_default_performance_config_when_not_overridden(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg-explicit", + region="westeurope", + ) + client = AzureClient( + runtime_config=self._build_runtime_config( + azure_config=azure_config, + perf_config=PerformanceConfig(), + ), + logger=MagicMock(), + ) + + assert client.perf_config["enable_batching"] is True + assert client.perf_config["enable_parallel"] is True + assert client.perf_config["max_workers"] == 10 + assert client.perf_config["enable_caching"] is True + assert client.perf_config["cache_ttl"] == 300 + + @pytest.mark.asyncio + async def test_validate_credentials_async_returns_true_and_sets_validation_flag(self): + client = self._build_partial_client() + async_credential = MagicMock() + async_credential.get_token = AsyncMock(return_value=MagicMock()) + client.get_async_credential = AsyncMock(return_value=async_credential) + + assert await AzureClient.validate_credentials_async(client) is True + async_credential.get_token.assert_awaited_once_with("https://management.azure.com/.default") + assert client._credentials_validated is True + + @pytest.mark.asyncio + async def test_validate_subscription_async_returns_false_for_known_azure_errors(self): + from azure.core.exceptions import ResourceNotFoundError + + client = self._build_partial_client() + client.subscription_id = "12345678-1234-1234-1234-123456789012" + async_subscription_client = MagicMock() + async_subscription_client.subscriptions.get = AsyncMock( + side_effect=ResourceNotFoundError("missing") + ) + client.get_async_subscription_client = AsyncMock(return_value=async_subscription_client) + + assert await AzureClient.validate_subscription_async(client) is False + + @pytest.mark.asyncio + async def test_validate_credentials_async_returns_false_for_authentication_error(self): + client = self._build_partial_client() + async_credential = MagicMock() + async_credential.get_token = AsyncMock(side_effect=AuthenticationError("bad credential")) + client.get_async_credential = AsyncMock(return_value=async_credential) + + assert await AzureClient.validate_credentials_async(client) is False + + @pytest.mark.asyncio + async def test_validate_credentials_async_reraises_unexpected_errors(self): + client = self._build_partial_client() + async_credential = MagicMock() + async_credential.get_token = AsyncMock(side_effect=RuntimeError("boom")) + client.get_async_credential = AsyncMock(return_value=async_credential) + + with pytest.raises(RuntimeError, match="boom"): + await AzureClient.validate_credentials_async(client) + + @pytest.mark.asyncio + async def test_validate_subscription_async_reraises_unexpected_errors(self): + client = self._build_partial_client() + client.subscription_id = "12345678-1234-1234-1234-123456789012" + async_subscription_client = MagicMock() + async_subscription_client.subscriptions.get = AsyncMock(side_effect=RuntimeError("boom")) + client.get_async_subscription_client = AsyncMock(return_value=async_subscription_client) + + with pytest.raises(RuntimeError, match="boom"): + await AzureClient.validate_subscription_async(client) + + def test_close_marks_client_closed_and_prevents_reuse(self): + logger = MagicMock() + client = self._build_client(logger=logger) + client._credentials_validated = True + client._closed = False + + AzureClient.close(client) + + assert client._credentials_validated is False + assert client._closed is True + + with pytest.raises(RuntimeError, match="AzureClient has been closed"): + client._ensure_open() + + def test_close_is_idempotent(self): + client = self._build_client() + client._credentials_validated = False + client._closed = False + + AzureClient.close(client) + AzureClient.close(client) + + @pytest.mark.asyncio + async def test_close_schedules_async_cleanup_when_event_loop_is_running(self): + logger = MagicMock() + client = self._build_client(logger=logger) + async_compute_client = MagicMock() + async_compute_client.close = AsyncMock() + client._async_compute_client = async_compute_client + client._closed = False + + AzureClient.close(client) + await asyncio.sleep(0) + await asyncio.sleep(0) + + async_compute_client.close.assert_awaited_once() + assert client._pending_async_close_task is None + + @pytest.mark.asyncio + async def test_aclose_raises_first_error_after_attempting_all_async_resources(self): + client = self._build_partial_client() + async_compute_client = MagicMock() + async_network_client = MagicMock() + async_compute_client.close = AsyncMock(side_effect=RuntimeError("compute close failed")) + async_network_client.close = AsyncMock(side_effect=RuntimeError("network close failed")) + client._async_compute_client = async_compute_client + client._async_network_client = async_network_client + + with pytest.raises(RuntimeError, match="network close failed"): + await AzureClient.aclose(client) + + async_compute_client.close.assert_awaited_once() + async_network_client.close.assert_awaited_once() + assert client._async_compute_client is None + assert client._async_network_client is None + + @pytest.mark.asyncio + async def test_get_async_compute_client_returns_live_client_when_other_task_wins_race(self): + client = self._build_partial_client() + existing_client = MagicMock() + created_client = MagicMock() + created_client.close = AsyncMock() + client._async_compute_client = None + + async def build_management_client_async(**kwargs): + _ = kwargs + client._async_compute_client = existing_client + return created_client + + client._build_management_client_async = AsyncMock(side_effect=build_management_client_async) + client._close_async_resource = AsyncMock() + + resolved = await client.get_async_compute_client() + + assert resolved is existing_client + client._close_async_resource.assert_awaited_once_with( + "async_compute_client", + created_client, + ) + + def test_context_manager_closes_on_exit(self): + client = self._build_client() + client._credentials_validated = False + client._closed = False + + with client as scoped_client: + assert scoped_client is client + + assert client._closed is True + + +class TestAzureClientNetworkResolution: + @staticmethod + def _build_partial_client() -> AzureClient: + client = object.__new__(AzureClient) + client._logger = MagicMock() + client._lazy_init_lock = threading.RLock() + client._closed = False + client._credentials_validated = False + client._async_credential = None + client._async_compute_client = None + client._async_network_client = None + client._async_resource_client = None + client._async_subscription_client = None + client._pending_async_close_task = None + client._arm_resource_id_parser = ArmResourceIdParser() + client._network_identity_resolver = AzureNetworkIdentityResolver( + async_network_client_getter=client.get_async_network_client, + logger=client._logger, + arm_resource_id_parser=client._arm_resource_id_parser, + network_lookup_error_types=client._network_lookup_error_types, + ) + return client + + @pytest.mark.asyncio + async def test_resolve_network_identity_from_vm_async_populates_ips_and_subnet(self): + azure_client = self._build_partial_client() + async_network_client = MagicMock() + async_network_client.network_interfaces.get = AsyncMock() + async_network_client.public_ip_addresses.get = AsyncMock() + azure_client.get_async_network_client = AsyncMock(return_value=async_network_client) + azure_client._network_identity_resolver = AzureNetworkIdentityResolver( + async_network_client_getter=azure_client.get_async_network_client, + logger=azure_client._logger, + arm_resource_id_parser=azure_client._arm_resource_id_parser, + network_lookup_error_types=azure_client._network_lookup_error_types, + ) + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + vm = MagicMock() + vm.network_profile.network_interfaces = [nic_ref] + + subnet = MagicMock() + subnet.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) + public_ip_ref = MagicMock() + public_ip_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/publicIPAddresses/pip-vm-1" + ) + ip_cfg = MagicMock() + ip_cfg.private_ip_address = "10.0.0.4" + ip_cfg.subnet = subnet + ip_cfg.public_ip_address = public_ip_ref + + nic = MagicMock() + nic.ip_configurations = [ip_cfg] + async_network_client.network_interfaces.get.return_value = nic + + pip = MagicMock() + pip.ip_address = "52.1.2.3" + async_network_client.public_ip_addresses.get.return_value = pip + + result = await AzureClient.resolve_network_identity_from_vm_async(azure_client, vm) + + assert result["private_ip"] == "10.0.0.4" + assert result["public_ip"] == "52.1.2.3" + assert result["subnet_id"].endswith("/subnets/default") + assert result["vnet_id"].endswith("/virtualNetworks/test-vnet") + assert result["nic_name"] == "nic-vm-1" diff --git a/tests/providers/azure/test_azure_config_and_infra.py b/tests/providers/azure/test_azure_config_and_infra.py new file mode 100644 index 000000000..1d5bdaf17 --- /dev/null +++ b/tests/providers/azure/test_azure_config_and_infra.py @@ -0,0 +1,640 @@ +"""Tests for Azure configuration and registration behavior.""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest +from pydantic import ValidationError + +from orb.application.dto.template import TemplateDTO +from orb.bootstrap.infrastructure_services import register_infrastructure_services +from orb.config import PerformanceConfig +from orb.domain.base.ports.logging_port import LoggingPort +from orb.domain.template.factory import TemplateFactory +from orb.infrastructure.di.container import DIContainer +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.configuration.template_extension import AzureTemplateExtensionConfig +from orb.providers.azure.configuration.validator import ( + validate_azure_config, + validate_azure_template, +) +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.infrastructure.azure_client import AzureClient +from orb.providers.azure.infrastructure.azure_handler_factory import AzureHandlerFactory +from orb.providers.azure.registration import ( + _build_azure_client_runtime_config, + create_azure_config, + create_azure_strategy, + register_azure_provider, + register_azure_services_with_di, +) +from orb.providers.registry import ProviderRegistry + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +class TestAzureProviderConfig: + def test_loads_provider_settings_from_azure_environment_prefix(self, monkeypatch): + monkeypatch.setenv("ORB_AZURE_REGION", "westeurope") + + assert AzureProviderConfig().region == "westeurope" + + def test_location_alias_populates_region(self): + config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + location="westeurope", + ) + + assert config.region == "westeurope" + assert config.location == "westeurope" + + def test_region_still_populates_location_accessor(self): + config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + region="westeurope", + ) + + assert config.region == "westeurope" + assert config.location == "westeurope" + + def test_conflicting_location_and_region_are_rejected(self): + with pytest.raises(ValueError, match="conflicting 'location' and 'region'"): + AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + location="westeurope", + region="eastus2", + ) + + def test_invalid_subscription_id(self): + with pytest.raises(ValueError, match="subscription_id"): + AzureProviderConfig(subscription_id="not-a-uuid") + + def test_invalid_resource_group_chars(self): + with pytest.raises(ValueError, match="resource_group"): + AzureProviderConfig(resource_group="rg with spaces!") + + def test_create_azure_config_rejects_invalid_subscription(self): + with pytest.raises(RuntimeError, match="subscription_id"): + create_azure_config( + { + "provider_type": "azure", + "subscription_id": "not-a-uuid", + } + ) + + def test_unknown_top_level_provider_fields_are_rejected(self): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + unexpected_option="value", + ) + + def test_cyclecloud_config_rejects_inline_basic_auth(self): + with pytest.raises(ValueError, match="credential_path"): + AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + cyclecloud={ + "url": "https://cc.example.com", + "username": "admin", + "password": "secret", + }, + ) + + def test_cyclecloud_config_accepts_credential_path_auth(self): + config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + cyclecloud={ + "url": "https://cc.example.com", + "credential_path": "config/cyclecloud-credentials.json", + }, + ) + + assert config.cyclecloud.credential_path == "config/cyclecloud-credentials.json" + + def test_cyclecloud_config_accepts_typed_tls_and_auth_fields(self): + config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + cyclecloud={ + "url": "https://cc.example.com", + "credential_path": "config/cyclecloud-credentials.json", + "verify_ssl": False, + "auth_mode": "bearer", + "aad_scope": "https://cc.example.com/.default", + }, + ) + + assert config.cyclecloud.verify_ssl is False + assert config.cyclecloud.auth_mode == "bearer" + assert config.cyclecloud.aad_scope == "https://cc.example.com/.default" + + def test_cyclecloud_config_rejects_unknown_fields(self): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + cyclecloud={ + "url": "https://cc.example.com", + "credential_path": "config/cyclecloud-credentials.json", + "extra_transport_option": "value", + }, + ) + + +# --------------------------------------------------------------------------- +# Template validation +# --------------------------------------------------------------------------- + + +class TestTemplateValidation: + def test_region_alias_is_accepted_for_location(self): + result = validate_azure_template( + { + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "region": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + assert result["valid"] is True + + def test_conflicting_location_and_region_is_invalid(self): + result = validate_azure_template( + { + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "region": "westeurope", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + + assert result["valid"] is False + assert any("conflicting 'location' and 'region'" in e for e in result["errors"]) + + def test_valid_template(self): + result = validate_azure_template( + { + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + assert result["valid"] is True + + def test_missing_image_source_is_invalid(self): + result = validate_azure_template( + { + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + } + ) + assert result["valid"] is False + assert any("image source is required" in e for e in result["errors"]) + + def test_missing_required_fields(self): + result = validate_azure_template({}) + assert result["valid"] is False + assert any("vm_size" in e for e in result["errors"]) + assert any("resource_group" in e for e in result["errors"]) + assert any("location" in e for e in result["errors"]) + + def test_uncommon_vm_size_warning(self): + result = validate_azure_template( + { + "template_id": "t1", + "vm_size": "custom_vm_size", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + assert result["valid"] is True + assert any("Uncommon" in w for w in result["warnings"]) + + def test_spot_regular_conflict(self): + result = validate_azure_template( + { + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + "priority": "Regular", + "eviction_policy": "Delete", + } + ) + assert result["valid"] is False + + +# --------------------------------------------------------------------------- +# Config validation +# --------------------------------------------------------------------------- + + +class TestConfigValidation: + def test_valid_config(self): + c = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg", + ) + result = validate_azure_config(c) + assert result["valid"] is True + + def test_missing_subscription(self): + c = AzureProviderConfig() + result = validate_azure_config(c) + assert result["valid"] is False + + +class TestAzureHandlerFactory: + def test_bootstrap_template_factory_preserves_azure_image_templates(self): + container = DIContainer() + container.register_instance(LoggingPort, MagicMock()) + register_infrastructure_services(container) + + factory = container.get(TemplateFactory) + template = factory.create_template( + { + "template_id": "azure-cheapest-vmss", + "provider_type": "azure", + "provider_api": "VMSS", + "vm_size": "Standard_B1s", + "resource_group": "orb-test-rg", + "location": "eastus2", + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + "admin_username": "azureuser", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "network_config": { + "subnet_id": ( + "/subscriptions/a6eb5b32-65bb-47c0-a2b8-34fa90400a4b/" + "resourceGroups/orb-test-rg/providers/Microsoft.Network/" + "virtualNetworks/orb-test-vnet/subnets/default" + ) + }, + } + ) + + assert template.image.publisher == "Canonical" + + def test_template_factory_accepts_azure_template_dto_shape(self): + container = DIContainer() + container.register_instance(LoggingPort, MagicMock()) + register_infrastructure_services(container) + + template_dto = TemplateDTO( + template_id="azure-cheapest-vmss", + provider_type="azure", + provider_api="VMSS", + provider_config=AzureTemplateExtensionConfig( + vm_size="Standard_B1s", + resource_group="orb-test-rg", + location="eastus2", + image={ + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + ), + provider_data={}, + ) + + template = container.get(TemplateFactory).create_template(template_dto.to_dict()) + + assert isinstance(template, AzureTemplate) + assert template.vm_size == "Standard_B1s" + assert template.provider_data == {} + + def test_template_factory_does_not_downgrade_invalid_azure_template(self): + container = DIContainer() + container.register_instance(LoggingPort, MagicMock()) + register_infrastructure_services(container) + + with pytest.raises(ValidationError, match="vm_size"): + container.get(TemplateFactory).create_template( + { + "template_id": "invalid-azure-template", + "provider_type": "azure", + "provider_api": "VMSS", + } + ) + + def test_registry_created_strategy_gets_azure_client_resolver(self): + registry = ProviderRegistry() + registry.clear_registrations() + perf_config = PerformanceConfig( + enable_batching=False, + enable_parallel=False, + max_workers=6, + caching={"request_status": {"enabled": False, "ttl_seconds": 17}}, + ) + config_port = MagicMock() + config_port.get_typed.side_effect = lambda config_type: ( + perf_config if config_type is PerformanceConfig else None + ) + register_azure_provider(registry=registry, config_port=config_port) + + strategy = registry.create_strategy( + "azure", + { + "provider_type": "azure", + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg-explicit", + "region": "westeurope", + }, + ) + + assert strategy.azure_client.subscription_id == "12345678-1234-1234-1234-123456789012" + assert strategy.azure_client.resource_group == "rg-explicit" + assert strategy.azure_client.perf_config["max_workers"] == 6 + assert strategy.azure_client.perf_config["cache_ttl"] == 17 + + def test_create_azure_strategy_resolves_performance_config_per_strategy_creation(self): + perf_configs = [ + PerformanceConfig( + enable_batching=False, + enable_parallel=False, + max_workers=6, + caching={"request_status": {"enabled": False, "ttl_seconds": 17}}, + ), + PerformanceConfig( + enable_batching=True, + enable_parallel=True, + max_workers=9, + caching={"request_status": {"enabled": True, "ttl_seconds": 44}}, + ), + ] + config_port = MagicMock() + config_port.get_typed.side_effect = list(perf_configs) + first_strategy = create_azure_strategy( + { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg-explicit", + "region": "westeurope", + }, + provider_instance_name="azure-test", + config_port=config_port, + ) + second_strategy = create_azure_strategy( + { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg-explicit", + "region": "westeurope", + }, + provider_instance_name="azure-test", + config_port=config_port, + ) + + assert first_strategy.azure_client.perf_config["max_workers"] == 6 + assert second_strategy.azure_client.perf_config["max_workers"] == 9 + assert config_port.get_typed.call_count == 2 + + def test_vmss_handler_receives_explicit_optional_services_from_factory(self): + azure_client = MagicMock(spec=AzureClient) + logger = MagicMock() + azure_native_spec_service = MagicMock() + azure_resource_manager = MagicMock() + + factory = AzureHandlerFactory( + azure_client=azure_client, + logger=logger, + azure_native_spec_service=azure_native_spec_service, + azure_resource_manager=azure_resource_manager, + ) + + handler = factory.create_handler(AzureProviderApi.VMSS) + + assert handler.azure_native_spec_service is azure_native_spec_service + assert handler.azure_resource_manager is azure_resource_manager + + def test_single_vm_handler_receives_explicit_native_spec_service_from_factory(self): + azure_client = MagicMock(spec=AzureClient) + logger = MagicMock() + azure_native_spec_service = MagicMock() + + factory = AzureHandlerFactory( + azure_client=azure_client, + logger=logger, + azure_native_spec_service=azure_native_spec_service, + ) + + handler = factory.create_handler(AzureProviderApi.SINGLE_VM) + + assert handler.azure_native_spec_service is azure_native_spec_service + + def test_create_azure_strategy_uses_explicit_native_spec_service(self): + azure_native_spec_service = MagicMock() + + strategy = create_azure_strategy( + { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg-explicit", + "region": "westeurope", + }, + provider_instance_name="azure-test", + azure_native_spec_service=azure_native_spec_service, + ) + + assert strategy._azure_native_spec_service is azure_native_spec_service + + +# --------------------------------------------------------------------------- +# Template extension +# --------------------------------------------------------------------------- + + +class TestTemplateExtension: + def test_defaults_do_not_include_vm_size_without_explicit_config(self): + ext = AzureTemplateExtensionConfig() + defaults = ext.to_template_defaults() + assert "vm_size" not in defaults + + def test_defaults_include_explicit_vm_size(self): + ext = AzureTemplateExtensionConfig(vm_size="Standard_B1s") + defaults = ext.to_template_defaults() + assert defaults["vm_size"] == "Standard_B1s" + + def test_defaults_include_structured_os_disk(self): + ext = AzureTemplateExtensionConfig( + os_disk={ + "storage_account_type": "StandardSSD_LRS", + "caching": "ReadOnly", + "ephemeral_os_disk": True, + } + ) + + defaults = ext.to_template_defaults() + + assert defaults["os_disk"] == { + "storage_account_type": "StandardSSD_LRS", + "caching": "ReadOnly", + "ephemeral_os_disk": True, + "ephemeral_placement": "CacheDisk", + } + + def test_defaults_include_legacy_os_disk_type_without_size(self): + ext = AzureTemplateExtensionConfig(os_disk_type="StandardSSD_LRS") + + defaults = ext.to_template_defaults() + + assert defaults["os_disk"] == {"storage_account_type": "StandardSSD_LRS"} + + def test_defaults_preserve_spot_placement_score_fields(self): + ext = AzureTemplateExtensionConfig( + spot_placement_score_enabled=True, + placement_split_strategy="hybrid", + placement_primary_share_percent=80, + placement_regions=["eastus2"], + placement_zones=["1", "2", "3"], + ) + + defaults = ext.to_template_defaults() + + assert defaults["spot_placement_score_enabled"] is True + assert defaults["placement_split_strategy"] == "hybrid" + assert defaults["placement_primary_share_percent"] == 80 + assert defaults["placement_regions"] == ["eastus2"] + assert defaults["placement_zones"] == ["1", "2", "3"] + + +class TestAzureRegistration: + def setup_method(self): + ProviderRegistry().clear_registrations() + + def test_register_azure_provider_with_di_uses_registry_instance_factory_contract(self): + provider_instance = Mock() + provider_instance.name = "azure-test" + provider_instance.config = { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg", + } + + logger = Mock() + container = Mock() + container.get.return_value = logger + + strategy = object() + + with ( + patch( + "orb.providers.azure.registration._register_azure_components_with_di", + return_value=None, + ), + patch( + "orb.providers.azure.registration._create_azure_strategy_with_di", + return_value=strategy, + ) as mock_create_strategy, + ): + from orb.domain.base.ports import LoggingPort + from orb.providers.azure.registration import register_azure_provider_with_di + + container.get.side_effect = lambda key: logger if key is LoggingPort else Mock() + + assert register_azure_provider_with_di(provider_instance, container) is True + + created = ProviderRegistry().create_strategy_by_instance( + "azure-test", + {"ignored": "config"}, + ) + + assert created is strategy + mock_create_strategy.assert_called_once() + + def test_build_azure_client_runtime_config_forwards_performance_config(self): + azure_config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="rg", + region="westeurope", + ) + perf_config = PerformanceConfig( + enable_batching=False, + enable_parallel=False, + max_workers=4, + caching={"request_status": {"enabled": False, "ttl_seconds": 21}}, + ) + fallback_config = Mock() + fallback_config.get_typed.side_effect = lambda config_type: ( + perf_config if config_type is PerformanceConfig else None + ) + + runtime_config = _build_azure_client_runtime_config( + azure_config, + logger=Mock(), + config_port=fallback_config, + ) + + assert runtime_config.azure_config is azure_config + assert runtime_config.performance_config is perf_config + fallback_config.get_typed.assert_called_once_with(PerformanceConfig) + + def test_create_azure_strategy_forwards_runtime_performance_config_to_client(self): + perf_config = PerformanceConfig( + enable_batching=False, + enable_parallel=False, + max_workers=4, + caching={"request_status": {"enabled": False, "ttl_seconds": 21}}, + ) + strategy = create_azure_strategy( + { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "rg", + "region": "westeurope", + }, + provider_instance_name="azure-test", + performance_config=perf_config, + ) + + assert strategy.azure_client.perf_config["enable_batching"] is False + assert strategy.azure_client.perf_config["enable_parallel"] is False + assert strategy.azure_client.perf_config["max_workers"] == 4 + assert strategy.azure_client.perf_config["enable_caching"] is False + assert strategy.azure_client.perf_config["cache_ttl"] == 21 + + def test_register_services_populates_template_example_generator_registry(self): + from orb.infrastructure.registry.template_example_generator_registry import ( + TemplateExampleGeneratorRegistry, + ) + + TemplateExampleGeneratorRegistry.clear() + logger = Mock() + container = Mock() + container.get.return_value = logger + container.is_registered.return_value = True + + register_azure_services_with_di(container) + + generator = TemplateExampleGeneratorRegistry.get("azure") + assert generator is not None + assert generator.generate_example_templates("azure-default") + + TemplateExampleGeneratorRegistry.clear() + + +def test_azure_template_example_generator_filters_by_provider_api() -> None: + from orb.providers.azure.infrastructure.adapters.template_example_generator_adapter import ( + AzureTemplateExampleGeneratorAdapter, + ) + + examples = AzureTemplateExampleGeneratorAdapter().generate_example_templates( + "azure-default", provider_api="SingleVM" + ) + + assert len(examples) == 1 + assert examples[0]["provider_api"] == "SingleVM" diff --git a/tests/providers/azure/test_azure_credential_factory.py b/tests/providers/azure/test_azure_credential_factory.py new file mode 100644 index 000000000..50259a572 --- /dev/null +++ b/tests/providers/azure/test_azure_credential_factory.py @@ -0,0 +1,214 @@ +"""Focused tests for Azure credential factory helpers.""" + +import sys +import types +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from orb.providers.azure.infrastructure.credential_factory import ( + AsyncAzureCredentialAccessTokenProvider, + AsyncDefaultAzureAccessTokenProvider, + AzureCredentialAccessTokenProvider, + DefaultAzureAccessTokenProvider, + create_default_azure_credential, + create_default_azure_credential_async, + get_default_azure_credential_error_types, +) + + +def test_azure_credential_access_token_provider_adapts_existing_credential(): + credential = MagicMock() + credential.get_token.return_value = types.SimpleNamespace(token="access-token") + + provider = AzureCredentialAccessTokenProvider(credential) + + assert provider.get_access_token("scope") == "access-token" + credential.get_token.assert_called_once_with("scope") + + +def test_create_default_azure_credential_passes_managed_identity_client_id_when_configured(): + fake_identity = types.ModuleType("azure.identity") + fake_ctor = MagicMock(return_value=MagicMock()) + fake_identity.DefaultAzureCredential = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": fake_identity, + }, + ): + credential = create_default_azure_credential( + client_id="managed-identity-client-id", + logger=MagicMock(), + ) + + assert credential is fake_ctor.return_value + fake_ctor.assert_called_once_with(managed_identity_client_id="managed-identity-client-id") + + +def test_create_default_azure_credential_omits_managed_identity_client_id_when_unset(): + fake_identity = types.ModuleType("azure.identity") + fake_ctor = MagicMock(return_value=MagicMock()) + fake_identity.DefaultAzureCredential = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": fake_identity, + }, + ): + create_default_azure_credential( + client_id=None, + logger=MagicMock(), + ) + + fake_ctor.assert_called_once_with() + + +def test_create_default_azure_credential_async_preserves_nested_import_error(): + fake_identity_aio = types.ModuleType("azure.identity.aio") + fake_ctor = MagicMock(side_effect=ImportError("aiohttp package is not installed")) + fake_identity_aio.DefaultAzureCredential = fake_ctor + logger = MagicMock() + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": types.ModuleType("azure.identity"), + "azure.identity.aio": fake_identity_aio, + }, + ): + with pytest.raises(ImportError, match="aiohttp package is not installed"): + create_default_azure_credential_async( + client_id=None, + logger=logger, + ) + + logger.error.assert_called_once_with( + "azure-identity dependency error: aiohttp package is not installed" + ) + + +@pytest.mark.asyncio +async def test_async_azure_credential_access_token_provider_adapts_existing_credential(): + credential = MagicMock() + credential.get_token = AsyncMock(return_value=types.SimpleNamespace(token="access-token")) + + provider = AsyncAzureCredentialAccessTokenProvider(credential) + + assert await provider.get_access_token("scope") == "access-token" + credential.get_token.assert_awaited_once_with("scope") + + +def test_create_default_azure_credential_async_passes_managed_identity_client_id_when_configured(): + fake_identity_aio = types.ModuleType("azure.identity.aio") + fake_ctor = MagicMock(return_value=MagicMock()) + fake_identity_aio.DefaultAzureCredential = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": types.ModuleType("azure.identity"), + "azure.identity.aio": fake_identity_aio, + }, + ): + credential = create_default_azure_credential_async( + client_id="managed-identity-client-id", + logger=MagicMock(), + ) + + assert credential is fake_ctor.return_value + fake_ctor.assert_called_once_with(managed_identity_client_id="managed-identity-client-id") + + +def test_default_azure_access_token_provider_closes_short_lived_credential(): + fake_identity = types.ModuleType("azure.identity") + fake_credential = MagicMock() + fake_credential.get_token.return_value = types.SimpleNamespace(token="access-token") + fake_ctor = MagicMock(return_value=fake_credential) + fake_identity.DefaultAzureCredential = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": fake_identity, + }, + ): + provider = DefaultAzureAccessTokenProvider( + logger=MagicMock(), + client_id=None, + ) + token = provider.get_access_token("https://management.azure.com/.default") + + assert token == "access-token" + fake_credential.close.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_async_default_azure_access_token_provider_closes_short_lived_credential(): + fake_identity_aio = types.ModuleType("azure.identity.aio") + fake_credential = MagicMock() + fake_credential.get_token = AsyncMock(return_value=types.SimpleNamespace(token="access-token")) + fake_credential.close = AsyncMock() + fake_ctor = MagicMock(return_value=fake_credential) + fake_identity_aio.DefaultAzureCredential = fake_ctor + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.identity": types.ModuleType("azure.identity"), + "azure.identity.aio": fake_identity_aio, + }, + ): + provider = AsyncDefaultAzureAccessTokenProvider( + logger=MagicMock(), + client_id=None, + ) + token = await provider.get_access_token("https://management.azure.com/.default") + + assert token == "access-token" + fake_credential.close.assert_awaited_once_with() + + +def test_get_default_azure_credential_error_types_returns_sdk_types_when_available(): + fake_core = types.ModuleType("azure.core.exceptions") + fake_identity = types.ModuleType("azure.identity") + + class FakeClientAuthenticationError(Exception): + pass + + class FakeCredentialUnavailableError(Exception): + pass + + fake_core.ClientAuthenticationError = FakeClientAuthenticationError + fake_identity.CredentialUnavailableError = FakeCredentialUnavailableError + + with patch.dict( + sys.modules, + { + "azure": types.ModuleType("azure"), + "azure.core": types.ModuleType("azure.core"), + "azure.core.exceptions": fake_core, + "azure.identity": fake_identity, + }, + ): + assert get_default_azure_credential_error_types() == ( + FakeCredentialUnavailableError, + FakeClientAuthenticationError, + ) + + +def test_default_azure_access_token_provider_reports_import_error_as_expected_auth_type(): + provider = DefaultAzureAccessTokenProvider( + logger=MagicMock(), + client_id=None, + ) + + assert ImportError in provider.get_auth_error_types() diff --git a/tests/providers/azure/test_azure_deployment_service.py b/tests/providers/azure/test_azure_deployment_service.py new file mode 100644 index 000000000..33ac5637a --- /dev/null +++ b/tests/providers/azure/test_azure_deployment_service.py @@ -0,0 +1,194 @@ +"""Focused tests for Azure deployment helpers.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.infrastructure.services.azure_deployment_service import ( + AzureDeploymentService, +) + + +def test_build_deployment_name_normalizes_and_truncates_deterministically(): + name = AzureDeploymentService.build_deployment_name( + "vm", + "req with spaces", + "template/with/slashes", + "x" * 80, + ) + + assert len(name) <= 64 + assert " " not in name + assert "/" not in name + assert name.startswith("vm-req-with-spaces-template-with-slashes-") + + +def test_resource_id_expression_targets_sibling_resource(): + assert ( + AzureDeploymentService.resource_id_expression( + "Microsoft.Network/networkInterfaces", + "nic-vm-1", + ) + == "[resourceId('Microsoft.Network/networkInterfaces', 'nic-vm-1')]" + ) + + +@pytest.mark.asyncio +async def test_submit_template_deployment_async_uses_resources_api_without_waiting(): + azure_client = MagicMock() + async_resource_client = MagicMock() + async_resource_client.resources.begin_create_or_update = AsyncMock() + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource_client) + logger = MagicMock() + service = AzureDeploymentService(azure_client=azure_client, logger=logger) + template = {"resources": []} + + deployment_name = await service.submit_template_deployment_async( + resource_group="test-rg", + deployment_name="dep-1", + template=template, + parameters={"vmName": {"value": "vm-1"}}, + ) + + assert deployment_name == "dep-1" + async_resource_client.resources.begin_create_or_update.assert_awaited_once_with( + resource_group_name="test-rg", + resource_provider_namespace="Microsoft.Resources", + parent_resource_path="", + resource_type="deployments", + resource_name="dep-1", + api_version="2025-04-01", + parameters={ + "properties": { + "mode": "Incremental", + "template": template, + "parameters": {"vmName": {"value": "vm-1"}}, + } + }, + ) + + +@pytest.mark.asyncio +async def test_get_deployment_status_async_extracts_provisioning_and_error_fields(): + azure_client = MagicMock() + async_resource_client = MagicMock() + async_resource_client.resources.get = AsyncMock() + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource_client) + logger = MagicMock() + service = AzureDeploymentService(azure_client=azure_client, logger=logger) + deployment = MagicMock() + deployment.properties = { + "provisioningState": "Failed", + "error": { + "code": "DeploymentFailed", + "message": "template validation failed", + }, + } + async_resource_client.resources.get.return_value = deployment + + status = await service.get_deployment_status_async( + resource_group="test-rg", + deployment_name="dep-1", + ) + + assert status == { + "provisioning_state": "Failed", + "error_code": "DeploymentFailed", + "error_message": "template validation failed", + } + + +@pytest.mark.asyncio +async def test_get_deployment_status_async_tolerates_non_mapping_properties(): + azure_client = MagicMock() + async_resource_client = MagicMock() + async_resource_client.resources.get = AsyncMock() + azure_client.get_async_resource_client = AsyncMock(return_value=async_resource_client) + logger = MagicMock() + service = AzureDeploymentService(azure_client=azure_client, logger=logger) + deployment = MagicMock() + deployment.properties = None + async_resource_client.resources.get.return_value = deployment + + status = await service.get_deployment_status_async( + resource_group="test-rg", + deployment_name="dep-1", + ) + + assert status == { + "provisioning_state": None, + "error_code": None, + "error_message": None, + } + + +def test_build_single_vm_deployment_template_attaches_public_ip_when_enabled(): + service = AzureDeploymentService(azure_client=MagicMock(), logger=MagicMock()) + deployment_template = service.build_single_vm_deployment_template( + location="eastus2", + subnet_id="/subscriptions/.../subnets/default", + vm_definitions=[ + { + "vm_name": "vm-test", + "nic_name": "nic-vm-test", + "public_ip_name": "pip-vm-test", + "vm_payload": { + "location": "eastus2", + "properties": { + "networkProfile": { + "networkInterfaces": [ + { + "id": service.resource_id_expression( + "Microsoft.Network/networkInterfaces", + "nic-vm-test", + ), + "properties": { + "primary": True, + "deleteOption": "Delete", + }, + } + ] + }, + "storageProfile": { + "osDisk": {"deleteOption": "Delete"}, + "dataDisks": [{"deleteOption": "Delete"}], + }, + }, + }, + } + ], + ) + + public_ip_resource = next( + resource + for resource in deployment_template["resources"] + if resource["type"] == "Microsoft.Network/publicIPAddresses" + ) + nic_resource = next( + resource + for resource in deployment_template["resources"] + if resource["type"] == "Microsoft.Network/networkInterfaces" + ) + vm_resource = next( + resource + for resource in deployment_template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines" + ) + + assert public_ip_resource == { + "type": "Microsoft.Network/publicIPAddresses", + "apiVersion": "2023-09-01", + "name": "pip-vm-test", + "location": "eastus2", + "sku": {"name": "Standard"}, + "properties": { + "publicIPAllocationMethod": "Static", + "deleteOption": "Delete", + }, + } + assert nic_resource["properties"]["ipConfigurations"][0]["properties"]["publicIPAddress"] == { + "id": "[resourceId('Microsoft.Network/publicIPAddresses', 'pip-vm-test')]", + "deleteOption": "Delete", + } + assert vm_resource["properties"]["storageProfile"]["osDisk"]["deleteOption"] == "Delete" + assert vm_resource["properties"]["storageProfile"]["dataDisks"][0]["deleteOption"] == "Delete" diff --git a/tests/providers/azure/test_azure_error_utils.py b/tests/providers/azure/test_azure_error_utils.py new file mode 100644 index 000000000..fbd2016b4 --- /dev/null +++ b/tests/providers/azure/test_azure_error_utils.py @@ -0,0 +1,117 @@ +"""Focused tests for Azure error normalization helpers.""" + +from types import SimpleNamespace + +from orb.providers.azure.infrastructure.error_utils import ( + canonical_azure_error_code, + extract_azure_error_details, +) + + +def test_extract_azure_error_details_reads_common_exception_shapes(): + exc = SimpleNamespace( + error_code="AllocationFailed", + status_code=409, + message="allocation failed in zone 1", + error=None, + response=None, + ) + + assert extract_azure_error_details(exc) == { + "raw_error_code": "AllocationFailed", + "status_code": 409, + "message": "allocation failed in zone 1", + "details": [], + } + + +def test_extract_azure_error_details_falls_back_to_nested_error_and_response(): + exc = SimpleNamespace( + error=SimpleNamespace(code="QuotaExceeded", message="quota exceeded"), + response=SimpleNamespace(status_code=429), + ) + + assert extract_azure_error_details(exc) == { + "raw_error_code": "QuotaExceeded", + "status_code": 429, + "message": "quota exceeded", + "details": [], + } + + +def test_extract_azure_error_details_preserves_nested_arm_details(): + exc = SimpleNamespace( + error=SimpleNamespace( + code="InvalidTemplateDeployment", + message="Deployment validation failed", + details=[ + { + "code": "InvalidParameter", + "message": "The supplied VM size is not available in this location.", + } + ], + ), + response=SimpleNamespace(status_code=400), + ) + + assert extract_azure_error_details(exc) == { + "raw_error_code": "InvalidTemplateDeployment", + "status_code": 400, + "message": "Deployment validation failed", + "details": [ + { + "code": "InvalidParameter", + "message": "The supplied VM size is not available in this location.", + } + ], + } + + +def test_extract_azure_error_details_reads_wrapped_launch_error_details(): + exc = SimpleNamespace( + error_code="InvalidRequest", + message="Deployment validation failed", + details={ + "raw_error_code": "InvalidTemplateDeployment", + "status_code": 400, + "details": [ + { + "code": "InvalidSubnet", + "message": "The subnet resource ID is invalid.", + } + ], + }, + ) + + assert extract_azure_error_details(exc) == { + "raw_error_code": "InvalidRequest", + "status_code": 400, + "message": "Deployment validation failed", + "details": [ + { + "code": "InvalidSubnet", + "message": "The subnet resource ID is invalid.", + } + ], + } + + +def test_canonical_azure_error_code_prefers_raw_error_code(): + exc = SimpleNamespace(error_code="SkuNotAvailable", status_code=409, message="ignored") + + assert canonical_azure_error_code(exc) == "SkuNotAvailable" + + +def test_canonical_azure_error_code_maps_status_codes_and_messages(): + assert ( + canonical_azure_error_code(SimpleNamespace(status_code=429, error=None, response=None)) + == "TooManyRequests" + ) + assert canonical_azure_error_code(Exception("quota exceeded for region")) == "QuotaExceeded" + assert ( + canonical_azure_error_code(Exception("insufficient capacity to allocate")) + == "AllocationFailed" + ) + assert ( + canonical_azure_error_code(Exception("validation failed for parameter")) == "InvalidRequest" + ) diff --git a/tests/providers/azure/test_azure_exceptions.py b/tests/providers/azure/test_azure_exceptions.py new file mode 100644 index 000000000..fbf6c784a --- /dev/null +++ b/tests/providers/azure/test_azure_exceptions.py @@ -0,0 +1,63 @@ +"""Tests for Azure's provider-neutral exception contract.""" + +import pytest + +from orb.providers.azure.exceptions import ( + AuthenticationError, + AzureConfigurationError, + AzureError, + AzureInfrastructureError, + AzureValidationError, + NetworkError, + QuotaExceededError, + RateLimitError, +) +from orb.providers.base.exceptions import ( + ProviderAuthError, + ProviderConfigError, + ProviderError, + ProviderPermanentError, + ProviderQuotaError, + ProviderTransientError, +) + + +@pytest.mark.parametrize( + ("azure_error_type", "provider_error_type", "is_retryable"), + [ + (AzureValidationError, ProviderPermanentError, False), + (AuthenticationError, ProviderAuthError, False), + (AzureConfigurationError, ProviderConfigError, False), + (QuotaExceededError, ProviderQuotaError, True), + (RateLimitError, ProviderQuotaError, True), + (NetworkError, ProviderTransientError, True), + (AzureInfrastructureError, ProviderTransientError, True), + ], +) +def test_azure_errors_declare_provider_failure_semantics( + azure_error_type: type[AzureError], + provider_error_type: type[ProviderError], + is_retryable: bool, +) -> None: + error = azure_error_type("operation failed") + + assert isinstance(error, provider_error_type) + assert error.provider_type == "azure" + assert error.is_retryable is is_retryable + + +def test_safe_serialization_omits_underlying_exception_and_preserves_error_code() -> None: + error = NetworkError( + "Azure API unavailable", + error_code="ServiceUnavailable", + underlying_exception=RuntimeError("secret-bearing SDK response"), + ) + + assert error.to_dict()["underlying_exception"] == "RuntimeError('secret-bearing SDK response')" + assert error.safe_to_dict() == { + "error_type": "NetworkError", + "message": "Azure API unavailable", + "provider_type": "azure", + "is_retryable": True, + "error_code": "ServiceUnavailable", + } diff --git a/tests/providers/azure/test_azure_hostfactory_field_mapping.py b/tests/providers/azure/test_azure_hostfactory_field_mapping.py new file mode 100644 index 000000000..58014bc97 --- /dev/null +++ b/tests/providers/azure/test_azure_hostfactory_field_mapping.py @@ -0,0 +1,44 @@ +"""Azure HostFactory field-mapping contracts.""" + +from orb.infrastructure.scheduler.hostfactory.field_mapper import HostFactoryFieldMapper +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.registration import register_azure_hostfactory_field_mapping + + +def test_azure_hostfactory_field_mapping_is_provider_owned(): + register_azure_hostfactory_field_mapping() + + mapped = HostFactoryFieldMapper("azure").map_input_fields( + { + "templateId": "azure-hf", + "resourceGroup": "rg-test", + "location": "uksouth", + "vmType": "Standard_D4s_v5", + "vmTypes": {"Standard_D4s_v5": 1, "Standard_D8s_v5": 1}, + "imageId": "Canonical:0001-com-ubuntu-server-jammy:22_04-lts:latest", + "maxNumber": 2, + } + ) + + assert mapped["template_id"] == "azure-hf" + assert mapped["resource_group"] == "rg-test" + assert mapped["vm_size"] == "Standard_D4s_v5" + assert mapped["vm_sizes"] == { + "Standard_D4s_v5": 1, + "Standard_D8s_v5": 1, + } + + +def test_azure_template_accepts_hostfactory_vm_types_mapping(): + template = AzureTemplate( + template_id="azure-hf", + resource_group="rg-test", + location="uksouth", + vm_size="Standard_D4s_v5", + vm_sizes={"Standard_D8s_v5": 1, "Standard_D16s_v5": 1}, + image_id="Canonical:0001-com-ubuntu-server-jammy:22_04-lts:latest", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABg azure@test"], + max_instances=2, + ) + + assert template.vm_sizes == ["Standard_D8s_v5", "Standard_D16s_v5"] diff --git a/tests/providers/azure/test_azure_native_spec_service.py b/tests/providers/azure/test_azure_native_spec_service.py new file mode 100644 index 000000000..3f5b3cc1e --- /dev/null +++ b/tests/providers/azure/test_azure_native_spec_service.py @@ -0,0 +1,146 @@ +"""Tests for Azure native spec processing.""" + +from unittest.mock import Mock, patch + +from pydantic import ValidationError + +from orb.application.services.native_spec_service import NativeSpecService +from orb.config.schemas.provider_strategy_schema import ProviderConfig +from orb.domain.request.aggregate import Request +from orb.domain.request.request_types import RequestType +from orb.providers.azure.infrastructure.services.azure_native_spec_service import ( + AzureNativeSpecService, +) +from tests.providers.azure.strategy_test_support import make_azure_template + + +def _make_template(**overrides): + return make_azure_template( + template_id="azure-native-spec-test", + provider_api="VMSS", + **overrides, + ) + + +def _make_request(): + return Request.create_new_request( + request_type=RequestType.ACQUIRE, + template_id="azure-native-spec-test", + machine_count=2, + provider_type="azure", + provider_name="azure-default", + ) + + +def _make_provider_config(*, azure_extensions=None): + return ProviderConfig( + providers=[ + { + "name": "azure-default", + "type": "azure", + "enabled": True, + "config": {}, + } + ], + provider_defaults={ + "azure": { + "extensions": azure_extensions, + } + } + if azure_extensions is not None + else {}, + ) + + +def test_process_provider_api_spec_with_merge_merges_rendered_spec(): + config_port = Mock() + config_port.get_native_spec_config.return_value = {"enabled": True, "merge_mode": "merge"} + config_port.get_package_info.return_value = {"name": "orb", "version": "1.0.0"} + config_port.get_provider_config.return_value = _make_provider_config() + + logger = Mock() + spec_renderer = Mock() + spec_renderer.render_spec.return_value = { + "tags": {"Rendered": "{{ request_id }}"}, + "sku": {"name": "Standard_D8s_v5"}, + } + + native_spec_service = NativeSpecService(config_port, spec_renderer, logger) + service = AzureNativeSpecService(native_spec_service, config_port) + + template = _make_template(provider_api_spec={"tags": {"Rendered": "{{ request_id }}"}}) + request = _make_request() + default_payload = {"location": "eastus2", "sku": {"capacity": 2}} + + result = service.process_provider_api_spec_with_merge(template, request, default_payload) + + assert result["location"] == "eastus2" + assert result["sku"]["capacity"] == 2 + assert result["sku"]["name"] == "Standard_D8s_v5" + spec_renderer.render_spec.assert_called_once() + + +def test_process_provider_api_spec_with_merge_replace_mode_replaces_default(): + config_port = Mock() + config_port.get_native_spec_config.return_value = {"enabled": True, "merge_mode": "replace"} + config_port.get_package_info.return_value = {"name": "orb", "version": "1.0.0"} + config_port.get_provider_config.return_value = _make_provider_config() + + logger = Mock() + spec_renderer = Mock() + spec_renderer.render_spec.return_value = {"location": "westus2"} + + native_spec_service = NativeSpecService(config_port, spec_renderer, logger) + service = AzureNativeSpecService(native_spec_service, config_port) + + template = _make_template(provider_api_spec={"location": "{{ location }}"}) + request = _make_request() + + result = service.process_provider_api_spec_with_merge( + template, + request, + {"location": "eastus2", "sku": {"capacity": 2}}, + ) + + assert result == {"location": "westus2"} + + +def test_load_spec_file_uses_typed_provider_config_extensions_path(): + config_port = Mock() + config_port.get_provider_config.return_value = _make_provider_config( + azure_extensions={"native_spec": {"spec_file_base_path": "config/specs/azure"}} + ) + + service = AzureNativeSpecService( + NativeSpecService(config_port, Mock(), Mock()), + config_port, + ) + + with patch( + "orb.providers.azure.infrastructure.services.azure_native_spec_service.read_json_file" + ) as mock_read: + mock_read.return_value = {"location": "eastus2"} + + result = service._load_spec_file("vmss.json") + + assert result == {"location": "eastus2"} + mock_read.assert_called_once_with("config/specs/azure/vmss.json") + + +def test_load_spec_file_rejects_non_object_native_spec_extensions(): + config_port = Mock() + config_port.get_provider_config.return_value = _make_provider_config( + azure_extensions={"native_spec": "config/specs/azure"} + ) + + service = AzureNativeSpecService( + NativeSpecService(config_port, Mock(), Mock()), + config_port, + ) + + try: + service._load_spec_file("vmss.json") + except ValidationError as exc: + assert "native_spec" in str(exc) + else: + raise AssertionError("Expected ValidationError for malformed Azure native_spec config") diff --git a/tests/providers/azure/test_azure_network_identity_resolver.py b/tests/providers/azure/test_azure_network_identity_resolver.py new file mode 100644 index 000000000..61f237d5b --- /dev/null +++ b/tests/providers/azure/test_azure_network_identity_resolver.py @@ -0,0 +1,319 @@ +"""Focused tests for ARM ID parsing and Azure network identity resolution.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.infrastructure.services.arm_resource_id_parser import ( + ArmResourceIdParser, +) +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentityResolver, +) + + +class TestArmResourceIdParser: + def test_extract_resource_group_and_name_rejects_incomplete_ids(self): + parser = ArmResourceIdParser() + + assert parser.extract_resource_group_and_name("/subscriptions/sub/resourceGroups") is None + + def test_extract_resource_group_and_name_rejects_malformed_ids(self): + parser = ArmResourceIdParser() + + assert ( + parser.extract_resource_group_and_name( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces//nic-vm-1" + ) + is None + ) + assert ( + parser.extract_resource_group_and_name( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces" + ) + is None + ) + + def test_extract_resource_group_and_name_accepts_child_resources(self): + parser = ArmResourceIdParser() + + assert parser.extract_resource_group_and_name( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) == ("test-rg", "default") + + def test_subnet_id_to_vnet_id_rejects_malformed_or_non_subnet_ids(self): + parser = ArmResourceIdParser() + + assert ( + parser.subnet_id_to_vnet_id( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets" + ) + is None + ) + assert ( + parser.subnet_id_to_vnet_id( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + is None + ) + + +class TestAzureNetworkIdentityResolver: + @staticmethod + def _build_resolver() -> tuple[AzureNetworkIdentityResolver, MagicMock]: + network_client = MagicMock() + network_client.network_interfaces.get = AsyncMock() + network_client.public_ip_addresses.get = AsyncMock() + resolver = AzureNetworkIdentityResolver( + async_network_client_getter=AsyncMock(return_value=network_client), + logger=MagicMock(), + arm_resource_id_parser=ArmResourceIdParser(), + network_lookup_error_types=lambda: (), + ) + return resolver, network_client + + @pytest.mark.asyncio + async def test_resolve_from_vm_populates_ips_and_subnet(self): + resolver, network_client = self._build_resolver() + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + vm = MagicMock() + vm.network_profile.network_interfaces = [nic_ref] + + subnet = MagicMock() + subnet.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) + public_ip_ref = MagicMock() + public_ip_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/publicIPAddresses/pip-vm-1" + ) + ip_cfg = MagicMock() + ip_cfg.private_ip_address = "10.0.0.4" + ip_cfg.subnet = subnet + ip_cfg.public_ip_address = public_ip_ref + + nic = MagicMock() + nic.ip_configurations = [ip_cfg] + network_client.network_interfaces.get.return_value = nic + + pip = MagicMock() + pip.ip_address = "52.1.2.3" + network_client.public_ip_addresses.get.return_value = pip + + result = await resolver.resolve_from_vm_async(vm) + + assert result["private_ip"] == "10.0.0.4" + assert result["public_ip"] == "52.1.2.3" + assert result["subnet_id"].endswith("/subnets/default") + assert result["vnet_id"].endswith("/virtualNetworks/test-vnet") + assert result["nic_name"] == "nic-vm-1" + + @pytest.mark.asyncio + async def test_resolve_from_vm_tolerates_missing_nested_property_bags(self): + resolver, network_client = self._build_resolver() + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties = None + + vm = MagicMock() + vm.network_profile.network_interfaces = [nic_ref] + + subnet = MagicMock() + subnet.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) + public_ip_ref = MagicMock() + public_ip_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/publicIPAddresses/pip-vm-1" + ) + ip_cfg = MagicMock() + ip_cfg.private_ip_address = "10.0.0.4" + ip_cfg.subnet = subnet + ip_cfg.public_ip_address = public_ip_ref + ip_cfg.properties = None + + nic = MagicMock() + nic.ip_configurations = [ip_cfg] + network_client.network_interfaces.get.return_value = nic + + pip = MagicMock() + pip.ip_address = "52.1.2.3" + network_client.public_ip_addresses.get.return_value = pip + + result = await resolver.resolve_from_vm_async(vm) + + assert result["private_ip"] == "10.0.0.4" + assert result["public_ip"] == "52.1.2.3" + assert result["subnet_id"].endswith("/subnets/default") + assert result["vnet_id"].endswith("/virtualNetworks/test-vnet") + assert result["nic_name"] == "nic-vm-1" + + @pytest.mark.asyncio + async def test_resolve_from_nic_refs_skips_known_nic_lookup_errors(self): + from azure.core.exceptions import ResourceNotFoundError + + network_client = MagicMock() + network_client.network_interfaces.get = AsyncMock( + side_effect=ResourceNotFoundError("missing") + ) + resolver = AzureNetworkIdentityResolver( + async_network_client_getter=AsyncMock(return_value=network_client), + logger=MagicMock(), + arm_resource_id_parser=ArmResourceIdParser(), + network_lookup_error_types=lambda: (ResourceNotFoundError,), + ) + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + assert await resolver.resolve_from_nic_refs_async([nic_ref]) == { + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + + @pytest.mark.asyncio + async def test_resolve_from_nic_refs_reraises_unexpected_nic_lookup_errors(self): + resolver, network_client = self._build_resolver() + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + network_client.network_interfaces.get.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + await resolver.resolve_from_nic_refs_async([nic_ref]) + + @pytest.mark.asyncio + async def test_resolve_from_vm_async_populates_ips_and_subnet(self): + network_client = MagicMock() + network_client.network_interfaces.get = AsyncMock() + network_client.public_ip_addresses.get = AsyncMock() + resolver = AzureNetworkIdentityResolver( + async_network_client_getter=AsyncMock(return_value=network_client), + logger=MagicMock(), + arm_resource_id_parser=ArmResourceIdParser(), + network_lookup_error_types=lambda: (), + ) + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + vm = MagicMock() + vm.network_profile.network_interfaces = [nic_ref] + + subnet = MagicMock() + subnet.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) + public_ip_ref = MagicMock() + public_ip_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/publicIPAddresses/pip-vm-1" + ) + ip_cfg = MagicMock() + ip_cfg.private_ip_address = "10.0.0.4" + ip_cfg.subnet = subnet + ip_cfg.public_ip_address = public_ip_ref + + nic = MagicMock() + nic.ip_configurations = [ip_cfg] + network_client.network_interfaces.get.return_value = nic + + pip = MagicMock() + pip.ip_address = "52.1.2.3" + network_client.public_ip_addresses.get.return_value = pip + + result = await resolver.resolve_from_vm_async(vm) + + assert result["private_ip"] == "10.0.0.4" + assert result["public_ip"] == "52.1.2.3" + assert result["subnet_id"].endswith("/subnets/default") + assert result["vnet_id"].endswith("/virtualNetworks/test-vnet") + assert result["nic_name"] == "nic-vm-1" + + @pytest.mark.asyncio + async def test_resolve_from_nic_refs_async_skips_known_public_ip_lookup_errors(self): + from azure.core.exceptions import ResourceNotFoundError + + network_client = MagicMock() + network_client.network_interfaces.get = AsyncMock() + network_client.public_ip_addresses.get = AsyncMock( + side_effect=ResourceNotFoundError("missing") + ) + resolver = AzureNetworkIdentityResolver( + async_network_client_getter=AsyncMock(return_value=network_client), + logger=MagicMock(), + arm_resource_id_parser=ArmResourceIdParser(), + network_lookup_error_types=lambda: (ResourceNotFoundError,), + ) + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + subnet = MagicMock() + subnet.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ) + public_ip_ref = MagicMock() + public_ip_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/publicIPAddresses/pip-vm-1" + ) + ip_cfg = MagicMock() + ip_cfg.private_ip_address = "10.0.0.4" + ip_cfg.subnet = subnet + ip_cfg.public_ip_address = public_ip_ref + + nic = MagicMock() + nic.ip_configurations = [ip_cfg] + network_client.network_interfaces.get.return_value = nic + + result = await resolver.resolve_from_nic_refs_async([nic_ref]) + + assert result["private_ip"] == "10.0.0.4" + assert result["public_ip"] is None + assert result["subnet_id"].endswith("/subnets/default") + assert result["nic_name"] == "nic-vm-1" diff --git a/tests/providers/azure/test_azure_resource_manager.py b/tests/providers/azure/test_azure_resource_manager.py new file mode 100644 index 000000000..cb37403fe --- /dev/null +++ b/tests/providers/azure/test_azure_resource_manager.py @@ -0,0 +1,130 @@ +"""Focused tests for AzureResourceManager.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.exceptions.azure_exceptions import AzureInfrastructureError +from orb.providers.azure.managers.azure_resource_manager import AzureResourceManager + + +def _make_manager(): + azure_client = MagicMock() + logger = MagicMock() + config = AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group="test-rg", + region="eastus2", + ) + return ( + AzureResourceManager(azure_client=azure_client, config=config, logger=logger), + azure_client, + logger, + ) + + +@pytest.mark.asyncio +async def test_get_vmss_capacity_async_uses_async_vmss_capacity_and_member_count(): + manager, azure_client, _ = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + + vmss = MagicMock() + vmss.sku.capacity = 5 + vmss.sku.name = "Standard_D4s_v5" + vmss.orchestration_mode = "Flexible" + vmss.provisioning_state = "Updating" + async_compute_client.virtual_machine_scale_sets.get = AsyncMock(return_value=vmss) + + async def _vm_iter(): + yield MagicMock( + virtual_machine_scale_set=MagicMock(id="/.../virtualMachineScaleSets/vmss-1") + ) + yield MagicMock( + virtual_machine_scale_set=MagicMock(id="/.../virtualMachineScaleSets/vmss-1") + ) + + async_compute_client.virtual_machines.list.return_value = _vm_iter() + + result = await manager.get_vmss_capacity_async("test-rg", "vmss-1") + + assert result == { + "vmss_name": "vmss-1", + "resource_group": "test-rg", + "capacity": 5, + "vm_size": "Standard_D4s_v5", + "provisioning_state": "Updating", + "provisioned_instance_count": 2, + } + + +@pytest.mark.asyncio +async def test_get_vmss_capacity_async_raises_infrastructure_error_when_lookup_fails(): + manager, azure_client, _ = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + async_compute_client.virtual_machine_scale_sets.get = AsyncMock( + side_effect=RuntimeError("boom") + ) + + with pytest.raises(AzureInfrastructureError, match="Failed to get VMSS capacity"): + await manager.get_vmss_capacity_async("test-rg", "vmss-1") + + +@pytest.mark.asyncio +async def test_get_vmss_member_count_async_counts_uniform_instances(): + manager, azure_client, _ = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + + async def _vm_iter(): + yield MagicMock() + yield MagicMock() + yield MagicMock() + + async_compute_client.virtual_machine_scale_set_vms.list.return_value = _vm_iter() + + assert ( + await manager.get_vmss_member_count_async("test-rg", "vmss-1", orchestration_mode="Uniform") + == 3 + ) + + +@pytest.mark.asyncio +async def test_get_vmss_member_count_async_tolerates_count_errors_and_returns_none(): + manager, azure_client, logger = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + async_compute_client.virtual_machine_scale_set_vms.list.side_effect = RuntimeError("boom") + + assert ( + await manager.get_vmss_member_count_async("test-rg", "vmss-1", orchestration_mode="Uniform") + is None + ) + logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_vmss_exists_async_recognizes_not_found_errors(): + manager, azure_client, _ = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + exc = RuntimeError("not found") + exc.status_code = 404 + async_compute_client.virtual_machine_scale_sets.get = AsyncMock(side_effect=exc) + + assert await manager.vmss_exists_async("test-rg", "vmss-1") is False + + +@pytest.mark.asyncio +async def test_vmss_exists_async_returns_none_for_unknown_lookup_error(): + manager, azure_client, logger = _make_manager() + async_compute_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute_client) + async_compute_client.virtual_machine_scale_sets.get = AsyncMock( + side_effect=RuntimeError("boom") + ) + + assert await manager.vmss_exists_async("test-rg", "vmss-1") is None + logger.warning.assert_called_once() diff --git a/tests/providers/azure/test_azure_strategy.py b/tests/providers/azure/test_azure_strategy.py new file mode 100644 index 000000000..0dcfa2e0b --- /dev/null +++ b/tests/providers/azure/test_azure_strategy.py @@ -0,0 +1,979 @@ +"""Focused tests for Azure strategy core behavior.""" + +import asyncio +import threading +import time +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.infrastructure.services.spot_placement_score_adapter import ( + AzureSpotPlacementScoreAdapter, +) +from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy +from orb.providers.base.strategy import ( + ProviderOperation, + ProviderOperationType, + ProviderResult, +) +from tests.providers.azure.strategy_test_support import run_operation + + +class TestInitialization: + def test_requires_azure_config(self, logger): + with pytest.raises(ValueError, match="AzureProviderConfig"): + AzureProviderStrategy( + config=cast(Any, {"region": "x"}), + logger=logger, + provider_instance_name="azure-default", + ) + + def test_not_initialized_returns_error(self, azure_config, logger): + s = AzureProviderStrategy( + config=azure_config, logger=logger, provider_instance_name="azure-default" + ) + # Do NOT call s.initialize() + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + result = run_operation(s.execute_operation(op)) + assert not result.success + assert result.error_code == "NOT_INITIALIZED" + + def test_execute_operation_propagates_cancellation(self, strategy, monkeypatch): + async def cancelled(_operation): + raise asyncio.CancelledError() + + monkeypatch.setattr(strategy, "_execute_operation_internal", cancelled) + + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + + with pytest.raises(asyncio.CancelledError): + run_operation(strategy.execute_operation(op)) + + def test_execute_operation_delegates_to_execute_operation_async(self, strategy, monkeypatch): + expected = ProviderResult.success_result({"ok": True}) + + async def delegated(_operation): + return expected + + monkeypatch.setattr(strategy, "execute_operation_async", delegated) + + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result == expected + + +class TestCapacityMetadata: + def test_describe_resource_instances_surfaces_vmss_errors_without_instances( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock( + return_value=[ + { + "error_code": "ProvisioningStateFailed", + "error_message": "VMSS provisioning failed", + } + ] + ) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + return_value={ + "capacity": 3, + "provisioned_instance_count": 0, + "provisioning_state": "Failed", + } + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["fleet_errors"][0]["error_code"] == "ProvisioningStateFailed" + + def test_describe_resource_instances_surfaces_single_vm_deployment_errors_without_instances( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["SingleVM"] = handler + deployment_service = MagicMock() + deployment_service.get_deployment_status_async = AsyncMock( + return_value={ + "provisioning_state": "Failed", + "error_code": "DeploymentFailed", + "error_message": "Deployment failed during validation", + } + ) + strategy_harness.deployment_service = deployment_service + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vm-a", "vm-b"], + "provider_api": "SingleVM", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "deployment_name": "dep-singlevm-1", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["deployment_name"] == "dep-singlevm-1" + assert result.metadata["deployment_provisioning_state"] == "Failed" + assert result.metadata["fleet_errors"][0]["error_code"] == "DeploymentFailed" + + def test_describe_resource_instances_returns_canonical_machine_shape(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "vmss-demo_000001", + "status": "running", + "private_ip": "10.0.0.4", + "public_ip": None, + "instance_type": "Standard_B1s", + "subnet_id": "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Network/virtualNetworks/vnet/subnets/default", + "vpc_id": "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Network/virtualNetworks/vnet", + "provider_data": {"vmss_name": "vmss-demo"}, + } + ] + ) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + return_value={ + "capacity": 1, + "provisioned_instance_count": 1, + "provisioning_state": "Succeeded", + } + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["instances"][0]["instance_id"] == "vmss-demo_000001" + assert result.data["instances"][0]["status"] == "running" + assert "InstanceId" not in result.data["instances"][0] + + def test_get_instance_status_reconciles_empty_flexible_vmss_return(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.compute_client = compute_client + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=0) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["vmss-demo_000001"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "resource_id": "vmss-demo", + "resource_mapping": {"vmss-demo_000001": ("vmss-demo", 1)}, + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vmss-demo_000001"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + compute_client.virtual_machine_scale_sets.begin_delete.assert_awaited_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-demo", + ) + assert result.metadata["termination_follow_up_pending"] is True + + def test_describe_resource_instances_adds_shortfall_summary(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "vm-1", + "status": "running", + "private_ip": None, + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": None, + "vpc_id": None, + "provider_data": { + "fleet_errors": [ + { + "error_code": "AllocationFailed", + "error_message": "No capacity in selected zone", + } + ] + }, + } + ] + ) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + return_value={ + "capacity": 3, + "provisioned_instance_count": 1, + "provisioning_state": "Updating", + } + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["capacity_shortfall"]["missing_capacity_units"] == 2 + assert result.metadata["capacity_shortfall"]["likely_causes"] == ["AllocationFailed"] + + def test_describe_resource_instances_uses_request_metadata_resource_group_for_capacity( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + return_value={ + "capacity": 2, + "provisioned_instance_count": 0, + "provisioning_state": "Updating", + } + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "custom-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + resource_manager.get_vmss_capacity_async.assert_awaited_once_with( + "custom-rg", + "vmss-demo", + ) + + def test_describe_resource_instances_uses_request_metadata_resource_group_when_present( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + return_value={ + "capacity": 2, + "provisioned_instance_count": 0, + "provisioning_state": "Updating", + } + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "context-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + resource_manager.get_vmss_capacity_async.assert_awaited_once_with( + "context-rg", + "vmss-demo", + ) + + def test_describe_resource_instances_aggregates_capacity_for_multiple_vmss( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_capacity_async = AsyncMock( + side_effect=[ + { + "capacity": 4, + "provisioned_instance_count": 2, + "provisioning_state": "Updating", + }, + { + "capacity": 3, + "provisioned_instance_count": 1, + "provisioning_state": "Updating", + }, + ] + ) + strategy_harness.resource_manager = resource_manager + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-a", "vmss-b"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["fleet_capacity_fulfilment"] == { + "target_capacity_units": 7, + "fulfilled_capacity_units": 3, + "provisioned_instance_count": 3, + "state": "Updating", + } + assert result.metadata["fleet_capacity_fulfilment_by_resource"] == { + "vmss-a": { + "target_capacity_units": 4, + "fulfilled_capacity_units": 2, + "provisioned_instance_count": 2, + "state": "Updating", + }, + "vmss-b": { + "target_capacity_units": 3, + "fulfilled_capacity_units": 1, + "provisioned_instance_count": 1, + "state": "Updating", + }, + } + + +# --------------------------------------------------------------------------- +# VALIDATE_TEMPLATE +# --------------------------------------------------------------------------- + + +class TestValidateTemplate: + def test_valid_template(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.VALIDATE_TEMPLATE, + parameters={ + "template_config": { + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + }, + ) + result = run_operation(strategy.execute_operation(op)) + assert result.success + assert result.data["valid"] is True + + def test_invalid_template_missing_fields(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.VALIDATE_TEMPLATE, + parameters={ + "template_config": {"template_id": "bad"}, + }, + ) + result = run_operation(strategy.execute_operation(op)) + assert result.success # validation itself doesn't fail the operation + assert result.data["valid"] is False + assert len(result.data["errors"]) > 0 + + def test_missing_template_config(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.VALIDATE_TEMPLATE, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_TEMPLATE_CONFIG" + + +# --------------------------------------------------------------------------- +# GET_AVAILABLE_TEMPLATES +# --------------------------------------------------------------------------- + + +class TestGetAvailableTemplates: + def test_returns_fallback_templates(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.GET_AVAILABLE_TEMPLATES, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert result.success + assert result.data["count"] >= 1 + + def test_fallback_templates_validate_as_azure_templates(self, strategy): + for template in strategy._template_catalog_service.get_fallback_templates(): + AzureTemplate.model_validate(template) + + +# --------------------------------------------------------------------------- +# UNSUPPORTED_OPERATION +# --------------------------------------------------------------------------- + + +class TestUnsupportedOperation: + def test_unsupported_operation(self, strategy): + op = ProviderOperation( + operation_type=cast(Any, "totally_unknown"), + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "UNSUPPORTED_OPERATION" + + +class TestSpotPlacementScoreAdapter: + def test_score_candidates_uses_template_location_value_object(self, logger): + adapter = AzureSpotPlacementScoreAdapter( + azure_client=MagicMock(), + logger=logger, + subscription_id="12345678-1234-1234-1234-123456789012", + base_location="westeurope", + ) + template = AzureTemplate( + template_id="azure-spot-score-test", + provider_api="VMSS", + vm_size="Standard_D4s_v5", + resource_group="test-rg", + location="eastus2", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + image={ + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + ) + + adapter._fetch_scores_async = AsyncMock(return_value={}) + + scores = adapter.score_candidates( + requested_count=2, + template=cast(Any, template), + ) + + assert [score.candidate.region for score in scores] == ["eastus2"] + adapter._fetch_scores_async.assert_awaited_once_with( + requested_count=2, + regions=["eastus2"], + vm_sizes=template.candidate_vm_sizes, + zones=[], + ) + + @pytest.mark.asyncio + async def test_score_candidates_async_uses_async_credential_and_http_payload( + self, logger, monkeypatch + ): + posted_requests = [] + + class FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return { + "placementScores": [ + { + "region": "eastus2", + "availabilityZone": "1", + "sku": "Standard_D4s_v5", + "score": "High", + "isQuotaAvailable": True, + } + ] + } + + class FakeAsyncClient: + def __init__(self, *, timeout): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def post(self, url, *, json, headers): + posted_requests.append( + { + "url": url, + "json": json, + "headers": headers, + "timeout": self.timeout, + } + ) + return FakeResponse() + + monkeypatch.setattr( + "orb.providers.azure.infrastructure.services." + "spot_placement_score_adapter.httpx.AsyncClient", + FakeAsyncClient, + ) + credential = MagicMock() + credential.get_token = AsyncMock(return_value=SimpleNamespace(token="token-1")) + azure_client = MagicMock() + azure_client.get_async_credential = AsyncMock(return_value=credential) + adapter = AzureSpotPlacementScoreAdapter( + azure_client=azure_client, + logger=logger, + subscription_id="sub-1", + base_location="westeurope", + ) + template = AzureTemplate( + template_id="azure-spot-score-http-test", + provider_api="VMSS", + vm_size="Standard_D4s_v5", + resource_group="test-rg", + location="eastus2", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + image={ + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + placement_zones=["1"], + ) + + scores = await adapter.score_candidates_async( + requested_count=3, + template=cast(Any, template), + ) + + assert len(scores) == 1 + assert scores[0].candidate.region == "eastus2" + assert scores[0].candidate.zone == "1" + assert scores[0].candidate.instance_type == "Standard_D4s_v5" + assert scores[0].raw_score == "High" + assert scores[0].normalized_score == 1.0 + assert scores[0].metadata["is_quota_available"] is True + azure_client.get_async_credential.assert_awaited_once_with() + credential.get_token.assert_awaited_once_with("https://management.azure.com/.default") + assert posted_requests == [ + { + "url": ( + "https://management.azure.com/subscriptions/sub-1/providers/" + "Microsoft.Compute/locations/westeurope/placementScores/spot/" + "generate?api-version=2025-02-01-preview" + ), + "json": { + "desiredLocations": ["eastus2"], + "desiredSizes": [{"sku": "Standard_D4s_v5"}], + "desiredCount": 3, + "availabilityZones": True, + }, + "headers": { + "Authorization": "Bearer token-1", + "Content-Type": "application/json", + }, + "timeout": 30.0, + } + ] + + @pytest.mark.asyncio + async def test_score_candidates_async_returns_low_scores_on_http_failure( + self, logger, monkeypatch + ): + class FakeResponse: + def raise_for_status(self): + raise RuntimeError("boom") + + class FakeAsyncClient: + def __init__(self, *, timeout): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def post(self, url, *, json, headers): + return FakeResponse() + + monkeypatch.setattr( + "orb.providers.azure.infrastructure.services." + "spot_placement_score_adapter.httpx.AsyncClient", + FakeAsyncClient, + ) + credential = MagicMock() + credential.get_token = AsyncMock(return_value=SimpleNamespace(token="token-1")) + azure_client = MagicMock() + azure_client.get_async_credential = AsyncMock(return_value=credential) + adapter = AzureSpotPlacementScoreAdapter( + azure_client=azure_client, + logger=logger, + subscription_id="sub-1", + base_location="westeurope", + ) + template = AzureTemplate( + template_id="azure-spot-score-failure-test", + provider_api="VMSS", + vm_size="Standard_D4s_v5", + resource_group="test-rg", + location="eastus2", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + image={ + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + ) + + scores = await adapter.score_candidates_async( + requested_count=1, + template=cast(Any, template), + ) + + assert len(scores) == 1 + assert scores[0].raw_score == "Low" + assert scores[0].normalized_score == 0.2 + assert scores[0].metadata["raw_entry"] == {} + + +# --------------------------------------------------------------------------- +# CREATE_INSTANCES (with missing config → error path) +# --------------------------------------------------------------------------- + + +class TestProviderNaming: + def test_generate_provider_name(self, strategy): + name = strategy.generate_provider_name( + { + "subscription_id": "12345678-abcd", + "region": "westeurope", + } + ) + assert name == "azure_12345678-abcd_westeurope" + + def test_parse_provider_name(self, strategy): + parsed = strategy.parse_provider_name("azure_12345678_westeurope") + assert parsed["type"] == "azure" + assert parsed["subscription_id"] == "12345678" + assert parsed["region"] == "westeurope" + + def test_get_provider_name_pattern(self, strategy): + pattern = strategy.get_provider_name_pattern() + assert "{type}" in pattern + assert "{region}" in pattern + + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + + +class TestCleanup: + def test_cleanup_closes_owned_azure_client(self, azure_config, logger): + client = MagicMock() + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=lambda: client, + ) + strategy.initialize() + assert strategy.azure_client is client + + strategy.cleanup() + + client.close.assert_called_once_with() + + @pytest.mark.asyncio + async def test_cleanup_async_awaits_owned_azure_client_aclose(self, azure_config, logger): + client = MagicMock() + client.aclose = AsyncMock() + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=lambda: client, + ) + strategy.initialize() + assert strategy.azure_client is client + + await strategy.cleanup_async() + + client.aclose.assert_awaited_once_with() + assert strategy.azure_client is None + + def test_cleanup_waits_for_in_flight_operation(self, azure_config, logger, monkeypatch): + client = MagicMock() + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=lambda: client, + ) + strategy.initialize() + assert strategy.azure_client is client + + operation_started = threading.Event() + release_operation = threading.Event() + cleanup_finished = threading.Event() + result_holder: list[ProviderResult] = [] + + async def block_operation(_operation): + operation_started.set() + while not release_operation.is_set(): + await asyncio.sleep(0.01) + return ProviderResult.success_result({"ok": True}) + + monkeypatch.setattr(strategy, "_execute_operation_internal", block_operation) + + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + + operation_thread = threading.Thread( + target=lambda: result_holder.append(run_operation(strategy.execute_operation(op))) + ) + operation_thread.start() + assert operation_started.wait(timeout=1) + + cleanup_thread = threading.Thread( + target=lambda: (strategy.cleanup(), cleanup_finished.set()) + ) + cleanup_thread.start() + + time.sleep(0.05) + assert not cleanup_finished.is_set() + client.close.assert_not_called() + + release_operation.set() + operation_thread.join(timeout=1) + cleanup_thread.join(timeout=1) + + assert cleanup_finished.is_set() + assert result_holder[0].success is True + client.close.assert_called_once_with() + + def test_cleanup_times_out_with_in_flight_operation(self, azure_config, logger, monkeypatch): + client = MagicMock() + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=lambda: client, + ) + strategy.initialize() + assert strategy.azure_client is client + strategy._cleanup_wait_timeout_seconds = 0.05 + + operation_started = threading.Event() + release_operation = threading.Event() + cleanup_finished = threading.Event() + result_holder: list[ProviderResult] = [] + + async def block_operation(_operation): + operation_started.set() + while not release_operation.is_set(): + await asyncio.sleep(0.01) + return ProviderResult.success_result({"ok": True}) + + monkeypatch.setattr(strategy, "_execute_operation_internal", block_operation) + + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + + operation_thread = threading.Thread( + target=lambda: result_holder.append(run_operation(strategy.execute_operation(op))) + ) + operation_thread.start() + assert operation_started.wait(timeout=1) + + cleanup_thread = threading.Thread( + target=lambda: (strategy.cleanup(), cleanup_finished.set()) + ) + cleanup_thread.start() + cleanup_thread.join(timeout=1) + + assert cleanup_finished.is_set() + assert strategy._cleanup_requested is True + assert strategy._runtime.azure_client is client + client.close.assert_not_called() + + release_operation.set() + operation_thread.join(timeout=1) + assert result_holder[0].success is True + + def test_cleanup_drops_cached_handler_factory_and_client(self, azure_config, logger): + client_one = MagicMock() + client_two = MagicMock() + resolver_calls = 0 + + def resolve_client(): + nonlocal resolver_calls + resolver_calls += 1 + return client_one if resolver_calls == 1 else client_two + + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=resolve_client, + ) + strategy.initialize() + + first_handler = strategy.resolve_handler(AzureProviderApi.VMSS) + assert first_handler is not None + assert strategy.azure_client is client_one + + strategy.cleanup() + strategy.initialize() + + second_handler = strategy.resolve_handler(AzureProviderApi.VMSS) + assert second_handler is not None + assert strategy.azure_client is client_two + assert second_handler is not first_handler + client_one.close.assert_called_once_with() + + def test_execute_operation_rejects_new_work_after_cleanup_starts( + self, azure_config, logger, monkeypatch + ): + client = MagicMock() + strategy = AzureProviderStrategy( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + azure_client_resolver=lambda: client, + ) + strategy.initialize() + assert strategy.azure_client is client + + operation_started = threading.Event() + release_operation = threading.Event() + + async def block_operation(_operation): + operation_started.set() + while not release_operation.is_set(): + await asyncio.sleep(0.01) + return ProviderResult.success_result({"ok": True}) + + monkeypatch.setattr(strategy, "_execute_operation_internal", block_operation) + + op = ProviderOperation( + operation_type=ProviderOperationType.HEALTH_CHECK, + parameters={}, + ) + + operation_thread = threading.Thread( + target=lambda: run_operation(strategy.execute_operation(op)) + ) + operation_thread.start() + assert operation_started.wait(timeout=1) + + cleanup_thread = threading.Thread(target=strategy.cleanup) + cleanup_thread.start() + time.sleep(0.05) + + rejected = run_operation(strategy.execute_operation(op)) + + release_operation.set() + operation_thread.join(timeout=1) + cleanup_thread.join(timeout=1) + + assert not rejected.success + assert rejected.error_code == "STRATEGY_SHUTTING_DOWN" diff --git a/tests/providers/azure/test_azure_strategy_lifecycle.py b/tests/providers/azure/test_azure_strategy_lifecycle.py new file mode 100644 index 000000000..7340b4010 --- /dev/null +++ b/tests/providers/azure/test_azure_strategy_lifecycle.py @@ -0,0 +1,1217 @@ +"""Focused tests for Azure strategy create and terminate flows.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import ( + AzureValidationError, + TerminationError, +) +from orb.providers.azure.infrastructure.cyclecloud_session import CycleCloudRequestContext +from orb.providers.azure.infrastructure.handlers.azure_handler import AzureReleaseContext +from orb.providers.azure.infrastructure.handlers.vmss_handler import VMSSHandler +from orb.providers.azure.infrastructure.services.spot_placement_score_adapter import ( + AzureSpotPlacementScoreAdapter, +) +from orb.providers.azure.services.spot_placement_planner import ( + PlacementCandidate, + PlacementPlanEntry, + PlacementScore, +) +from orb.providers.base.strategy import ProviderOperation, ProviderOperationType +from tests.providers.azure.strategy_test_support import build_strategy_harness, run_operation + + +class TestCreateInstances: + def test_missing_template_config_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_TEMPLATE_CONFIG" + + def test_missing_create_handler_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": "VMSS", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "HANDLER_NOT_FOUND" + + def test_dry_run_short_circuits_before_handler(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock() + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": "VMSS", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 2, + }, + context={"dry_run": True}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["resource_ids"] == ["dry-run-resource-id"] + assert result.data["count"] == 2 + assert result.metadata["method"] == "dry_run" + assert result.metadata["provider_data"] == {"dry_run": True} + handler.acquire_hosts_async.assert_not_awaited() + + def test_create_instances_preserves_named_provider_instance_on_synthesized_request( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness( + config=azure_config, + logger=logger, + provider_instance_name="azure-test", + ) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["vmss-demo"], + "instances": [], + "provider_data": {"resource_group": "test-rg"}, + } + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": "VMSS", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + "request_id": "12345678-1234-1234-1234-123456789012", + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["resource_ids"] == ["vmss-demo"] + + @pytest.mark.parametrize("provider_api", ["VMSS", "SingleVM"]) + def test_create_instances_rejects_missing_vm_size(self, azure_config, logger, provider_api): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["azure-resource"], + "instances": [], + "provider_data": {"resource_group": "test-rg"}, + } + ) + strategy_harness.handlers[provider_api] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-minimal-test", + "provider_api": provider_api, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert "vm_size" in result.error_message + handler.acquire_hosts_async.assert_not_called() + + def test_create_instances_uses_image_from_template_dto_metadata_roundtrip( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["azure-resource"], + "instances": [], + } + ) + strategy_harness.handlers["VMSS"] = handler + + azure_template = AzureTemplate( + template_id="azure-roundtrip-test", + provider_api="VMSS", + vm_size="Standard_D4s_v5", + resource_group="test-rg", + location="eastus2", + ssh_public_keys=["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + image={ + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + ) + template_config = azure_template.model_dump(mode="json", exclude_none=True) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": template_config, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.acquire_hosts_async.assert_awaited_once() + + def test_create_instances_accepts_template_dto_provider_config_roundtrip( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["azure-resource"], + "instances": [], + } + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-dto-roundtrip-test", + "provider_type": "azure", + "provider_api": "VMSS", + "image_id": "Canonical:0001-com-ubuntu-server-jammy:22_04-lts-gen2:latest", + "provider_config": { + "resource_group": "template-rg", + "location": "westus2", + "vm_size": "Standard_B1s", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + "node_attributes": {}, + }, + "provider_data": {}, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + forwarded_template = handler.acquire_hosts_async.await_args.args[1] + assert forwarded_template.resource_group.value == "template-rg" + assert forwarded_template.location.value == "westus2" + assert forwarded_template.vm_size == "Standard_B1s" + + def test_create_instances_rejects_non_mapping_provider_config(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["azure-resource"], + "instances": [], + } + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-invalid-provider-config", + "provider_type": "azure", + "provider_api": "VMSS", + "provider_config": object(), + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "AzureValidationError" + assert "Azure provider_config must be a mapping" in result.error_message + handler.acquire_hosts_async.assert_not_called() + + def test_create_instances_accepts_enum_provider_api_in_template_config( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["vmss-demo"], + "instances": [], + } + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": AzureProviderApi.VMSS, + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["provider_api"] == "VMSS" + + def test_create_instances_preserves_azure_validation_failures(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + side_effect=AzureValidationError( + "VM size is not available in this region", + error_code="SkuNotAvailable", + ) + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": "VMSS", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "SkuNotAvailable" + assert result.metadata["error_class"] == "AzureValidationError" + assert result.metadata["provider_error"]["error_code"] == "SkuNotAvailable" + + def test_create_instances_reports_missing_vmss_subnet_as_validation_error( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + strategy_harness.handlers["VMSS"] = VMSSHandler( + azure_client=MagicMock(), + logger=logger, + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "template_config": { + "template_id": "azure-vmss-test", + "provider_api": "VMSS", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "network_config": None, + "subnet_ids": [], + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + "count": 1, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "InvalidParameter" + assert result.metadata["error_class"] == "AzureValidationError" + assert result.metadata["provider_error"]["error_code"] == "InvalidParameter" + assert "No subnet specified" in result.error_message + + +# --------------------------------------------------------------------------- +# TERMINATE_INSTANCES (with missing ids → error path) +# --------------------------------------------------------------------------- + + +class TestTerminateInstances: + def test_missing_instance_ids_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_INSTANCE_IDS" + + def test_missing_provider_api_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={"instance_ids": ["orb-1"]}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_PROVIDER_API" + + def test_missing_resource_id_returns_error(self, strategy_harness): + strategy = strategy_harness.strategy + strategy_harness.handlers["VMSS"] = MagicMock() + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={"instance_ids": ["orb-1"], "provider_api": "VMSS"}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "MISSING_RESOURCE_ID" + + def test_missing_terminate_handler_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={"instance_ids": ["orb-1"], "provider_api": "VMSS"}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "HANDLER_NOT_FOUND" + + def test_fallback_handler_uses_grouped_resource_ids(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock(return_value=None) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": "VMSS", + "resource_mapping": { + "orb-1": ("vmss-prod-b", 1), + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.release_hosts_async.assert_awaited_once() + _, kwargs = handler.release_hosts_async.await_args + assert kwargs["machine_ids"] == ["orb-1"] + assert kwargs["resource_id"] == "vmss-prod-b" + assert kwargs["context"] == AzureReleaseContext( + resource_group="test-rg", + resource_id="vmss-prod-b", + ) + + def test_dry_run_short_circuits_before_release(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock() + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1", "orb-2"], + "provider_api": "VMSS", + }, + context={"dry_run": True}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["terminated_count"] == 2 + assert result.metadata["method"] == "dry_run" + handler.release_hosts_async.assert_not_awaited() + + def test_terminate_instances_records_pending_vmss_cleanup(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + return_value={ + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-1"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + } + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": "VMSS", + "resource_mapping": { + "orb-1": ("vmss-prod-b", 1), + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["provider_data"]["termination_requests"] == [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-1"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + ] + + def test_terminate_instances_merges_pending_vmss_cleanup_for_same_vmss( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=[ + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-1"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + }, + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-2"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + }, + ] + ) + strategy_harness.handlers["VMSS"] = handler + + first_op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": "VMSS", + "resource_mapping": { + "orb-1": ("vmss-prod-b", 1), + }, + }, + ) + second_op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-2"], + "provider_api": "VMSS", + "resource_mapping": { + "orb-2": ("vmss-prod-b", 1), + }, + }, + ) + + first_result = run_operation(strategy.execute_operation(first_op)) + second_result = run_operation(strategy.execute_operation(second_op)) + + assert first_result.success + assert second_result.success + assert first_result.metadata["provider_data"]["termination_requests"] == [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-1"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + ] + assert second_result.metadata["provider_data"]["termination_requests"] == [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-2"], + "delete_vmss_when_empty": False, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + } + ] + + def test_terminate_instances_preserves_provider_failures(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock( + side_effect=TerminationError( + "Azure rejected the delete request", + resource_ids=["vmss-prod-b"], + ) + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": "VMSS", + "resource_mapping": { + "orb-1": ("vmss-prod-b", 1), + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "TerminationError" + assert result.metadata["error_class"] == "TerminationError" + assert result.metadata["provider_error"]["details"]["resource_ids"] == ["vmss-prod-b"] + + def test_get_instance_status_restores_pending_vmss_cleanup_from_request_metadata( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=0) + strategy_harness.resource_manager = resource_manager + + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": "VMSS", + "resource_id": "vmss-prod-b", + "resource_mapping": {"orb-1": ("vmss-prod-b", 1)}, + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-prod-b", + "machine_ids": ["orb-1"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + compute_client.virtual_machine_scale_sets.begin_delete.assert_awaited_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-prod-b", + ) + assert result.metadata["termination_follow_up_pending"] is True + + def test_terminate_instances_forwards_cyclecloud_secret_reference_request_metadata( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock(return_value=None) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["node-1"], + "provider_api": "CycleCloud", + "request_metadata": { + "cluster_name": "my-cluster", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_credential_path": "config/cc.json", + "cyclecloud_verify_ssl": False, + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.release_hosts_async.assert_awaited_once() + _, kwargs = handler.release_hosts_async.await_args + assert kwargs["machine_ids"] == ["node-1"] + assert kwargs["resource_id"] == "my-cluster" + assert kwargs["context"] == AzureReleaseContext( + resource_group="test-rg", + resource_id="my-cluster", + cyclecloud_request_context=CycleCloudRequestContext( + cluster_name="my-cluster", + cyclecloud_url="https://cc.example.com", + cyclecloud_credential_path="config/cc.json", + cyclecloud_verify_ssl=False, + cyclecloud_auth_mode="bearer", + cyclecloud_aad_scope="https://cc.example.com/.default", + ), + ) + + def test_terminate_instances_uses_cyclecloud_context_from_request_metadata( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + ) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock(return_value=None) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["node-1"], + "provider_api": "CycleCloud", + "request_id": "req-11111111-1111-4111-8111-111111111111", + "request_metadata": { + "cluster_name": "my-cluster", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_credential_path": "config/cc.json", + "cyclecloud_verify_ssl": False, + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.release_hosts_async.assert_awaited_once() + _, kwargs = handler.release_hosts_async.await_args + assert kwargs["machine_ids"] == ["node-1"] + assert kwargs["resource_id"] == "my-cluster" + assert kwargs["context"] == AzureReleaseContext( + resource_group="test-rg", + resource_id="my-cluster", + cyclecloud_request_context=CycleCloudRequestContext( + cluster_name="my-cluster", + cyclecloud_url="https://cc.example.com", + cyclecloud_credential_path="config/cc.json", + cyclecloud_verify_ssl=False, + cyclecloud_auth_mode="bearer", + cyclecloud_aad_scope="https://cc.example.com/.default", + ), + ) + + def test_terminate_instances_accepts_enum_provider_api(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.release_hosts_async = AsyncMock(return_value=None) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.TERMINATE_INSTANCES, + parameters={ + "instance_ids": ["orb-1"], + "provider_api": AzureProviderApi.VMSS, + "resource_mapping": {"orb-1": ("vmss-prod-b", 1)}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.release_hosts_async.assert_awaited_once() + _, kwargs = handler.release_hosts_async.await_args + assert kwargs["machine_ids"] == ["orb-1"] + assert kwargs["resource_id"] == "vmss-prod-b" + assert kwargs["context"] == AzureReleaseContext( + resource_group="test-rg", + resource_id="vmss-prod-b", + ) + + +class TestSpotPlacementPlanning: + def test_create_instances_returns_generic_planning_error_for_capacity_exhaustion( + self, strategy_harness, monkeypatch + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": False, + "error_message": "AllocationFailed: No capacity in selected zone", + "provider_data": {"error_codes": ["AllocationFailed"]}, + } + ) + strategy_harness.handlers["VMSS"] = handler + + monkeypatch.setattr( + strategy, + "_build_spot_placement_plan", + lambda template, count: [ + PlacementPlanEntry( + score=PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:1:Standard_D4s_v5", + instance_type="Standard_D4s_v5", + region="eastus2", + zone="1", + ), + raw_score="High", + normalized_score=1.0, + ), + planned_count=2, + ), + ], + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "count": 2, + "request_id": "req-11111111-1111-4111-8111-111111111111", + "template_config": { + "template_id": "tmpl-1", + "provider_api": "VMSS", + "resource_group": "rg1", + "location": "eastus2", + "image": { + "image_id": "/subscriptions/x/resourceGroups/rg/providers/Microsoft.Compute/images/img" + }, + "vm_size": "Standard_D4s_v5", + "vm_sizes": ["Standard_D8s_v5"], + "price_type": "spot", + "priority": "Spot", + "spot_placement_score_enabled": True, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "PROVISIONING_ADAPTER_ERROR" + assert result.error_message == "Spot placement plan could not provision any instances" + + def test_create_instances_spot_plan_uses_async_handler_contract( + self, strategy_harness, monkeypatch + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={"success": True, "resource_ids": ["vmss-b"], "instances": []} + ) + strategy_harness.handlers["VMSS"] = handler + + monkeypatch.setattr( + strategy, + "_build_spot_placement_plan", + lambda template, count: [ + PlacementPlanEntry( + score=PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:1:Standard_D4s_v5", + instance_type="Standard_D4s_v5", + region="eastus2", + zone="1", + ), + raw_score="High", + normalized_score=1.0, + ), + planned_count=1, + ), + ], + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "count": 1, + "request_id": "req-11111111-1111-4111-8111-111111111111", + "template_config": { + "template_id": "tmpl-1", + "provider_api": "VMSS", + "resource_group": "rg1", + "location": "eastus2", + "image": { + "image_id": "/subscriptions/x/resourceGroups/rg/providers/Microsoft.Compute/images/img" + }, + "vm_size": "Standard_D4s_v5", + "vm_sizes": ["Standard_D8s_v5"], + "price_type": "spot", + "priority": "Spot", + "spot_placement_score_enabled": True, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + handler.acquire_hosts_async.assert_awaited_once() + + def test_create_instances_returns_terminal_planning_error_for_non_capacity_failure( + self, strategy_harness, monkeypatch + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": False, + "error_message": "insufficient capacity", + "provider_data": {"error_codes": ["OtherFailure"]}, + } + ) + strategy_harness.handlers["VMSS"] = handler + + monkeypatch.setattr( + strategy, + "_build_spot_placement_plan", + lambda template, count: [ + PlacementPlanEntry( + score=PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:1:Standard_D4s_v5", + instance_type="Standard_D4s_v5", + region="eastus2", + zone="1", + ), + raw_score="High", + normalized_score=1.0, + ), + planned_count=2, + ), + ], + ) + + monkeypatch.setattr( + strategy, + "_is_capacity_like_failure", + lambda result: False, + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "count": 2, + "request_id": "req-11111111-1111-4111-8111-111111111111", + "template_config": { + "template_id": "tmpl-1", + "provider_api": "VMSS", + "resource_group": "rg1", + "location": "eastus2", + "image": { + "image_id": "/subscriptions/x/resourceGroups/rg/providers/Microsoft.Compute/images/img" + }, + "vm_size": "Standard_D4s_v5", + "vm_sizes": ["Standard_D8s_v5"], + "price_type": "spot", + "priority": "Spot", + "spot_placement_score_enabled": True, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "PROVISIONING_ADAPTER_ERROR" + assert result.error_message == "Provisioning failed: insufficient capacity" + + def test_create_instances_uses_planned_handler_path(self, strategy_harness, monkeypatch): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + side_effect=[ + { + "success": False, + "error_message": "AllocationFailed: No capacity in selected zone", + "provider_data": {"error_codes": ["AllocationFailed"]}, + }, + {"success": True, "resource_ids": ["vmss-b"], "instances": []}, + ] + ) + strategy_harness.handlers["VMSS"] = handler + + monkeypatch.setattr( + strategy, + "_build_spot_placement_plan", + lambda template, count: [ + PlacementPlanEntry( + score=PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:1:Standard_D4s_v5", + instance_type="Standard_D4s_v5", + region="eastus2", + zone="1", + ), + raw_score="High", + normalized_score=1.0, + ), + planned_count=2, + ), + PlacementPlanEntry( + score=PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:2:Standard_D8s_v5", + instance_type="Standard_D8s_v5", + region="eastus2", + zone="2", + ), + raw_score="Medium", + normalized_score=0.6, + ), + planned_count=1, + ), + ], + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "count": 2, + "request_id": "req-11111111-1111-4111-8111-111111111111", + "template_config": { + "template_id": "tmpl-1", + "provider_api": "VMSS", + "resource_group": "rg1", + "location": "eastus2", + "image": { + "image_id": "/subscriptions/x/resourceGroups/rg/providers/Microsoft.Compute/images/img" + }, + "vm_size": "Standard_D4s_v5", + "vm_sizes": ["Standard_D8s_v5"], + "price_type": "spot", + "priority": "Spot", + "spot_placement_score_enabled": True, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["resource_ids"] == ["vmss-b"] + assert result.metadata["method"] == "planned_handler" + assert result.metadata["provider_data"]["unfulfilled_count"] == 0 + + def test_create_instances_falls_back_when_scores_are_stale(self, strategy_harness, monkeypatch): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.acquire_hosts_async = AsyncMock( + return_value={ + "success": True, + "resource_ids": ["vmss-fallback"], + "instances": [], + } + ) + strategy_harness.handlers["VMSS"] = handler + + monkeypatch.setattr( + AzureSpotPlacementScoreAdapter, + "score_candidates", + lambda self, requested_count, template: [ + PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:regional:Standard_D4s_v5", + instance_type="Standard_D4s_v5", + region="eastus2", + ), + raw_score="DataNotFoundOrStale", + normalized_score=0.0, + metadata={"raw_entry": {"score": "DataNotFoundOrStale"}}, + ), + PlacementScore( + candidate=PlacementCandidate( + candidate_id="azure:eastus2:regional:Standard_D8s_v5", + instance_type="Standard_D8s_v5", + region="eastus2", + ), + raw_score="DataNotFoundOrStale", + normalized_score=0.0, + metadata={"raw_entry": {"score": "DataNotFoundOrStale"}}, + ), + ], + ) + + op = ProviderOperation( + operation_type=ProviderOperationType.CREATE_INSTANCES, + parameters={ + "count": 2, + "template_config": { + "template_id": "tmpl-stale", + "provider_api": "VMSS", + "resource_group": "rg1", + "location": "eastus2", + "image": { + "image_id": "/subscriptions/x/resourceGroups/rg/providers/Microsoft.Compute/images/img" + }, + "vm_size": "Standard_D4s_v5", + "vm_sizes": ["Standard_D8s_v5"], + "price_type": "spot", + "priority": "Spot", + "spot_placement_score_enabled": True, + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCu"], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["resource_ids"] == ["vmss-fallback"] + assert result.metadata["method"] == "planned_handler" + serialized_plan = result.metadata["provider_data"]["placement_plan"] + assert serialized_plan[0]["instance_type"] == "Standard_D4s_v5" + assert serialized_plan[0]["planned_count"] == 2 + assert serialized_plan[0]["approximate"] is True + assert serialized_plan[0]["metadata"]["fallback_reason"] == "no_viable_provider_scores" diff --git a/tests/providers/azure/test_azure_strategy_status.py b/tests/providers/azure/test_azure_strategy_status.py new file mode 100644 index 000000000..dedcbff7a --- /dev/null +++ b/tests/providers/azure/test_azure_strategy_status.py @@ -0,0 +1,987 @@ +"""Focused tests for Azure strategy status and discovery flows.""" + +from unittest.mock import AsyncMock, MagicMock + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.value_objects import AzureProviderApi +from orb.providers.azure.exceptions.azure_exceptions import CycleCloudConnectionError +from orb.providers.azure.strategy.azure_provider_strategy import AzureProviderStrategy +from orb.providers.base.strategy import ProviderOperation, ProviderOperationType +from tests.providers.azure.strategy_test_support import build_strategy_harness, run_operation + + +class TestGetInstanceStatus: + def test_missing_instance_ids_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_INSTANCE_IDS" + + def test_missing_resource_group_returns_error_when_not_in_request_or_config(self, logger): + strategy = AzureProviderStrategy( + config=AzureProviderConfig( + subscription_id="12345678-1234-1234-1234-123456789012", + resource_group=None, + region="eastus2", + ), + logger=logger, + provider_instance_name="azure-default", + ) + strategy.initialize() + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={"instance_ids": ["vm-1"]}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "MISSING_RESOURCE_GROUP" + + def test_get_instance_status_uses_request_metadata_resource_group(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "vm-1", + "status": "running", + "private_ip": "10.0.0.4", + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": "subnet-1", + "vpc_id": "vnet-1", + "provider_type": "azure", + "provider_data": { + "resource_group": "context-rg", + "vm_name": "vm-1", + }, + } + ] + ) + strategy_harness.handlers["SingleVM"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["vm-1"], + "provider_api": "SingleVM", + "request_metadata": {"resource_group": "context-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["queried_count"] == 1 + assert result.metadata["provider_fulfilment"].state == "fulfilled" + assert result.metadata["provider_fulfilment"].target_units == 1 + handler.check_hosts_status_async.assert_awaited_once() + + def test_dry_run_short_circuits_status_lookup(self, azure_config, logger): + strategy = AzureProviderStrategy( + config=azure_config, logger=logger, provider_instance_name="azure-default" + ) + strategy.initialize() + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={"instance_ids": ["vm-1", "vm-2"]}, + context={"dry_run": True}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["queried_count"] == 2 + assert [m["instance_id"] for m in result.data["instances"]] == ["vm-1", "vm-2"] + assert result.metadata["method"] == "dry_run" + assert result.metadata["provider_fulfilment"].state == "in_progress" + + def test_single_vm_provider_api_routes_status_via_handler(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "vm-1", + "status": "running", + "private_ip": "10.0.0.4", + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": "/subscriptions/.../subnets/default", + "vpc_id": "/subscriptions/.../virtualNetworks/test-vnet", + "provider_type": "azure", + "provider_data": {"vm_name": "vm-1"}, + } + ] + ) + strategy_harness.handlers["SingleVM"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["vm-1"], + "provider_api": "SingleVM", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["method"] == "handler" + handler.check_hosts_status_async.assert_awaited_once() + assert result.data["instances"][0]["instance_id"] == "vm-1" + + def test_vmss_provider_api_routes_status_via_handler_with_resource_mapping( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "3", + "status": "running", + "private_ip": "10.0.0.7", + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": "/subscriptions/.../subnets/default", + "vpc_id": "/subscriptions/.../virtualNetworks/test-vnet", + "provider_type": "azure", + "provider_data": { + "vmss_instance_id": "3", + "vm_id": "vm-guid-3", + }, + }, + { + "instance_id": "9", + "status": "running", + "private_ip": "10.0.0.9", + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": "/subscriptions/.../subnets/default", + "vpc_id": "/subscriptions/.../virtualNetworks/test-vnet", + "provider_type": "azure", + "provider_data": { + "vmss_instance_id": "9", + "vm_id": "vm-guid-9", + }, + }, + ] + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["3"], + "provider_api": "VMSS", + "request_metadata": {"resource_group": "test-rg"}, + "resource_mapping": {"3": ("vmss-demo", 2)}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["method"] == "handler" + handler.check_hosts_status_async.assert_awaited_once() + assert [m["instance_id"] for m in result.data["instances"]] == ["3"] + + def test_vmss_resource_mapping_routes_status_via_handler_with_provider_api( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "3", + "status": "running", + "private_ip": "10.0.0.7", + "public_ip": None, + "launch_time": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": "/subscriptions/.../subnets/default", + "vpc_id": "/subscriptions/.../virtualNetworks/test-vnet", + "provider_type": "azure", + "provider_data": { + "vmss_instance_id": "3", + "vm_id": "vm-guid-3", + }, + } + ] + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["3"], + "provider_api": "VMSS", + "request_metadata": {"resource_group": "test-rg"}, + "resource_mapping": {"3": ("vmss-demo", 2)}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["method"] == "handler" + handler.check_hosts_status_async.assert_awaited_once() + assert [m["instance_id"] for m in result.data["instances"]] == ["3"] + + def test_cyclecloud_status_handler_failure_surfaces_error(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + side_effect=CycleCloudConnectionError( + "cyclecloud auth failed", + url="https://cc.example.com", + ) + ) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["node-1"], + "provider_api": "CycleCloud", + "resource_id": "my-cluster", + "request_metadata": { + "resource_group": "test-rg", + "cyclecloud_url": "https://cc.example.com", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "CycleCloudConnectionError" + assert "cyclecloud auth failed" in result.error_message + assert result.metadata["error_class"] == "CycleCloudConnectionError" + + def test_status_preserves_handler_identity_fields(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "vm-1", + "status": "running", + "private_ip": "10.0.0.4", + "public_ip": None, + "instance_type": "Standard_D4s_v5", + "subnet_id": ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ), + "vpc_id": ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet" + ), + "provider_type": "azure", + "provider_data": {"vm_name": "vm-1"}, + } + ] + ) + strategy_harness.handlers["SingleVM"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["vm-1"], + "provider_api": "SingleVM", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + machine = result.data["instances"][0] + assert machine["instance_id"] == "vm-1" + assert machine["status"] == "running" + assert machine["private_ip"] == "10.0.0.4" + assert machine["instance_type"] == "Standard_D4s_v5" + assert machine["subnet_id"].endswith("/subnets/default") + assert machine["provider_data"]["vm_name"] == "vm-1" + + def test_status_query_without_provider_api_is_rejected(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["vm-1"], + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "MISSING_PROVIDER_API" + + def test_get_instance_status_accepts_enum_provider_api(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "3", + "status": "running", + "provider_type": "azure", + "provider_data": {"vmss_instance_id": "3"}, + } + ] + ) + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["3"], + "provider_api": AzureProviderApi.VMSS, + "request_metadata": {"resource_group": "test-rg"}, + "resource_mapping": {"3": ("vmss-demo", 2)}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert [m["instance_id"] for m in result.data["instances"]] == ["3"] + handler.check_hosts_status_async.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# DESCRIBE_RESOURCE_INSTANCES (with missing resource_ids → error path) +# --------------------------------------------------------------------------- + + +class TestDescribeResourceInstances: + def test_missing_resource_ids_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_RESOURCE_IDS" + + def test_missing_provider_api_returns_error(self, strategy): + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={"resource_ids": ["vmss-demo"]}, + ) + result = run_operation(strategy.execute_operation(op)) + assert not result.success + assert result.error_code == "MISSING_PROVIDER_API" + + def test_describe_resource_instances_cleans_up_empty_vmss_when_requested_members_are_gone( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=0) + strategy_harness.resource_manager = resource_manager + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + compute_client.virtual_machine_scale_sets.begin_delete.assert_awaited_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-demo", + ) + assert result.metadata["termination_follow_up_pending"] is True + assert result.metadata["termination_follow_up_details"] == [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": True, + "delete_retry_pending": False, + } + ] + + def test_describe_resource_instances_clears_pending_cleanup_after_vmss_is_gone( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.vmss_exists_async = AsyncMock(return_value=False) + strategy_harness.resource_manager = resource_manager + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "delete_submitted": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + resource_manager.vmss_exists_async.assert_awaited_once_with( + resource_group="test-rg", + vmss_name="vmss-demo", + ) + compute_client.virtual_machine_scale_sets.begin_delete.assert_not_awaited() + assert result.metadata["termination_follow_up_pending"] is False + + def test_describe_resource_instances_does_not_cleanup_when_strict_vmss_status_fails( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + side_effect=RuntimeError( + "Failed to list instances for VMSS 'vmss-demo': transient ARM failure" + ) + ) + strategy_harness.handlers["VMSS"] = handler + strategy_harness.resource_manager = MagicMock() + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert not result.success + assert result.error_code == "DESCRIBE_RESOURCE_INSTANCES_ERROR" + strategy_harness.resource_manager.get_vmss_member_count_async.assert_not_called() + forwarded_request = handler.check_hosts_status_async.await_args.args[0] + assert forwarded_request.metadata["raise_on_status_error"] is True + + def test_describe_resource_instances_leaves_cleanup_pending_when_other_members_remain( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=1) + strategy_harness.resource_manager = resource_manager + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-b"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-b", + "machine_ids": ["vm-b"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + resource_manager.get_vmss_member_count_async.assert_awaited_once_with( + resource_group="test-rg", + vmss_name="vmss-b", + ) + assert result.metadata["termination_follow_up_pending"] is True + assert result.metadata["termination_follow_up_details"] == [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-b", + "machine_ids": ["vm-b"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + ] + compute_client.virtual_machine_scale_sets.begin_delete.assert_not_awaited() + + def test_describe_resource_instances_leaves_cleanup_pending_when_member_count_is_unknown( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=None) + strategy_harness.resource_manager = resource_manager + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + resource_manager.get_vmss_member_count_async.assert_awaited_once_with( + resource_group="test-rg", + vmss_name="vmss-demo", + ) + assert result.metadata["termination_follow_up_pending"] is True + assert result.metadata["termination_follow_up_details"] == [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + ] + compute_client.virtual_machine_scale_sets.begin_delete.assert_not_awaited() + + def test_describe_resource_instances_surfaces_retry_pending_when_vmss_delete_retry_fails( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + resource_manager.get_vmss_member_count_async = AsyncMock(return_value=0) + strategy_harness.resource_manager = resource_manager + compute_client = MagicMock() + compute_client.virtual_machine_scale_sets.begin_delete = AsyncMock( + side_effect=RuntimeError("delete still blocked") + ) + azure_client = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=compute_client) + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + ], + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["termination_follow_up_pending"] is True + assert result.metadata["termination_follow_up_details"] == [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "delete_retry_count": 1, + "last_delete_error": "delete still blocked", + } + ] + + def test_describe_resource_instances_clears_pending_cleanup_when_no_delete_is_required( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + resource_manager = MagicMock() + strategy_harness.resource_manager = resource_manager + azure_client = MagicMock() + strategy_harness.azure_client = azure_client + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": "VMSS", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": False, + } + } + ], + }, + }, + ) + + first_result = run_operation(strategy.execute_operation(op)) + second_result = run_operation(strategy.execute_operation(op)) + + assert first_result.success + assert first_result.metadata["termination_follow_up_pending"] is False + assert second_result.success + assert second_result.metadata["termination_follow_up_pending"] is False + resource_manager.get_vmss_member_count_async.assert_not_called() + azure_client.get_async_compute_client.assert_not_called() + + def test_describe_resource_instances_forwards_cyclecloud_request_metadata( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["req-12345678-1234-1234-1234-123456789012"], + "provider_api": "CycleCloud", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "cluster_name": "my-cluster", + "node_array": "execute", + "node_ids": ["node-1"], + "operation_id": "op-123", + "operation_location": "https://cc.example.com/operations/op-123", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + "cyclecloud_verify_ssl": False, + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + forwarded_request = handler.check_hosts_status_async.await_args.args[0] + assert forwarded_request.resource_ids == ["req-12345678-1234-1234-1234-123456789012"] + assert forwarded_request.metadata["cluster_name"] == "my-cluster" + assert forwarded_request.metadata["node_array"] == "execute" + assert forwarded_request.metadata["node_ids"] == ["node-1"] + assert forwarded_request.metadata["operation_id"] == "op-123" + assert ( + forwarded_request.metadata["operation_location"] + == "https://cc.example.com/operations/op-123" + ) + assert forwarded_request.metadata["cyclecloud_url"] == "https://cc.example.com" + assert forwarded_request.metadata["cyclecloud_auth_mode"] == "bearer" + assert ( + forwarded_request.metadata["cyclecloud_aad_scope"] == "https://cc.example.com/.default" + ) + assert forwarded_request.metadata["cyclecloud_verify_ssl"] is False + + def test_get_instance_status_matches_cyclecloud_node_name_alias(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock( + return_value=[ + { + "instance_id": "6ecc44d4-417d-41e4-a729-3d504d651fd3", + "name": "dynamic-1", + "status": "running", + "provider_type": "azure", + "provider_data": { + "cluster_name": "contoso-slurm-lab-cluster", + "node_id": "6ecc44d4-417d-41e4-a729-3d504d651fd3", + "node_name": "dynamic-1", + }, + } + ] + ) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["dynamic-1"], + "provider_api": "CycleCloud", + "resource_id": "contoso-slurm-lab-cluster", + "resource_mapping": {"dynamic-1": ("contoso-slurm-lab-cluster", 1)}, + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "cluster_name": "contoso-slurm-lab-cluster", + "node_ids": ["dynamic-1"], + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.data["queried_count"] == 1 + assert len(result.data["instances"]) == 1 + assert result.data["instances"][0]["name"] == "dynamic-1" + + def test_get_instance_status_uses_cyclecloud_cluster_name_when_resource_id_missing( + self, strategy_harness + ): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["dynamic-1"], + "provider_api": "CycleCloud", + "template_id": "tmpl-1", + "request_metadata": { + "resource_group": "test-rg", + "cluster_name": "contoso-slurm-lab-cluster", + "node_ids": ["dynamic-1"], + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + forwarded_request = handler.check_hosts_status_async.await_args.args[0] + assert forwarded_request.resource_ids == ["contoso-slurm-lab-cluster"] + + def test_get_instance_status_uses_cyclecloud_context_from_request_metadata( + self, azure_config, logger + ): + strategy_harness = build_strategy_harness( + config=azure_config, + logger=logger, + provider_instance_name="azure-default", + ) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + strategy_harness.handlers["CycleCloud"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.GET_INSTANCE_STATUS, + parameters={ + "instance_ids": ["dynamic-1"], + "provider_api": "CycleCloud", + "template_id": "tmpl-1", + "request_id": "req-11111111-1111-4111-8111-111111111111", + "request_metadata": { + "resource_group": "test-rg", + "cluster_name": "contoso-slurm-lab-cluster", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + forwarded_request = handler.check_hosts_status_async.await_args.args[0] + assert forwarded_request.resource_ids == ["contoso-slurm-lab-cluster"] + assert forwarded_request.metadata["cluster_name"] == "contoso-slurm-lab-cluster" + + def test_dry_run_short_circuits_resource_discovery(self, azure_config, logger): + strategy_harness = build_strategy_harness(config=azure_config, logger=logger) + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock() + strategy_harness.handlers["VMSS"] = handler + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-test"], + "provider_api": "VMSS", + }, + context={"dry_run": True}, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["method"] == "dry_run" + handler.check_hosts_status_async.assert_not_awaited() + + def test_describe_resource_instances_accepts_enum_provider_api(self, strategy_harness): + strategy = strategy_harness.strategy + handler = MagicMock() + handler.check_hosts_status_async = AsyncMock(return_value=[]) + handler.get_vmss_resource_errors_async = AsyncMock(return_value=[]) + strategy_harness.handlers["VMSS"] = handler + strategy_harness.resource_manager = MagicMock() + + op = ProviderOperation( + operation_type=ProviderOperationType.DESCRIBE_RESOURCE_INSTANCES, + parameters={ + "resource_ids": ["vmss-demo"], + "provider_api": AzureProviderApi.VMSS, + "template_id": "tmpl-1", + "request_metadata": {"resource_group": "test-rg"}, + }, + ) + + result = run_operation(strategy.execute_operation(op)) + + assert result.success + assert result.metadata["provider_api"] == "VMSS" + handler.check_hosts_status_async.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Provider naming +# --------------------------------------------------------------------------- diff --git a/tests/providers/azure/test_azure_template.py b/tests/providers/azure/test_azure_template.py new file mode 100644 index 000000000..7d12a96f9 --- /dev/null +++ b/tests/providers/azure/test_azure_template.py @@ -0,0 +1,596 @@ +"""Tests for the Azure domain template aggregate, value objects, and ARM mapper.""" + +import pytest +from pydantic import ValidationError + +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.domain.template.value_objects import ( + AzureAllocationStrategy, + AzureDataDisk, + AzureEvictionPolicy, + AzureImageReference, + AzureNetworkConfig, + AzureOSDiskConfig, + AzureOSDiskType, + AzurePriority, + AzureProviderApi, + AzureVMSSOrchestrationMode, +) +from orb.providers.azure.infrastructure.services.arm_payload_mapper import ArmPayloadMapper + +# --------------------------------------------------------------------------- +# Required fields +# --------------------------------------------------------------------------- + +_BASE_FIELDS = { + "template_id": "test-template", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, +} + +_SPOT_DEFAULTS = { + "priority": "Spot", + "eviction_policy": "Deallocate", +} + + +# --------------------------------------------------------------------------- +# AzureTemplate basic construction +# --------------------------------------------------------------------------- + + +class TestAzureTemplateConstruction: + def test_minimal_template(self): + t = AzureTemplate(**_BASE_FIELDS) + assert t.template_id == "test-template" + assert t.vm_size == "Standard_D4s_v5" + assert t.resource_group.value == "test-rg" + assert t.location.value == "eastus2" + assert t.provider_api == AzureProviderApi.VMSS + + def test_accepts_template_dto_payload_with_provider_config_and_version(self): + payload = { + "template_id": "test-template", + "provider_type": "azure", + "provider_api": "SingleVM", + "version": "catalog-v1", + "provider_config": { + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + }, + } + + t = AzureTemplate(**payload) + + assert t.provider_api == AzureProviderApi.SINGLE_VM + assert t.version == "catalog-v1" + assert t.image is not None + assert t.image.publisher == "Canonical" + + def test_rejects_missing_ssh_keys(self): + """SSH access is required — no password fallback (mirrors AWS key_name pattern).""" + fields = {**_BASE_FIELDS} + fields.pop("ssh_public_keys", None) + with pytest.raises(ValueError, match="SSH access is required"): + AzureTemplate(**fields) + + def test_rejects_empty_ssh_keys_without_key_name(self): + """Empty ssh_public_keys without ssh_key_name should be rejected.""" + fields = {**_BASE_FIELDS, "ssh_public_keys": []} + with pytest.raises(ValueError, match="SSH access is required"): + AzureTemplate(**fields) + + def test_rejects_missing_image_source(self): + fields = {**_BASE_FIELDS} + fields.pop("image") + with pytest.raises(ValueError, match="image source is required"): + AzureTemplate(**fields) + + def test_ssh_key_name_accepted(self): + """ssh_key_name alone (without inline keys) should pass validation.""" + fields = { + "template_id": "test-template", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "ssh_key_name": "my-azure-ssh-key", + "image": { + "publisher": "Canonical", + "offer": "0001-com-ubuntu-server-jammy", + "sku": "22_04-lts-gen2", + "version": "latest", + }, + } + t = AzureTemplate(**fields) + assert t.ssh_key_name == "my-azure-ssh-key" + + def test_with_custom_image_id(self): + fields = {**_BASE_FIELDS} + fields.pop("image", None) + t = AzureTemplate( + **fields, + image={"image_id": "/subscriptions/.../images/my-image"}, + ) + assert t.image is not None + assert t.image.image_id == "/subscriptions/.../images/my-image" + + def test_rejects_invalid_location_slug(self): + fields = {**_BASE_FIELDS, "location": "East US 2"} + with pytest.raises(ValueError, match="location must be an Azure location slug"): + AzureTemplate(**fields) + + def test_rejects_invalid_arm_identifier_type(self): + with pytest.raises(ValueError, match="diskEncryptionSets"): + AzureTemplate( + **_BASE_FIELDS, + disk_encryption_set_id=( + "/subscriptions/sub/resourceGroups/rg/providers/" + "Microsoft.Compute/proximityPlacementGroups/ppg-1" + ), + ) + + def test_accepts_mixed_case_arm_identifier_segments(self): + t = AzureTemplate( + **_BASE_FIELDS, + disk_encryption_set_id=( + "/SUBSCRIPTIONS/sub/RESOURCEGROUPS/rg/PROVIDERS/" + "MICROSOFT.COMPUTE/DISKENCRYPTIONSETS/des-1" + ), + ) + assert ( + str(t.disk_encryption_set_id) == "/SUBSCRIPTIONS/sub/RESOURCEGROUPS/rg/PROVIDERS/" + "MICROSOFT.COMPUTE/DISKENCRYPTIONSETS/des-1" + ) + + def test_rejects_invalid_upgrade_policy_mode(self): + with pytest.raises(ValueError, match="upgrade_policy_mode"): + AzureTemplate(**_BASE_FIELDS, upgrade_policy_mode="BlueGreen") + + def test_model_dump_json_serialises_azure_value_objects_as_strings(self): + t = AzureTemplate( + **_BASE_FIELDS, + proximity_placement_group_id=( + "/subscriptions/sub/resourceGroups/rg/providers/" + "Microsoft.Compute/proximityPlacementGroups/ppg-1" + ), + upgrade_policy_mode="Rolling", + ) + + payload = t.model_dump(mode="json", exclude_none=True) + + assert payload["resource_group"] == "test-rg" + assert payload["location"] == "eastus2" + assert ( + payload["proximity_placement_group_id"] + == "/subscriptions/sub/resourceGroups/rg/providers/" + "Microsoft.Compute/proximityPlacementGroups/ppg-1" + ) + assert payload["upgrade_policy_mode"] == "Rolling" + + def test_vmss_uniform_rejects_flexible_orchestration_mode(self): + with pytest.raises(ValueError, match="VMSSUniform"): + AzureTemplate( + **_BASE_FIELDS, + provider_api=AzureProviderApi.VMSS_UNIFORM, + ) + + def test_vmss_uniform_accepts_uniform_orchestration_mode(self): + t = AzureTemplate( + **_BASE_FIELDS, + provider_api=AzureProviderApi.VMSS_UNIFORM, + orchestration_mode=AzureVMSSOrchestrationMode.UNIFORM, + ) + assert t.provider_api == AzureProviderApi.VMSS_UNIFORM + assert t.orchestration_mode == AzureVMSSOrchestrationMode.UNIFORM + + def test_provider_type_forced_to_azure(self): + t = AzureTemplate(**_BASE_FIELDS, provider_type="wrong") + assert t.provider_type == "azure" + + def test_template_is_frozen_after_construction(self): + template = AzureTemplate(**_BASE_FIELDS) + + with pytest.raises(ValidationError): + template.vm_size = "Standard_D8s_v5" + + def test_from_azure_format(self): + data = { + **_BASE_FIELDS, + "vmSize": "Standard_D2s_v5", + } + data.pop("vm_size") + t = AzureTemplate.from_azure_format(data) + assert t.vm_size == "Standard_D2s_v5" + + def test_from_azure_format_accepts_max_number(self): + data = { + **_BASE_FIELDS, + "max_number": 7, + } + t = AzureTemplate.from_azure_format(data) + assert t.max_instances == 7 + + def test_from_azure_format_prefers_explicit_max_instances_over_max_number(self): + data = { + **_BASE_FIELDS, + "max_instances": 3, + "max_number": 7, + } + t = AzureTemplate.from_azure_format(data) + assert t.max_instances == 3 + + def test_rejects_both_provider_api_spec_and_file(self): + with pytest.raises(ValueError, match="provider_api_spec and provider_api_spec_file"): + AzureTemplate( + **_BASE_FIELDS, + provider_api_spec={"location": "eastus2"}, + provider_api_spec_file="vmss.json", + ) + + +# --------------------------------------------------------------------------- +# Spot / priority validation +# --------------------------------------------------------------------------- + + +class TestSpotValidation: + def test_spot_defaults_eviction_policy(self): + """mode='before' validator fills in implied defaults for Spot.""" + t = AzureTemplate(**_BASE_FIELDS, priority="Spot") + assert t.eviction_policy == AzureEvictionPolicy.DEALLOCATE + + def test_spot_accepts_explicit_values(self): + t = AzureTemplate(**_BASE_FIELDS, **_SPOT_DEFAULTS) + assert t.eviction_policy == AzureEvictionPolicy.DEALLOCATE + + def test_spot_percentage_implies_spot_priority(self): + t = AzureTemplate(**_BASE_FIELDS, spot_percentage=70) + assert t.priority == AzurePriority.SPOT + assert t.eviction_policy == AzureEvictionPolicy.DEALLOCATE + + def test_spot_percentage_overrides_regular_priority_default(self): + t = AzureTemplate(**_BASE_FIELDS, priority="Regular", spot_percentage=70) + assert t.priority == AzurePriority.SPOT + assert t.eviction_policy == AzureEvictionPolicy.DEALLOCATE + + def test_spot_percentage_requires_flexible(self): + with pytest.raises(ValueError, match="Flexible orchestration mode"): + AzureTemplate( + **_BASE_FIELDS, + spot_percentage=70, + orchestration_mode=AzureVMSSOrchestrationMode.UNIFORM, + ) + + def test_spot_percentage_rejects_single_placement_group(self): + with pytest.raises(ValueError, match="single_placement_group"): + AzureTemplate( + **_BASE_FIELDS, + spot_percentage=70, + single_placement_group=True, + ) + + def test_regular_rejects_eviction_policy(self): + with pytest.raises(ValueError, match="eviction_policy"): + AzureTemplate(**_BASE_FIELDS, priority="Regular", eviction_policy="Delete") + + def test_regular_rejects_billing_max_price(self): + with pytest.raises(ValueError, match="billing_profile_max_price"): + AzureTemplate(**_BASE_FIELDS, priority="Regular", billing_profile_max_price=1.0) + + +# --------------------------------------------------------------------------- +# Security validation +# --------------------------------------------------------------------------- + + +class TestSecurityValidation: + def test_trusted_launch_defaults(self): + """mode='before' validator fills in secure_boot and vtpm for TrustedLaunch.""" + t = AzureTemplate(**_BASE_FIELDS, security_type="TrustedLaunch") + assert t.secure_boot_enabled is True + assert t.vtpm_enabled is True + + def test_trusted_launch_accepts_explicit_values(self): + t = AzureTemplate( + **_BASE_FIELDS, + security_type="TrustedLaunch", + secure_boot_enabled=True, + vtpm_enabled=True, + ) + assert t.secure_boot_enabled is True + assert t.vtpm_enabled is True + + +# --------------------------------------------------------------------------- +# Zone validation +# --------------------------------------------------------------------------- + + +class TestZoneValidation: + def test_zone_balance_requires_zones(self): + with pytest.raises(ValueError, match="zone_balance"): + AzureTemplate(**_BASE_FIELDS, zone_balance=True) + + def test_zone_balance_with_zones(self): + t = AzureTemplate(**_BASE_FIELDS, zone_balance=True, zones=["1", "2"]) + assert t.zone_balance is True + + def test_overprovision_rejected_for_flexible(self): + with pytest.raises(ValueError, match="overprovision"): + AzureTemplate( + **_BASE_FIELDS, + orchestration_mode=AzureVMSSOrchestrationMode.FLEXIBLE, + overprovision=True, + ) + + def test_overprovision_allowed_for_uniform(self): + t = AzureTemplate( + **_BASE_FIELDS, + orchestration_mode=AzureVMSSOrchestrationMode.UNIFORM, + overprovision=True, + ) + assert t.overprovision is True + + +# --------------------------------------------------------------------------- +# ARM payload generation +# --------------------------------------------------------------------------- + + +class TestArmPayload: + def test_basic_arm_payload(self): + t = AzureTemplate( + **_BASE_FIELDS, + network_config={"subnet_id": "/subscriptions/.../subnets/default"}, + ) + arm = ArmPayloadMapper.vmss_payload(t) + + assert arm["location"] == "eastus2" + assert arm["sku"]["name"] == "Standard_D4s_v5" + assert "virtualMachineProfile" in arm["properties"] + vm_profile = arm["properties"]["virtualMachineProfile"] + assert vm_profile["storageProfile"]["osDisk"]["deleteOption"] == "Delete" + nic_config = vm_profile["networkProfile"]["networkInterfaceConfigurations"][0] + assert nic_config["properties"]["deleteOption"] == "Delete" + + def test_spot_arm_payload(self): + t = AzureTemplate( + **_BASE_FIELDS, + priority="Spot", + billing_profile_max_price=-1.0, + ) + arm = ArmPayloadMapper.vmss_payload(t) + vm_profile = arm["properties"]["virtualMachineProfile"] + assert vm_profile["priority"] == "Spot" + assert vm_profile["billingProfile"]["maxPrice"] == -1.0 + + def test_vmss_mix_payload_uses_sku_profile(self): + t = AzureTemplate( + **_BASE_FIELDS, + vm_sizes=["Standard_D8s_v5", "Standard_D16s_v5"], + vmss_allocation_strategy=AzureAllocationStrategy.CAPACITY_OPTIMIZED, + ) + arm = ArmPayloadMapper.vmss_payload(t) + + assert arm["sku"]["name"] == "Mix" + assert arm["properties"]["skuProfile"]["vmSizes"] == [ + {"name": "Standard_D4s_v5"}, + {"name": "Standard_D8s_v5"}, + {"name": "Standard_D16s_v5"}, + ] + assert arm["properties"]["skuProfile"]["allocationStrategy"] == "CapacityOptimized" + assert ( + "vmSizeProperties" not in arm["properties"]["virtualMachineProfile"]["hardwareProfile"] + ) + + def test_vmss_prioritized_mix_includes_ranks(self): + t = AzureTemplate( + **_BASE_FIELDS, + vm_size_preferences=[ + {"name": "Standard_D8s_v5", "rank": 2}, + {"name": "Standard_D16s_v5"}, + ], + vmss_allocation_strategy=AzureAllocationStrategy.PRIORITIZED, + ) + + arm = ArmPayloadMapper.vmss_payload(t) + + assert arm["sku"]["name"] == "Mix" + assert arm["properties"]["skuProfile"]["allocationStrategy"] == "Prioritized" + assert arm["properties"]["skuProfile"]["vmSizes"] == [ + {"name": "Standard_D4s_v5", "rank": 0}, + {"name": "Standard_D8s_v5", "rank": 2}, + {"name": "Standard_D16s_v5", "rank": 1}, + ] + + def test_prioritized_mix_requires_ranked_preferences(self): + with pytest.raises(ValueError, match="requires at least one vm_size_preference"): + AzureTemplate( + **_BASE_FIELDS, + vm_sizes=["Standard_D8s_v5", "Standard_D16s_v5"], + vmss_allocation_strategy=AzureAllocationStrategy.PRIORITIZED, + ) + + def test_vm_size_preferences_are_rank_only(self): + with pytest.raises(ValueError, match="vm_size_preferences is only valid"): + AzureTemplate( + **_BASE_FIELDS, + vm_size_preferences=[{"name": "Standard_D8s_v5"}], + vmss_allocation_strategy=AzureAllocationStrategy.CAPACITY_OPTIMIZED, + ) + + def test_vm_sizes_and_vm_size_preferences_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="Specify either vm_sizes or vm_size_preferences"): + AzureTemplate( + **_BASE_FIELDS, + vm_sizes=["Standard_D8s_v5"], + vm_size_preferences=[{"name": "Standard_D16s_v5", "rank": 1}], + vmss_allocation_strategy=AzureAllocationStrategy.PRIORITIZED, + ) + + def test_vmss_mix_without_vmss_allocation_strategy_omits_allocation_strategy(self): + t = AzureTemplate(**_BASE_FIELDS, priority="Spot", vm_sizes=["Standard_D8s_v5"]) + + arm = ArmPayloadMapper.vmss_payload(t) + + assert "allocationStrategy" not in arm["properties"]["skuProfile"] + + def test_spot_percentage_populates_priority_mix_policy(self): + t = AzureTemplate( + **_BASE_FIELDS, + vm_sizes=["Standard_D8s_v5"], + priority="Spot", + spot_percentage=70, + base_regular_priority_count=2, + ) + arm = ArmPayloadMapper.vmss_payload(t) + + vm_profile = arm["properties"]["virtualMachineProfile"] + assert vm_profile["priority"] == "Spot" + assert arm["properties"]["priorityMixPolicy"] == { + "baseRegularPriorityCount": 2, + "regularPriorityPercentageAboveBase": 30, + } + + def test_zones_in_arm_payload(self): + t = AzureTemplate( + **_BASE_FIELDS, + zones=["1", "2", "3"], + ) + arm = ArmPayloadMapper.vmss_payload(t) + assert arm["zones"] == ["1", "2", "3"] + + def test_identity_in_arm_payload(self): + t = AzureTemplate( + **_BASE_FIELDS, + system_assigned_identity=True, + user_assigned_identity_ids=["/subscriptions/.../identities/my-id"], + ) + arm = ArmPayloadMapper.vmss_payload(t) + assert arm["identity"]["type"] == "SystemAssigned, UserAssigned" + + def test_disk_encryption_set_is_applied_to_vmss_os_and_data_disks(self): + t = AzureTemplate( + **_BASE_FIELDS, + disk_encryption_set_id="/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Compute/diskEncryptionSets/des-1", + data_disks=[{"lun": 0, "disk_size_gb": 256}], + ) + arm = ArmPayloadMapper.vmss_payload(t) + storage_profile = arm["properties"]["virtualMachineProfile"]["storageProfile"] + + assert storage_profile["osDisk"]["managedDisk"]["diskEncryptionSet"] == { + "id": "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Compute/diskEncryptionSets/des-1" + } + assert storage_profile["dataDisks"][0]["managedDisk"]["diskEncryptionSet"] == { + "id": "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Compute/diskEncryptionSets/des-1" + } + + def test_single_vm_payload_reuses_storage_profile_builder_with_single_vm_defaults(self): + t = AzureTemplate( + **_BASE_FIELDS, + provider_api=AzureProviderApi.SINGLE_VM, + ) + + params = ArmPayloadMapper.single_vm_payload( + template=t, + vm_name="vm-test", + nic_id="/subscriptions/.../networkInterfaces/nic-vm-test", + ) + storage_profile = params["properties"]["storageProfile"] + + assert storage_profile["imageReference"]["publisher"] == "Canonical" + assert storage_profile["osDisk"]["managedDisk"]["storageAccountType"] == "Standard_LRS" + assert "caching" not in storage_profile["osDisk"] + + +# --------------------------------------------------------------------------- +# Value objects +# --------------------------------------------------------------------------- + + +class TestValueObjects: + def test_image_reference_marketplace(self): + img = AzureImageReference( + publisher="Canonical", + offer="0001-com-ubuntu-server-jammy", + sku="22_04-lts-gen2", + ) + arm = img.to_arm_dict() + assert arm["publisher"] == "Canonical" + assert arm["version"] == "latest" + + def test_image_reference_custom(self): + img = AzureImageReference(image_id="/subscriptions/.../images/myimg") + arm = img.to_arm_dict() + assert arm["id"] == "/subscriptions/.../images/myimg" + + def test_image_reference_requires_one_source(self): + with pytest.raises(ValueError): + AzureImageReference() + + def test_os_disk_config(self): + disk = AzureOSDiskConfig( + disk_size_gb=128, + storage_account_type=AzureOSDiskType.PREMIUM_LRS, + ) + arm = disk.to_arm_dict() + assert arm["diskSizeGB"] == 128 + assert arm["deleteOption"] == "Delete" + assert arm["managedDisk"]["storageAccountType"] == "Premium_LRS" + + def test_ephemeral_os_disk_defaults_placement(self): + disk = AzureOSDiskConfig(ephemeral_os_disk=True) + + assert disk.ephemeral_placement == "CacheDisk" + + def test_data_disk(self): + dd = AzureDataDisk(lun=0, disk_size_gb=256) + arm = dd.to_arm_dict() + assert arm["lun"] == 0 + assert arm["diskSizeGB"] == 256 + assert arm["createOption"] == "Empty" + assert arm["deleteOption"] == "Delete" + + def test_network_config(self): + nc = AzureNetworkConfig( + subnet_id="/subscriptions/.../subnets/default", + public_ip_enabled=True, + load_balancer_backend_pool_ids=["/subscriptions/.../backendAddressPools/pool-a"], + load_balancer_inbound_nat_pool_ids=["/subscriptions/.../inboundNatPools/nat-a"], + application_gateway_backend_pool_ids=["/subscriptions/.../backendAddressPools/appgw-a"], + ) + arm = nc.to_arm_dict() + assert arm["properties"]["deleteOption"] == "Delete" + assert arm["properties"]["primary"] is True + ip_config = arm["properties"]["ipConfigurations"][0]["properties"] + assert ip_config["subnet"]["id"] == ("/subscriptions/.../subnets/default") + assert ip_config["publicIPAddressConfiguration"]["properties"]["deleteOption"] == "Delete" + assert ip_config["loadBalancerBackendAddressPools"][0]["id"].endswith("/pool-a") + assert ip_config["loadBalancerInboundNatPools"][0]["id"].endswith("/nat-a") + assert ip_config["applicationGatewayBackendAddressPools"][0]["id"].endswith("/appgw-a") + + def test_rejects_aws_shaped_allocation_strategy(self): + with pytest.raises(ValueError, match="Azure templates do not support allocation_strategy"): + AzureTemplate(**_BASE_FIELDS, allocation_strategy="spotPlacementScore") + + def test_priority_from_price_type(self): + from orb.domain.base.value_objects import PriceType + + mapped = AzurePriority.from_price_type(PriceType.SPOT) + assert mapped == AzurePriority.SPOT diff --git a/tests/providers/azure/test_azure_validation_adapter.py b/tests/providers/azure/test_azure_validation_adapter.py new file mode 100644 index 000000000..35106f3f6 --- /dev/null +++ b/tests/providers/azure/test_azure_validation_adapter.py @@ -0,0 +1,103 @@ +"""Focused tests for AzureValidationAdapter.""" + +from unittest.mock import MagicMock + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.infrastructure.adapters.azure_validation_adapter import ( + AzureValidationAdapter, +) + + +def test_validate_template_configuration_uses_azure_rules(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + result = adapter.validate_template_configuration( + { + "provider_api": "VMSS", + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + + assert result["valid"] is True + assert result["errors"] == [] + assert "provider_api" in result["validated_fields"] + + +def test_validate_template_configuration_rejects_unsupported_provider_api(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + result = adapter.validate_template_configuration( + { + "provider_api": "BogusApi", + "template_id": "t1", + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + + assert result["valid"] is False + assert any("Unsupported provider API" in error for error in result["errors"]) + + +def test_validate_provider_api_fallback_accepts_vmss_uniform(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + assert adapter.validate_provider_api("VMSSUniform") is True + + +def test_validate_provider_api_does_not_accept_dead_config_only_api_names(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + assert adapter.validate_provider_api("AzureFleet") is False + + +def test_validate_template_configuration_rejects_spot_percentage_for_uniform(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + result = adapter.validate_template_configuration( + { + "provider_api": "VMSS", + "template_id": "t1", + "orchestration_mode": "Uniform", + "spot_percentage": 70, + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + + assert result["valid"] is False + assert any( + "spot_percentage requires Flexible orchestration mode" in error + for error in result["errors"] + ) + + +def test_validate_template_configuration_accepts_spot_percentage_without_spot_priority(): + adapter = AzureValidationAdapter(config=AzureProviderConfig(), logger=MagicMock()) + + result = adapter.validate_template_configuration( + { + "provider_api": "VMSS", + "template_id": "t1", + "spot_percentage": 70, + "vm_size": "Standard_D4s_v5", + "resource_group": "rg", + "location": "eastus2", + "ssh_public_keys": ["ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC7 test@host"], + "image": {"publisher": "C", "offer": "o", "sku": "s"}, + } + ) + + assert result["valid"] is True + assert result["errors"] == [] diff --git a/tests/providers/azure/test_cyclecloud_handler.py b/tests/providers/azure/test_cyclecloud_handler.py new file mode 100644 index 000000000..687700afd --- /dev/null +++ b/tests/providers/azure/test_cyclecloud_handler.py @@ -0,0 +1,798 @@ +"""Tests for the CycleCloud handler and related template/exception additions.""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.domain.template.azure_template_aggregate import AzureTemplate +from orb.providers.azure.exceptions.azure_exceptions import ( + CycleCloudConnectionError, + CycleCloudNodeError, + TerminationError, +) +from orb.providers.azure.infrastructure.cyclecloud_session import ( + CycleCloudCredentialData, + CycleCloudRequestContext, +) +from orb.providers.azure.infrastructure.cyclecloud_session_builder import ( + CycleCloudSessionBuilder, +) +from orb.providers.azure.infrastructure.handlers.azure_handler import AzureReleaseContext +from orb.providers.azure.infrastructure.handlers.cyclecloud_handler import ( + CycleCloudHandler, + resolve_cc_state, +) +from tests.providers.azure.strategy_test_support import make_azure_template, run_operation + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +_CC_TEMPLATE_FIELDS = { + "template_id": "cc-test", + "vm_size": "Standard_D4s_v5", + "resource_group": "test-rg", + "location": "eastus2", + "provider_api": "CycleCloud", + "cluster_name": "my-cluster", + "node_array": "execute", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + "cyclecloud_verify_ssl": False, +} + + +def _make_template(**overrides): + return make_azure_template(**{**_CC_TEMPLATE_FIELDS, **overrides}) + + +def _make_handler(): + azure_client = MagicMock() + azure_client.get_provider_config.return_value = None + azure_client.get_async_credential = AsyncMock(return_value=None) + logger = MagicMock() + return CycleCloudHandler(azure_client=azure_client, logger=logger) + + +def _make_request(count=2, resource_ids=None, metadata=None): + req = MagicMock() + req.request_id = "req-12345678-1234-1234-1234-123456789012" + req.requested_count = count + req.resource_ids = resource_ids or [] + req.metadata = metadata or {} + return req + + +def _make_cc_request_context(**values): + return CycleCloudRequestContext.from_mapping(values) + + +class _AsyncContextManager: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, exc_type, exc, tb): + return None + + +def _make_async_session_context( + *, + base_url: str = "https://cc.example.com", + credential_path: str | None = None, + verify_ssl: bool = False, + auth_mode: str | None = "bearer", +): + session_context = MagicMock() + session_context.client = object() + session_context.base_url = base_url + session_context.credential_path = credential_path + session_context.verify_ssl = verify_ssl + session_context.auth_mode = auth_mode + return session_context + + +def _wire_async_cyclecloud_calls( + handler: CycleCloudHandler, + *, + responses: list[dict[str, object]], + session_context=None, +): + session_context = session_context or _make_async_session_context() + handler._async_cc_session_scope = MagicMock(return_value=_AsyncContextManager(session_context)) + handler._cc_request_async = AsyncMock(side_effect=responses) + return session_context + + +@pytest.mark.asyncio +async def test_cc_request_async_wraps_invalid_json_with_connection_error(): + handler = _make_handler() + response = MagicMock() + response.request.url = "https://cc.example.com/api/nodes" + response.headers = {"Content-Type": "application/json"} + response.status_code = 200 + response.content = b"{invalid" + response.raise_for_status = MagicMock() + response.json.side_effect = json.JSONDecodeError("bad json", "{invalid", 1) + + client = MagicMock() + client.request = AsyncMock(return_value=response) + + with pytest.raises(CycleCloudConnectionError, match="invalid JSON"): + await handler._cc_request_async( + client, + "GET", + "https://cc.example.com/api/nodes", + include_metadata=True, + ) + + +# --------------------------------------------------------------------------- +# AzureTemplate CycleCloud fields +# --------------------------------------------------------------------------- + + +class TestCycleCloudTemplate: + def test_cyclecloud_template_omitted_verify_ssl_is_unset(self): + fields = {**_CC_TEMPLATE_FIELDS} + del fields["cyclecloud_verify_ssl"] + t = AzureTemplate(**fields) + assert t.cyclecloud_verify_ssl is None + + def test_cyclecloud_template_accepts_explicit_verify_ssl_true(self): + t = _make_template(cyclecloud_verify_ssl=True) + assert t.cyclecloud_verify_ssl is True + + def test_cyclecloud_template_accepts_credential_path(self): + t = _make_template( + cyclecloud_auth_mode=None, + cyclecloud_credential_path="config/cyclecloud-credentials.json", + ) + assert t.cyclecloud_credential_path == "config/cyclecloud-credentials.json" + + def test_cyclecloud_template_accepts_aad_scope_auth_fields(self): + t = _make_template( + cyclecloud_auth_mode="bearer", + cyclecloud_aad_scope="https://example/.default", + ) + assert t.cyclecloud_auth_mode == "bearer" + assert t.cyclecloud_aad_scope == "https://example/.default" + + def test_cyclecloud_template_rejects_inline_bearer_token_field(self): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _make_template(cyclecloud_bearer_token="token-123") + + def test_cyclecloud_template_rejects_inline_basic_auth_fields(self): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _make_template(cyclecloud_username="admin") + + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _make_template(cyclecloud_password="secret") + + def test_cyclecloud_template_requires_cluster_name(self): + with pytest.raises(ValueError, match="cluster_name is required"): + _make_template(cluster_name=None) + + def test_cyclecloud_default_node_array(self): + fields = {**_CC_TEMPLATE_FIELDS} + del fields["node_array"] + t = AzureTemplate(**fields) + assert t.node_array == "execute" + + def test_template_rejects_unknown_extra_fields(self): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _make_template(made_up_setting=True) + + +# --------------------------------------------------------------------------- +# State mapping +# --------------------------------------------------------------------------- + + +class TestStateMapping: + @pytest.mark.parametrize( + "cc_state,expected", + [ + ("Off", "stopped"), + ("Acquiring", "pending"), + ("Preparing", "pending"), + ("Starting", "pending"), + ("Software Configuration", "pending"), + ("Ready", "running"), + ("Deallocating", "shutting-down"), + ("Deallocated", "stopped"), + ("Terminated", "terminated"), + ("Failed", "failed"), + ("SomeUnknownState", "unknown"), + ], + ) + def test_resolve_cc_state(self, cc_state, expected): + assert resolve_cc_state(cc_state) == expected + + +# --------------------------------------------------------------------------- +# CycleCloudHandler - acquire_hosts +# --------------------------------------------------------------------------- + + +class TestCycleCloudHandlerAcquire: + def test_acquire_hosts_returns_submitted_operation_metadata(self): + handler = _make_handler() + template = _make_template() + request = _make_request(count=2) + session_context = _make_async_session_context( + credential_path="config/cc.json", + verify_ssl=False, + auth_mode="bearer", + ) + _wire_async_cyclecloud_calls( + handler, + session_context=session_context, + responses=[ + {"state": "Started"}, + { + "body": { + "operationId": "op-123", + "sets": [ + { + "added": 2, + "nodes": [ + {"name": "node-1", "status": "Acquiring"}, + {"name": "node-2", "status": "Acquiring"}, + ], + } + ], + }, + "headers": {"Location": "https://cc.example.com/operations/op-123"}, + }, + ], + ) + + result = run_operation(handler.acquire_hosts_async(request, template)) + + assert result["success"] is True + assert result["resource_ids"] == ["req-12345678-1234-1234-1234-123456789012"] + assert result["provider_data"]["cluster_name"] == "my-cluster" + assert result["provider_data"]["operation_id"] == "op-123" + assert ( + result["provider_data"]["operation_location"] + == "https://cc.example.com/operations/op-123" + ) + assert result["provider_data"]["added_count"] == 2 + assert result["provider_data"]["submitted_count"] == 2 + assert result["provider_data"]["operation_status"] == "submitted" + request_json = handler._cc_request_async.await_args_list[1].kwargs["json"] + assert request_json["requestId"] == "req-12345678-1234-1234-1234-123456789012" + + @pytest.mark.asyncio + async def test_acquire_hosts_async_returns_submitted_operation_metadata(self): + handler = _make_handler() + template = _make_template() + request = _make_request(count=2) + session_context = MagicMock() + session_context.client = object() + session_context.base_url = "https://cc.example.com" + session_context.credential_path = "config/cc.json" + session_context.verify_ssl = False + session_context.auth_mode = "bearer" + handler._async_cc_session_scope = MagicMock( + return_value=_AsyncContextManager(session_context) + ) + handler._cc_request_async = AsyncMock( + side_effect=[ + {"state": "Started"}, + { + "body": { + "operationId": "op-123", + "sets": [ + {"added": 2, "nodes": [{"name": "node-1", "status": "Acquiring"}]} + ], + }, + "headers": {"Location": "https://cc.example.com/operations/op-123"}, + }, + ] + ) + + result = await handler.acquire_hosts_async(request, template) + + assert result["success"] is True + assert result["provider_data"]["operation_id"] == "op-123" + assert result["provider_data"]["cyclecloud_auth_mode"] == "bearer" + assert result["provider_data"]["submitted_count"] == 2 + + def test_acquire_hosts_missing_cluster_name(self): + """Should raise CycleCloudNodeError if cluster_name is missing.""" + handler = _make_handler() + # Bypass template validation by setting cluster_name to a truthy value + # then overriding via object.__setattr__ + template = _make_template() + object.__setattr__(template, "cluster_name", None) + request = _make_request() + + with pytest.raises(CycleCloudNodeError, match="cluster_name is required"): + run_operation(handler.acquire_hosts_async(request, template)) + + def test_acquire_hosts_missing_url(self): + """Should raise CycleCloudConnectionError if cyclecloud_url is missing.""" + handler = _make_handler() + template = _make_template() + object.__setattr__(template, "cyclecloud_url", None) + request = _make_request() + + with pytest.raises(CycleCloudConnectionError, match="cyclecloud_url is required"): + run_operation(handler.acquire_hosts_async(request, template)) + + +# --------------------------------------------------------------------------- +# CycleCloudHandler - check_hosts_status +# --------------------------------------------------------------------------- + + +class TestCycleCloudHandlerStatus: + def test_check_hosts_status_returns_filtered_results(self): + handler = _make_handler() + _wire_async_cyclecloud_calls( + handler, + responses=[ + { + "nodes": [ + { + "name": "node-1", + "nodeId": "id-1", + "nodeArray": "execute", + "state": "Ready", + "privateIp": "10.0.0.1", + "machineType": "Standard_D4s_v5", + }, + { + "name": "node-2", + "nodeId": "id-2", + "nodeArray": "execute", + "state": "Preparing", + "privateIp": "10.0.0.2", + "machineType": "Standard_D4s_v5", + }, + ] + } + ], + ) + + request = _make_request( + resource_ids=["req-12345678-1234-1234-1234-123456789012"], + metadata={ + "cluster_name": "my-cluster", + "node_array": "execute", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + }, + ) + + results = run_operation(handler.check_hosts_status_async(request)) + + assert len(results) == 2 + assert results[0]["instance_id"] == "node-1" + assert results[0]["name"] == "node-1" + assert results[0]["resource_id"] == "my-cluster" + assert results[0]["status"] == "running" + assert results[0]["private_ip"] == "10.0.0.1" + assert results[1]["status"] == "pending" + + def test_check_hosts_status_no_cc_url(self): + handler = _make_handler() + request = _make_request( + resource_ids=["req-12345678-1234-1234-1234-123456789012"], + metadata={"cluster_name": "my-cluster"}, + ) + with pytest.raises(CycleCloudConnectionError, match="cyclecloud_url is required"): + run_operation(handler.check_hosts_status_async(request)) + + def test_check_hosts_status_requires_cyclecloud_request_identity(self): + handler = _make_handler() + request = _make_request( + resource_ids=[""], + metadata={ + "cluster_name": "my-cluster", + "cyclecloud_url": "https://cc.example.com", + }, + ) + with pytest.raises(CycleCloudConnectionError, match="request identity is required"): + run_operation(handler.check_hosts_status_async(request)) + + def test_check_hosts_status_requires_cluster_name(self): + handler = _make_handler() + request = _make_request( + resource_ids=["req-12345678-1234-1234-1234-123456789012"], + metadata={"cyclecloud_url": "https://cc.example.com"}, + ) + + with pytest.raises(CycleCloudConnectionError, match="cluster_name is required"): + run_operation(handler.check_hosts_status_async(request)) + + def test_check_hosts_status_request_failure_raises(self): + handler = _make_handler() + _wire_async_cyclecloud_calls( + handler, + responses=[ + CycleCloudConnectionError("Cannot connect to CycleCloud at x: boom", url="x") + ], + ) + handler._cc_request_async.side_effect = CycleCloudConnectionError( + "Cannot connect to CycleCloud at https://cc.example.com: boom", + url="https://cc.example.com", + ) + + request = _make_request( + resource_ids=["req-12345678-1234-1234-1234-123456789012"], + metadata={ + "cluster_name": "my-cluster", + "cyclecloud_url": "https://cc.example.com", + }, + ) + + with pytest.raises(CycleCloudConnectionError, match="Cannot connect to CycleCloud"): + run_operation(handler.check_hosts_status_async(request)) + + assert handler._logger.error.call_count == 1 + assert handler._logger.error.call_args.args[0] == ( + "Failed to build CycleCloud session for status check (cluster '%s'): %s" + ) + assert handler._logger.error.call_args.args[1] == "my-cluster" + + @pytest.mark.asyncio + async def test_check_hosts_status_async_returns_filtered_results(self): + handler = _make_handler() + session_context = MagicMock() + session_context.client = object() + session_context.base_url = "https://cc.example.com" + handler._async_cc_session_scope = MagicMock( + return_value=_AsyncContextManager(session_context) + ) + handler._cc_request_async = AsyncMock( + return_value={ + "nodes": [ + { + "name": "node-1", + "nodeId": "id-1", + "nodeArray": "execute", + "state": "Ready", + "machineType": "Standard_D4s_v5", + } + ] + } + ) + request = _make_request( + resource_ids=["req-123"], + metadata={ + "cluster_name": "my-cluster", + "node_array": "execute", + "node_ids": ["node-1"], + "cyclecloud_url": "https://cc.example.com", + }, + ) + + result = await handler.check_hosts_status_async(request) + + assert len(result) == 1 + assert result[0]["status"] == "running" + assert result[0]["instance_id"] == "node-1" + + +# --------------------------------------------------------------------------- +# CycleCloudHandler - release_hosts +# --------------------------------------------------------------------------- + + +class TestCycleCloudHandlerRelease: + def test_release_hosts_returns_submitted_release_metadata(self): + handler = _make_handler() + _wire_async_cyclecloud_calls( + handler, + responses=[ + {"nodes": [{"name": "node-1", "nodeId": "node-1"}]}, + { + "body": {"operationId": "op-release"}, + "headers": {"Location": "https://cc.example.com/operations/op-release"}, + }, + ], + ) + + result = run_operation( + handler.release_hosts_async( + machine_ids=["node-1", "node-2"], + resource_id="my-cluster", + context=AzureReleaseContext( + cyclecloud_request_context=CycleCloudRequestContext( + cyclecloud_url="https://cc.example.com", + cyclecloud_auth_mode="bearer", + cyclecloud_aad_scope="https://cc.example.com/.default", + ) + ), + ) + ) + assert result["provider_data"]["operation_status"] == "submitted" + assert result["provider_data"]["terminate_operation_location"] == ( + "https://cc.example.com/operations/op-release" + ) + + def test_release_hosts_missing_url(self): + handler = _make_handler() + with pytest.raises(TerminationError, match="cyclecloud_url is required"): + run_operation( + handler.release_hosts_async( + machine_ids=["node-1"], + resource_id="my-cluster", + context=AzureReleaseContext(), + ) + ) + + @pytest.mark.asyncio + async def test_release_hosts_async_returns_submitted_release_metadata(self): + handler = _make_handler() + session_context = MagicMock() + session_context.client = object() + session_context.base_url = "https://cc.example.com" + handler._async_cc_session_scope = MagicMock( + return_value=_AsyncContextManager(session_context) + ) + handler._resolve_release_node_targets_async = AsyncMock(return_value={"names": ["node-1"]}) + handler._cc_request_async = AsyncMock( + return_value={ + "body": {"operationId": "op-release"}, + "headers": {"Location": "https://cc.example.com/operations/op-release"}, + } + ) + + result = await handler.release_hosts_async( + machine_ids=["node-1"], + resource_id="my-cluster", + context=AzureReleaseContext( + cyclecloud_request_context=_make_cc_request_context( + cluster_name="my-cluster", + cyclecloud_url="https://cc.example.com", + ) + ), + ) + + assert result is not None + assert result["provider_data"]["operation_status"] == "submitted" + assert ( + result["provider_data"]["terminate_operation_location"] + == "https://cc.example.com/operations/op-release" + ) + + +# --------------------------------------------------------------------------- +# CycleCloudHandler - auth modes +# --------------------------------------------------------------------------- + + +class TestCycleCloudAuthModes: + @pytest.mark.asyncio + async def test_cc_request_async_uses_provider_configured_timeouts(self): + handler = _make_handler() + handler.azure_client.get_provider_config.return_value = AzureProviderConfig( + region="eastus2", + connect_timeout=5, + read_timeout=13, + ) + client = MagicMock() + response = MagicMock() + response.content = b"{}" + response.json.return_value = {} + response.raise_for_status = MagicMock() + client.request = AsyncMock(return_value=response) + + await handler._cc_request_async( + client, + "GET", + "https://cc.example.com/clusters/my-cluster/status", + ) + + client.request.assert_awaited_once_with( + "GET", + "https://cc.example.com/clusters/my-cluster/status", + ) + + @pytest.mark.asyncio + async def test_build_async_cc_session_uses_async_credential(self): + handler = _make_handler() + async_credential = MagicMock() + async_credential.get_token = AsyncMock(return_value=MagicMock(token="tok-123")) + handler.azure_client.get_async_credential = AsyncMock(return_value=async_credential) + + class _CredentialProperty: + def __get__(self, instance, owner): + raise AssertionError("sync credential should not be used") + + with patch.object( + type(handler.azure_client), + "credential", + new=_CredentialProperty(), + create=True, + ): + session_context = await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=False, + request_context=_make_cc_request_context(cyclecloud_auth_mode="bearer"), + ) + + assert session_context.auth_mode == "bearer" + await session_context.client.aclose() + + @pytest.mark.asyncio + async def test_build_async_cc_session_reads_timeout_tuple_once(self): + handler = _make_handler() + handler._get_cc_request_timeout = MagicMock(return_value=(5, 13)) + handler.azure_client.get_async_credential = AsyncMock(return_value=None) + + with patch.object( + CycleCloudSessionBuilder, + "resolve_async_auth", + new=AsyncMock(return_value=({}, None, "none")), + ): + session_context = await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=False, + request_context=_make_cc_request_context(), + ) + + handler._get_cc_request_timeout.assert_called_once_with() + await session_context.client.aclose() + + @pytest.mark.asyncio + async def test_build_async_cc_session_uses_azure_bearer_when_no_basic_auth(self): + handler = _make_handler() + async_credential = MagicMock() + async_credential.get_token = AsyncMock(return_value=MagicMock(token="tok-123")) + handler.azure_client.get_async_credential = AsyncMock(return_value=async_credential) + + session_context = await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=True, + request_context=_make_cc_request_context(cyclecloud_auth_mode="bearer"), + ) + + assert session_context.base_url == "https://cc.example.com" + assert session_context.auth_mode == "bearer" + assert session_context.client.headers["Authorization"] == "Bearer tok-123" + await session_context.client.aclose() + + @pytest.mark.asyncio + async def test_build_async_cc_session_rejects_ssh_auth_mode(self): + handler = _make_handler() + + with pytest.raises( + CycleCloudConnectionError, + match="cyclecloud_auth_mode=ssh is not supported", + ): + await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=True, + request_context=_make_cc_request_context(cyclecloud_auth_mode="ssh"), + ) + + @pytest.mark.asyncio + async def test_build_async_cc_session_propagates_auth_resolution_failures(self): + handler = _make_handler() + + with patch.object( + CycleCloudSessionBuilder, + "resolve_async_auth", + new=AsyncMock( + side_effect=CycleCloudConnectionError( + "cyclecloud_auth_mode=bearer requested but no bearer token could be resolved" + ) + ), + ): + with pytest.raises( + CycleCloudConnectionError, + match="cyclecloud_auth_mode=bearer requested but no bearer token could be resolved", + ): + await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=True, + request_context=_make_cc_request_context(cyclecloud_auth_mode="bearer"), + ) + + @pytest.mark.asyncio + async def test_build_async_cc_session_propagates_client_construction_failures(self): + handler = _make_handler() + + with ( + patch.object( + CycleCloudSessionBuilder, + "resolve_async_auth", + new=AsyncMock(return_value=({}, None, "none")), + ), + patch( + "orb.providers.azure.infrastructure.handlers.cyclecloud_handler.httpx.AsyncClient", + side_effect=RuntimeError("client setup failed"), + ), + ): + with pytest.raises(RuntimeError, match="client setup failed"): + await handler._build_async_cc_session( + cc_url="https://cc.example.com", + verify_ssl=True, + request_context=_make_cc_request_context(cyclecloud_auth_mode="none"), + ) + + @pytest.mark.asyncio + async def test_async_cc_session_scope_builds_session_on_enter(self): + handler = _make_handler() + + with patch.object( + CycleCloudSessionBuilder, + "resolve_async_auth", + new=AsyncMock(return_value=({}, None, "none")), + ): + async with handler._async_cc_session_scope( + cc_url="https://cc.example.com", + verify_ssl=True, + request_context=_make_cc_request_context(cyclecloud_auth_mode="none"), + ) as session_context: + assert session_context.base_url == "https://cc.example.com" + client = session_context.client + assert client.is_closed is False + + assert client.is_closed is True + + @pytest.mark.asyncio + async def test_build_async_cc_session_loads_cyclecloud_config_from_provider(self): + handler = _make_handler() + handler.azure_client.get_provider_config.return_value = AzureProviderConfig( + region="eastus2", + resource_group="orb-test-rg", + cyclecloud={ + "credential_path": "config/cyclecloud-credentials.json", + "url": "https://cc.example.com", + "verify_ssl": False, + }, + ) + with patch.object( + CycleCloudSessionBuilder, + "_load_credential_file", + return_value=CycleCloudCredentialData( + username="cc_admin", + password="changeme", + ), + ): + session_context = await handler._build_async_cc_session( + cc_url=None, + verify_ssl=None, + ) + + assert session_context.base_url == "https://cc.example.com" + assert session_context.verify_ssl is False + assert session_context.auth_mode == "basic" + assert session_context.credential_path == "config/cyclecloud-credentials.json" + await session_context.client.aclose() + + def test_build_settings_does_not_carry_raw_credentials_in_repr(self): + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + {"cyclecloud_credential_path": "/tmp/cyclecloud.json"} + ), + provider_cfg=None, + ) + with patch.object( + CycleCloudSessionBuilder, + "_load_credential_file", + return_value=CycleCloudCredentialData( + username="cc_admin", + password="changeme", + ), + ): + settings = builder.build_settings() + + settings_repr = repr(settings) + assert "cc_admin" not in settings_repr + assert "changeme" not in settings_repr diff --git a/tests/providers/azure/test_cyclecloud_session_builder.py b/tests/providers/azure/test_cyclecloud_session_builder.py new file mode 100644 index 000000000..174302797 --- /dev/null +++ b/tests/providers/azure/test_cyclecloud_session_builder.py @@ -0,0 +1,404 @@ +"""Focused tests for CycleCloud session settings resolution.""" + +import json +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from azure.core.exceptions import ClientAuthenticationError +from azure.identity import CredentialUnavailableError + +from orb.providers.azure.configuration.config import AzureProviderConfig +from orb.providers.azure.exceptions.azure_exceptions import CycleCloudConnectionError +from orb.providers.azure.infrastructure.cyclecloud_session import ( + CycleCloudCredentialData, + CycleCloudRequestContext, +) +from orb.providers.azure.infrastructure.cyclecloud_session_builder import ( + CycleCloudSessionBuilder, +) + + +def _make_builder(*, provider_cfg=None, template=None, request_context=None, credential=None): + async_token_provider = None + if credential is not None: + async_token_provider = MagicMock() + async_token_provider.get_access_token = AsyncMock( + side_effect=lambda scope: credential.get_token(scope).token + ) + async_token_provider.get_auth_error_types.return_value = ( + CredentialUnavailableError, + ClientAuthenticationError, + ) + return CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=None, + template=template, + request_context=request_context or CycleCloudRequestContext(), + provider_cfg=provider_cfg, + async_token_provider=async_token_provider, + ) + + +def test_cyclecloud_request_context_round_trips_metadata(): + context = CycleCloudRequestContext.from_mapping( + { + "cluster_name": "my-cluster", + "node_array": "execute", + "node_ids": ["node-1", "node-2"], + "operation_id": "op-123", + "operation_location": "https://cc.example.com/operations/op-123", + "added_count": "2", + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_credential_path": "config/cc.json", + "cyclecloud_verify_ssl": False, + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + } + ) + + assert context.added_count == 2 + assert context.node_ids == ("node-1", "node-2") + assert context.to_metadata() == { + "cluster_name": "my-cluster", + "node_array": "execute", + "node_ids": ["node-1", "node-2"], + "operation_id": "op-123", + "operation_location": "https://cc.example.com/operations/op-123", + "added_count": 2, + "cyclecloud_url": "https://cc.example.com", + "cyclecloud_credential_path": "config/cc.json", + "cyclecloud_verify_ssl": False, + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + } + + +def test_build_settings_returns_bearer_mode_when_no_token_provider_is_available(): + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping({"cyclecloud_auth_mode": "bearer"}), + provider_cfg=None, + ) + + settings = builder.build_settings() + + assert settings.auth_mode == "bearer" + + +@pytest.mark.asyncio +async def test_resolve_async_auth_skips_expected_auth_failures_before_returning_bearer_token(): + credential = MagicMock() + credential.get_token.side_effect = [ + CredentialUnavailableError("missing"), + ClientAuthenticationError(message="bad token"), + MagicMock(token="tok-123"), + ] + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + { + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://scope-1/.default", + } + ), + provider_cfg=None, + async_token_provider=MagicMock( + get_access_token=AsyncMock(side_effect=lambda scope: credential.get_token(scope).token), + get_auth_error_types=MagicMock( + return_value=(CredentialUnavailableError, ClientAuthenticationError) + ), + ), + ) + + settings = builder.build_settings() + headers, auth, resolved_auth_mode = await builder.resolve_async_auth(settings=settings) + + assert resolved_auth_mode == "bearer" + assert auth is None + assert headers["Authorization"] == "Bearer tok-123" + + +@pytest.mark.asyncio +async def test_resolve_async_auth_propagates_unexpected_token_errors(): + credential = MagicMock() + credential.get_token.side_effect = RuntimeError("boom") + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping({"cyclecloud_auth_mode": "bearer"}), + provider_cfg=None, + async_token_provider=MagicMock( + get_access_token=AsyncMock(side_effect=lambda scope: credential.get_token(scope).token), + get_auth_error_types=MagicMock( + return_value=(CredentialUnavailableError, ClientAuthenticationError) + ), + ), + ) + + with pytest.raises(RuntimeError, match="boom"): + settings = builder.build_settings() + await builder.resolve_async_auth(settings=settings) + + +@pytest.mark.asyncio +async def test_resolve_async_auth_uses_async_token_provider(): + async_token_provider = MagicMock() + async_token_provider.get_access_token = AsyncMock( + side_effect=[ + CredentialUnavailableError("missing"), + "tok-async-123", + ] + ) + async_token_provider.get_auth_error_types.return_value = ( + CredentialUnavailableError, + ClientAuthenticationError, + ) + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + { + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://scope-1/.default", + } + ), + provider_cfg=None, + async_token_provider=async_token_provider, + ) + + settings = builder.build_settings() + + headers, auth, resolved_auth_mode = await builder.resolve_async_auth(settings=settings) + + assert resolved_auth_mode == "bearer" + assert auth is None + assert headers["Authorization"] == "Bearer tok-async-123" + + +@pytest.mark.asyncio +async def test_resolve_async_auth_rejects_ssh_mode(): + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping({"cyclecloud_auth_mode": "ssh"}), + provider_cfg=None, + ) + + with pytest.raises(CycleCloudConnectionError, match="not supported"): + settings = builder.build_settings() + await builder.resolve_async_auth(settings=settings) + + +@pytest.mark.asyncio +async def test_resolve_async_auth_errors_when_bearer_requested_but_unavailable(): + async_token_provider = MagicMock() + async_token_provider.get_access_token = AsyncMock( + side_effect=CredentialUnavailableError("missing") + ) + async_token_provider.get_auth_error_types.return_value = ( + CredentialUnavailableError, + ClientAuthenticationError, + ) + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping({"cyclecloud_auth_mode": "bearer"}), + provider_cfg=None, + async_token_provider=async_token_provider, + ) + + with pytest.raises(CycleCloudConnectionError, match="no bearer token could be resolved"): + settings = builder.build_settings() + await builder.resolve_async_auth(settings=settings) + + +@pytest.mark.asyncio +async def test_resolve_async_auth_errors_when_no_auth_method_resolves(): + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext(), + provider_cfg=None, + ) + + with pytest.raises(CycleCloudConnectionError, match="No CycleCloud auth method resolved"): + settings = builder.build_settings() + await builder.resolve_async_auth(settings=settings) + + +def test_build_settings_loads_cyclecloud_config_from_provider(): + provider_cfg = AzureProviderConfig( + region="eastus2", + resource_group="orb-test-rg", + cyclecloud={ + "credential_path": "config/cyclecloud-credentials.json", + "url": "https://cc.example.com", + "verify_ssl": False, + }, + ) + builder = _make_builder(provider_cfg=provider_cfg) + builder._load_credential_file = MagicMock( # type: ignore[method-assign] + return_value=CycleCloudCredentialData( + username="cc_admin", + password="changeme", + ) + ) + + settings = builder.build_settings() + + assert settings.base_url == "https://cc.example.com" + assert settings.verify_ssl is False + assert settings.auth_mode is None + assert settings.credential_path == "config/cyclecloud-credentials.json" + + +def test_cyclecloud_credential_data_repr_masks_secret_fields(): + credential_data = CycleCloudCredentialData( + url="https://cc.example.com", + auth_mode="bearer", + username="cc_admin", + password="changeme", + bearer_token="tok-123", + aad_scope="https://cc.example.com/.default", + ) + + credential_repr = repr(credential_data) + + assert "cc_admin" not in credential_repr + assert "changeme" not in credential_repr + assert "tok-123" not in credential_repr + assert "https://cc.example.com" in credential_repr + assert "bearer" in credential_repr + + +@pytest.mark.asyncio +async def test_resolve_async_auth_loads_credentials_from_file(tmp_path: Path): + credential_file = tmp_path / "cyclecloud-credentials.json" + credential_file.write_text( + json.dumps( + { + "username": "file-admin", + "password": "file-secret", + "auth_mode": "basic", + } + ), + encoding="utf-8", + ) + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=False, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + {"cyclecloud_credential_path": str(credential_file)} + ), + provider_cfg=None, + ) + + settings = builder.build_settings() + + assert settings.base_url == "https://cc.example.com" + assert settings.verify_ssl is False + assert settings.credential_path == str(credential_file) + + headers, auth, resolved_auth_mode = await builder.resolve_async_auth(settings=settings) + assert headers == {} + assert resolved_auth_mode == "basic" + assert isinstance(auth, httpx.BasicAuth) + + +def test_build_settings_parses_verify_ssl_string_from_request_context(): + builder = CycleCloudSessionBuilder( + cc_url="https://cc.example.com", + verify_ssl=None, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + { + "cyclecloud_verify_ssl": "false", + "cyclecloud_auth_mode": "bearer", + "cyclecloud_aad_scope": "https://cc.example.com/.default", + } + ), + provider_cfg=None, + ) + settings = builder.build_settings() + + assert settings.verify_ssl is False + + +def test_build_settings_takes_verify_ssl_from_credential_file(tmp_path: Path): + credential_file = tmp_path / "cyclecloud-credentials.json" + credential_file.write_text( + json.dumps( + { + "url": "https://cc.example.com", + "username": "file-admin", + "password": "file-secret", + "verify_ssl": "false", + } + ), + encoding="utf-8", + ) + builder = CycleCloudSessionBuilder( + cc_url=None, + verify_ssl=None, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + {"cyclecloud_credential_path": str(credential_file)} + ), + provider_cfg=None, + ) + + settings = builder.build_settings() + + assert settings.base_url == "https://cc.example.com" + assert settings.verify_ssl is False + assert settings.credential_path == str(credential_file) + + +def test_resolve_cascaded_value_skips_blank_values_and_uses_default(): + builder = CycleCloudSessionBuilder( + cc_url=None, + verify_ssl=None, + template=None, + request_context=CycleCloudRequestContext(), + provider_cfg=None, + ) + + resolved = builder._resolve_cascaded_value(None, "", False, default=True) + + assert resolved is False + + +def test_build_settings_resolves_url_from_request_context_before_provider_config(): + provider_cfg = AzureProviderConfig( + region="eastus2", + resource_group="orb-test-rg", + cyclecloud={ + "url": "https://provider.example.com", + "verify_ssl": True, + }, + ) + builder = CycleCloudSessionBuilder( + cc_url=None, + verify_ssl=None, + template=None, + request_context=CycleCloudRequestContext.from_mapping( + {"cyclecloud_url": "https://request.example.com"} + ), + provider_cfg=provider_cfg, + ) + + settings = builder.build_settings() + + assert settings.base_url == "https://request.example.com" diff --git a/tests/providers/azure/test_import_guards.py b/tests/providers/azure/test_import_guards.py new file mode 100644 index 000000000..8d6004903 --- /dev/null +++ b/tests/providers/azure/test_import_guards.py @@ -0,0 +1,98 @@ +"""Import guards for Azure optional runtime packages.""" + +from __future__ import annotations + +import sys +from contextlib import contextmanager + +_MISSING_AZURE_MODULES = { + "azure": None, + "azure.core": None, + "azure.core.exceptions": None, + "azure.identity": None, + "azure.mgmt": None, + "azure.mgmt.compute": None, + "azure.mgmt.network": None, + "azure.mgmt.resource": None, + "azure.mgmt.resource.subscriptions": None, +} + + +@contextmanager +def _isolated_azure_provider_import(): + """Re-import Azure provider modules with Azure SDK packages hidden.""" + saved_orb_modules = { + k: v for k, v in sys.modules.items() if k.startswith("orb.providers.azure") + } + saved_azure_modules = {k: sys.modules[k] for k in _MISSING_AZURE_MODULES if k in sys.modules} + + for key in list(sys.modules): + if key.startswith("orb.providers.azure"): + del sys.modules[key] + + for key, value in _MISSING_AZURE_MODULES.items(): + # None is Python's import-blocking sentinel, omitted from typeshed's module map. + sys.modules[key] = value # type: ignore[assignment] + + try: + yield + finally: + for key in list(sys.modules): + if key.startswith("orb.providers.azure"): + del sys.modules[key] + for key in _MISSING_AZURE_MODULES: + if key in saved_azure_modules: + sys.modules[key] = saved_azure_modules[key] + else: + sys.modules.pop(key, None) + sys.modules.update(saved_orb_modules) + for module_name, module in sorted( + saved_orb_modules.items(), key=lambda item: item[0].count(".") + ): + parent_name, attribute_name = module_name.rsplit(".", maxsplit=1) + parent_module = sys.modules.get(parent_name) + if parent_module is not None: + setattr(parent_module, attribute_name, module) + + +def test_azure_package_import_does_not_require_azure_sdk() -> None: + with _isolated_azure_provider_import(): + import orb.providers.azure as azure_provider + + assert azure_provider is not None + + +def test_azure_registration_import_does_not_require_azure_sdk() -> None: + with _isolated_azure_provider_import(): + from orb.providers.azure.registration import register_azure_provider + + assert callable(register_azure_provider) + + +def test_register_azure_provider_does_not_require_azure_sdk() -> None: + with _isolated_azure_provider_import(): + from orb.providers.azure.registration import register_azure_provider + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + registry.clear_registrations() + + register_azure_provider(registry=registry) + + assert registry.is_provider_registered("azure") is True + + +def test_create_azure_strategy_does_not_require_azure_sdk_until_runtime() -> None: + with _isolated_azure_provider_import(): + from orb.providers.azure.registration import create_azure_strategy + + strategy = create_azure_strategy( + { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "orb-rg", + "region": "eastus2", + }, + provider_instance_name="azure-default", + ) + + assert strategy.is_initialized is False diff --git a/tests/providers/azure/test_provider_registration_contracts.py b/tests/providers/azure/test_provider_registration_contracts.py new file mode 100644 index 000000000..c97df9bc5 --- /dev/null +++ b/tests/providers/azure/test_provider_registration_contracts.py @@ -0,0 +1,172 @@ +"""Azure provider registration and default-loading contracts.""" + +from unittest.mock import MagicMock, patch + +from orb.config.schemas.provider_strategy_schema import ProviderInstanceConfig + + +def _raw_config_with_azure() -> dict: + return { + "provider": { + "providers": [ + { + "name": "azure-default", + "type": "azure", + "enabled": True, + "config": { + "subscription_id": "12345678-1234-1234-1234-123456789012", + "resource_group": "orb-test-rg", + "location": "eastus2", + }, + } + ], + "active_provider": "azure-default", + } + } + + +def test_load_strategy_defaults_includes_azure_defaults_without_provider_bootstrap(): + """Static defaults loading must include Azure without bootstrapping providers.""" + from orb.config.loader import ConfigurationLoader + + with ( + patch("orb.providers.registration.register_all_provider_types") as register_all, + patch("orb.providers.registry.get_provider_registry") as get_provider_registry, + ): + defaults = ConfigurationLoader._load_strategy_defaults() + + register_all.assert_not_called() + get_provider_registry.assert_not_called() + assert "azure" in defaults["provider"]["provider_defaults"] + + +def test_load_strategy_defaults_includes_azure_handler_capabilities_without_provider_bootstrap(): + """Config-based template validation must see Azure provider APIs from defaults.""" + from orb.config.loader import ConfigurationLoader + from orb.config.schemas.provider_strategy_schema import ProviderConfig, ProviderInstanceConfig + + with ( + patch("orb.providers.registration.register_all_provider_types") as register_all, + patch("orb.providers.registry.get_provider_registry") as get_provider_registry, + ): + defaults = ConfigurationLoader._load_strategy_defaults() + + register_all.assert_not_called() + get_provider_registry.assert_not_called() + provider_config = ProviderConfig.model_validate(defaults["provider"]) + provider_instance = ProviderInstanceConfig( + name="azure-default", + type="azure", + enabled=True, + config={}, + ) + + handlers = provider_instance.get_effective_handlers(provider_config.provider_defaults["azure"]) + + assert set(handlers) == {"VMSS", "VMSSUniform", "SingleVM", "CycleCloud"} + assert handlers["VMSS"].model_dump()["supports_spot"] is True + assert handlers["CycleCloud"].model_dump()["supports_spot"] is False + + +def test_register_all_provider_types_includes_azure(): + """Canonical provider bootstrap must register Azure.""" + from orb.providers.registration import register_all_provider_types + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + registry.clear_registrations() + + register_all_provider_types() + + assert registry.is_provider_registered("azure") is True + + +def test_register_all_provider_types_registers_azure_auth_strategy(): + """Azure auth strategy must be reachable through AuthRegistry lookup.""" + from orb.config.schemas.server_schema import AuthConfig + from orb.infrastructure.auth.registry import get_auth_registry + from orb.providers.azure.auth.azure_auth_strategy import AzureAuthStrategy + from orb.providers.registration import register_all_provider_types + + registry = get_auth_registry() + if registry.is_registered("azure"): + registry.unregister_type("azure") + + register_all_provider_types() + + assert registry.is_registered("azure") is True + strategy = registry.get_strategy("azure", AuthConfig(enabled=True, strategy="azure")) + assert isinstance(strategy, AzureAuthStrategy) + assert strategy.is_enabled() is True + + +def test_provider_config_builder_accepts_azure_provider_instance_config(): + """Azure config creation must accept the canonical ProviderInstanceConfig input.""" + from orb.providers.config_builder import ProviderConfigBuilder + from orb.providers.registration import register_all_provider_types + from orb.providers.registry import get_provider_registry + + registry = get_provider_registry() + registry.clear_registrations() + register_all_provider_types() + + logger = MagicMock() + builder = ProviderConfigBuilder(logger, registry) + provider_instance = ProviderInstanceConfig( # type: ignore[call-arg] + name="azure-default", + type="azure", + enabled=True, + config={ + "subscription_id": "11111111-1111-1111-1111-111111111111", + "client_id": "test-client", + "region": "uksouth", + }, + ) + + azure_config = builder.build_config(provider_instance) + + assert azure_config.subscription_id == "11111111-1111-1111-1111-111111111111" + assert azure_config.region == "uksouth" + + +def test_get_typed_azure_provider_config_via_registry(): + """get_typed(AzureProviderConfig) resolves through ProviderSettingsRegistry.""" + from orb.config.managers.type_converter import ConfigTypeConverter + from orb.providers.azure.configuration.config import AzureProviderConfig + from orb.providers.azure.registration import register_azure_provider_settings + + register_azure_provider_settings() + + converter = ConfigTypeConverter(_raw_config_with_azure()) + result = converter.get_typed(AzureProviderConfig) + + assert isinstance(result, AzureProviderConfig) + assert result.subscription_id == "12345678-1234-1234-1234-123456789012" + assert result.resource_group == "orb-test-rg" + assert result.location == "eastus2" + + +def test_ensure_provider_instance_registered_from_config_supports_azure(): + """Registry auto-registration must work for Azure instances.""" + from orb.providers.registry.provider_registry import ProviderRegistry + + registry = ProviderRegistry() + registry.clear_registrations() + provider_instance = ProviderInstanceConfig( # type: ignore[call-arg] + name="azure-default", + type="azure", + enabled=True, + config={ + "subscription_id": "test-subscription", + "tenant_id": "test-tenant", + "client_id": "test-client", + "client_secret_path": "/tmp/test-secret", + "region": "uksouth", + }, + ) + + result = registry.ensure_provider_instance_registered_from_config(provider_instance) + + assert result is True + assert registry.is_provider_registered("azure") is True + assert registry.is_provider_instance_registered("azure-default") is True diff --git a/tests/providers/azure/test_resource_metadata_service.py b/tests/providers/azure/test_resource_metadata_service.py new file mode 100644 index 000000000..79834f5db --- /dev/null +++ b/tests/providers/azure/test_resource_metadata_service.py @@ -0,0 +1,155 @@ +"""Focused tests for Azure resource metadata enrichment.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.services.resource_metadata_service import ( + AzureResourceMetadataService, +) + + +@pytest.mark.asyncio +async def test_augment_vmss_capacity_metadata_async_collects_multiple_vmss_concurrently(): + service = AzureResourceMetadataService( + default_resource_group="test-rg", + logger=MagicMock(), + ) + resource_manager = MagicMock() + + async def get_capacity(resource_group: str, vmss_name: str) -> dict[str, object]: + _ = resource_group + if vmss_name == "vmss-a": + await asyncio.sleep(0.02) + return { + "capacity": 2, + "provisioned_instance_count": 1, + "provisioning_state": "Updating", + } + await asyncio.sleep(0.0) + return { + "capacity": 3, + "provisioned_instance_count": 3, + "provisioning_state": "Succeeded", + } + + resource_manager.get_vmss_capacity_async = AsyncMock(side_effect=get_capacity) + metadata: dict[str, object] = {} + + await service.augment_vmss_capacity_metadata_async( + metadata, + ["vmss-a", "vmss-b"], + resource_manager=resource_manager, + ) + + assert metadata["fleet_capacity_fulfilment"] == { + "target_capacity_units": 5, + "fulfilled_capacity_units": 4, + "provisioned_instance_count": 4, + "state": "multiple", + } + assert metadata["fleet_capacity_fulfilment_by_resource"] == { + "vmss-a": { + "target_capacity_units": 2, + "fulfilled_capacity_units": 1, + "provisioned_instance_count": 1, + "state": "Updating", + }, + "vmss-b": { + "target_capacity_units": 3, + "fulfilled_capacity_units": 3, + "provisioned_instance_count": 3, + "state": "Succeeded", + }, + } + + +@pytest.mark.asyncio +async def test_augment_vmss_capacity_metadata_async_skips_failed_snapshot_and_preserves_order(): + logger = MagicMock() + service = AzureResourceMetadataService( + default_resource_group="test-rg", + logger=logger, + ) + resource_manager = MagicMock() + + async def get_capacity(resource_group: str, vmss_name: str) -> dict[str, object]: + _ = resource_group + if vmss_name == "vmss-a": + raise RuntimeError("boom") + return { + "capacity": 1, + "provisioned_instance_count": 1, + "provisioning_state": "Succeeded", + } + + resource_manager.get_vmss_capacity_async = AsyncMock(side_effect=get_capacity) + metadata: dict[str, object] = {} + + await service.augment_vmss_capacity_metadata_async( + metadata, + ["vmss-a", "vmss-b"], + resource_manager=resource_manager, + ) + + assert metadata["fleet_capacity_fulfilment"] == { + "target_capacity_units": 1, + "fulfilled_capacity_units": 1, + "provisioned_instance_count": 1, + "state": "Succeeded", + } + assert "fleet_capacity_fulfilment_by_resource" not in metadata + logger.warning.assert_called_once() + + +def test_attach_provider_fulfilment_uses_vmss_capacity_metadata(): + service = AzureResourceMetadataService( + default_resource_group="test-rg", + logger=MagicMock(), + ) + metadata = { + "fleet_capacity_fulfilment": { + "target_capacity_units": 3, + "fulfilled_capacity_units": 2, + "provisioned_instance_count": 2, + "state": "Updating", + } + } + + service.attach_provider_fulfilment( + metadata, + instances=[], + target_units=None, + ) + + fulfilment = metadata["provider_fulfilment"] + assert fulfilment.state == "in_progress" + assert fulfilment.target_units == 3 + assert fulfilment.fulfilled_units == 2 + + +def test_attach_provider_fulfilment_reports_failed_single_vm_deployment(): + service = AzureResourceMetadataService( + default_resource_group="test-rg", + logger=MagicMock(), + ) + metadata = { + "fleet_errors": [ + { + "error_code": "OperationNotAllowed", + "error_message": "quota exceeded", + } + ] + } + + service.attach_provider_fulfilment( + metadata, + instances=[], + target_units=1, + ) + + fulfilment = metadata["provider_fulfilment"] + assert fulfilment.state == "failed" + assert fulfilment.target_units == 1 + assert fulfilment.fulfilled_units == 0 diff --git a/tests/providers/azure/test_single_vm_handler.py b/tests/providers/azure/test_single_vm_handler.py new file mode 100644 index 000000000..3f14c8f76 --- /dev/null +++ b/tests/providers/azure/test_single_vm_handler.py @@ -0,0 +1,767 @@ +"""Focused tests for SingleVM handler behavior.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError + +from orb.providers.azure.exceptions.azure_exceptions import LaunchError, TerminationError +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + RAISE_ON_STATUS_ERROR_METADATA_KEY, + AzureReleaseContext, +) +from orb.providers.azure.infrastructure.handlers.single_vm_handler import ( + SingleVMHandler, + _format_azure_error_message, +) +from orb.providers.azure.infrastructure.sdk_shapes import AzureVmRuntimeStatusProtocol +from orb.providers.azure.infrastructure.services.azure_network_identity_resolver import ( + AzureNetworkIdentity, +) +from tests.providers.azure.strategy_test_support import ( + AsyncPager, + make_azure_template, + make_single_vm_azure_client, + run_operation, +) + + +def test_format_azure_error_message_includes_nested_details(): + message = _format_azure_error_message( + { + "message": "Deployment validation failed", + "details": [ + { + "code": "InvalidParameter", + "message": "The supplied VM size is not available in this location.", + } + ], + } + ) + + assert message == ("InvalidParameter: The supplied VM size is not available in this location.") + + +def test_format_azure_error_message_falls_back_to_primary_message(): + assert ( + _format_azure_error_message({"error_message": "template is invalid"}) + == "template is invalid" + ) + + +def _make_request(*, count: int = 1, request_id: str = "req-1", metadata=None): + request = MagicMock() + request.requested_count = count + request.request_id = request_id + request.metadata = metadata or {} + return request + + +def _deleted_vm_names(azure_client: MagicMock) -> list[str]: + return [ + str(call.kwargs["vm_name"]) + for call in azure_client.compute_client.virtual_machines.begin_delete.call_args_list + ] + + +def _make_template(**overrides): + return make_azure_template( + template_id="azure-singlevm-test", + provider_api="SingleVM", + **overrides, + ) + + +def _make_azure_client() -> MagicMock: + return make_single_vm_azure_client() + + +def test_status_result_omits_absent_vm_id(): + vm = MagicMock(spec=AzureVmRuntimeStatusProtocol) + vm.name = "vm-1" + vm.vm_id = None + vm.location = "eastus2" + vm.zones = None + vm.hardware_profile = None + vm.tags = None + network_identity: AzureNetworkIdentity = { + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + + result = SingleVMHandler._build_status_result( + vm=vm, + resource_group="rg-1", + status="running", + network_identity=network_identity, + ) + + provider_data = result.get("provider_data") + assert provider_data is not None + assert provider_data.get("cloud_host_id") == "vm-1" + assert "vm_id" not in provider_data + + +def test_acquire_hosts_submits_one_batched_deployment_and_returns_submitted_status(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + azure_client.resource_client.resources.begin_create_or_update.return_value = MagicMock() + + result = run_operation( + handler.acquire_hosts_async(_make_request(count=2, request_id="req-2"), _make_template()) + ) + + assert result["success"] is True + assert len(result["resource_ids"]) == 2 + assert result["provider_data"]["operation_status"] == "submitted" + deployment_call = azure_client.resource_client.resources.begin_create_or_update.call_args.kwargs + deployment_template = deployment_call["parameters"]["properties"]["template"] + resource_types = [resource["type"] for resource in deployment_template["resources"]] + assert resource_types.count("Microsoft.Network/networkInterfaces") == 2 + assert resource_types.count("Microsoft.Compute/virtualMachines") == 2 + assert result["provider_data"]["deployment_name"] == deployment_call["resource_name"] + assert len(result["provider_data"]["submitted_vms"]) == 2 + + +def test_acquire_hosts_creates_public_ips_when_enabled(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + azure_client.resource_client.resources.begin_create_or_update.return_value = MagicMock() + + template = _make_template( + network_config={ + "subnet_id": "/subscriptions/.../subnets/default", + "public_ip_enabled": True, + } + ) + + result = run_operation( + handler.acquire_hosts_async(_make_request(count=2, request_id="req-pip"), template) + ) + + assert result["success"] is True + deployment_template = ( + azure_client.resource_client.resources.begin_create_or_update.call_args.kwargs[ + "parameters" + ]["properties"]["template"] + ) + public_ip_resources = [ + resource + for resource in deployment_template["resources"] + if resource["type"] == "Microsoft.Network/publicIPAddresses" + ] + nic_resources = [ + resource + for resource in deployment_template["resources"] + if resource["type"] == "Microsoft.Network/networkInterfaces" + ] + assert len(public_ip_resources) == 2 + assert len(nic_resources) == 2 + public_ip_ref = nic_resources[0]["properties"]["ipConfigurations"][0]["properties"][ + "publicIPAddress" + ] + assert public_ip_resources[0]["name"].startswith("pip-vm-") + assert "Microsoft.Network/publicIPAddresses" in public_ip_ref["id"] + assert public_ip_ref["deleteOption"] == "Delete" + + +@pytest.mark.asyncio +async def test_acquire_hosts_async_submits_one_batched_deployment_and_returns_submitted_status(): + azure_client = _make_azure_client() + azure_client.get_async_compute_client = AsyncMock(return_value=MagicMock()) + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + handler.azure_deployment_service.submit_template_deployment_async = AsyncMock( + return_value="dep-async" + ) + + result = await handler.acquire_hosts_async( + _make_request(count=2, request_id="req-async"), + _make_template(), + ) + + assert result["success"] is True + assert result["provider_data"]["deployment_name"] == "dep-async" + assert result["provider_data"]["submitted_count"] == 2 + + +def test_acquire_hosts_falls_back_to_alternate_vm_size_for_the_whole_batch(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + first_failure = Exception("primary size unavailable") + first_failure.error_code = "AllocationFailed" + azure_client.resource_client.resources.begin_create_or_update.side_effect = [ + first_failure, + MagicMock(), + ] + + template = _make_template(vm_sizes=["Standard_D8s_v5"]) + + result = run_operation( + handler.acquire_hosts_async(_make_request(count=2, request_id="req-4"), template) + ) + + assert result["success"] is True + assert result["provider_data"]["submitted_count"] == 2 + assert all( + submitted_vm["selected_vm_size"] == "Standard_D8s_v5" + for submitted_vm in result["provider_data"]["submitted_vms"] + ) + + +def test_acquire_hosts_stops_after_non_capacity_error(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + failure = Exception("template is invalid") + failure.error_code = "InvalidParameter" + azure_client.resource_client.resources.begin_create_or_update.side_effect = failure + + template = _make_template(vm_sizes=["Standard_D8s_v5"]) + + with pytest.raises(LaunchError, match="template is invalid"): + run_operation( + handler.acquire_hosts_async(_make_request(count=2, request_id="req-invalid"), template) + ) + + assert azure_client.resource_client.resources.begin_create_or_update.call_count == 1 + + +def test_acquire_hosts_classifies_quota_runtime_error(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + handler.azure_deployment_service.submit_template_deployment_async = AsyncMock( + side_effect=RuntimeError("quota exceeded") + ) + + with pytest.raises(LaunchError) as exc_info: + run_operation( + handler.acquire_hosts_async( + _make_request(count=1, request_id="req-quota"), _make_template() + ) + ) + + assert exc_info.value.error_code == "QuotaExceeded" + + +def test_acquire_hosts_does_not_misclassify_resource_name_containing_quota(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + fake_error = MagicMock() + fake_error.code = "InvalidParameter" + fake_error.message = "subnet 'team-quota-subnet-1' does not exist" + + response = MagicMock() + response.status_code = 400 + + exc = HttpResponseError(response=response) + exc.error = fake_error + handler.azure_deployment_service.submit_template_deployment_async = AsyncMock(side_effect=exc) + + with pytest.raises(LaunchError) as exc_info: + run_operation( + handler.acquire_hosts_async( + _make_request(count=1, request_id="req-quota-in-name"), + _make_template(), + ) + ) + + assert exc_info.value.error_code == "InvalidParameter" + + +def test_acquire_hosts_requires_subnet_id(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + template = _make_template(network_config=None, subnet_ids=[]) + + with pytest.raises(LaunchError, match="No subnet specified"): + run_operation(handler.acquire_hosts_async(_make_request(), template)) + + +def test_acquire_hosts_resolves_ssh_key_name_without_mutating_original_template(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + azure_client.resource_client.resources.begin_create_or_update.return_value = MagicMock() + + template = _make_template( + ssh_public_keys=[], + ssh_key_name="orb-key", + ) + + with patch( + "orb.providers.azure.infrastructure.services.ssh_key_resolver.resolve_ssh_keys_async", + new=AsyncMock(return_value=["ssh-rsa resolved test@host"]), + ) as mock_resolve: + result = run_operation( + handler.acquire_hosts_async(_make_request(request_id="req-ssh"), template) + ) + + assert result["success"] is True + assert template.ssh_public_keys == [] + mock_resolve.assert_called_once_with( + ssh_key_name="orb-key", + ssh_public_keys=[], + resource_group="test-rg", + compute_client=mock_resolve.call_args.kwargs["compute_client"], + ) + + +def test_status_populates_network_identity(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vm-1" + ) + nic_ref.properties.primary = True + + vm = MagicMock() + vm.name = "vm-1" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + vm.network_profile.network_interfaces = [nic_ref] + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": "10.0.0.4", + "public_ip": "52.1.2.3", + "subnet_id": "/subscriptions/sub/.../subnets/default", + "vnet_id": "/subscriptions/sub/.../virtualNetworks/test-vnet", + "nic_id": nic_ref.id, + "nic_name": "nic-vm-1", + } + ) + azure_client.compute_client.virtual_machines.get.return_value = vm + + request = MagicMock() + request.resource_ids = ["vm-1"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert result[0]["private_ip"] == "10.0.0.4" + assert result[0]["public_ip"] == "52.1.2.3" + assert result[0]["subnet_id"].endswith("/subnets/default") + assert result[0]["vpc_id"].endswith("/virtualNetworks/test-vnet") + + +@pytest.mark.asyncio +async def test_check_hosts_status_async_populates_network_identity(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm = MagicMock() + vm.name = "vm-1" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + + async_compute = MagicMock() + async_compute.virtual_machines.get = AsyncMock(return_value=vm) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._resolve_vm_names_async = AsyncMock(return_value=["vm-1"]) + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": "10.0.0.4", + "public_ip": "52.1.2.3", + "subnet_id": "/subscriptions/sub/.../subnets/default", + "vnet_id": "/subscriptions/sub/.../virtualNetworks/test-vnet", + "nic_id": "nic-id", + "nic_name": "nic-name", + } + ) + + request = MagicMock() + request.resource_ids = ["vm-1"] + request.metadata = {"resource_group": "test-rg"} + + result = await handler.check_hosts_status_async(request) + + assert result[0]["instance_id"] == "vm-1" + assert result[0]["public_ip"] == "52.1.2.3" + + +def test_status_still_returns_instance_when_network_identity_resolution_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + side_effect=AttributeError("missing network property") + ) + + vm = MagicMock() + vm.name = "vm-single" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + + azure_client.compute_client.virtual_machines.get.return_value = vm + + request = MagicMock() + request.resource_ids = ["vm-single"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert len(result) == 1 + assert result[0]["instance_id"] == "vm-single" + assert result[0]["private_ip"] is None + assert result[0]["provider_data"]["nic_id"] is None + logger.warning.assert_called() + + +@pytest.mark.asyncio +async def test_status_best_effort_returns_partial_results_when_one_vm_lookup_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm = MagicMock() + vm.name = "vm-1" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + + async_compute = MagicMock() + async_compute.virtual_machines.get = AsyncMock(side_effect=[vm, RuntimeError("boom")]) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._resolve_vm_names_async = AsyncMock(return_value=["vm-1", "vm-2"]) + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + ) + + request = MagicMock() + request.resource_ids = ["vm-1", "vm-2"] + request.metadata = {"resource_group": "test-rg"} + + result = await handler.check_hosts_status_async(request) + + assert [entry["instance_id"] for entry in result] == ["vm-1"] + logger.error.assert_called_once() + + +@pytest.mark.asyncio +async def test_status_strict_mode_raises_when_one_vm_lookup_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm = MagicMock() + vm.name = "vm-1" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + + async_compute = MagicMock() + async_compute.virtual_machines.get = AsyncMock(side_effect=[vm, RuntimeError("boom")]) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._resolve_vm_names_async = AsyncMock(return_value=["vm-1", "vm-2"]) + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": None, + "public_ip": None, + "subnet_id": None, + "vnet_id": None, + "nic_id": None, + "nic_name": None, + } + ) + + request = MagicMock() + request.resource_ids = ["vm-1", "vm-2"] + request.metadata = { + "resource_group": "test-rg", + RAISE_ON_STATUS_ERROR_METADATA_KEY: True, + } + + with pytest.raises(RuntimeError, match="Failed to get status for VM 'vm-2'"): + await handler.check_hosts_status_async(request) + + +@pytest.mark.asyncio +async def test_status_best_effort_raises_when_all_vm_lookups_fail(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + async_compute = MagicMock() + async_compute.virtual_machines.get = AsyncMock(side_effect=RuntimeError("boom")) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._resolve_vm_names_async = AsyncMock(return_value=["vm-1"]) + + request = MagicMock() + request.resource_ids = ["vm-1"] + request.metadata = {"resource_group": "test-rg"} + + with pytest.raises(RuntimeError, match="Failed to get status for VM 'vm-1'"): + await handler.check_hosts_status_async(request) + + +@pytest.mark.asyncio +async def test_release_hosts_async_submits_deletions_for_resolved_vm_names(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machines.begin_delete = AsyncMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._resolve_vm_names_async = AsyncMock(return_value=["vm-1"]) + + result = await handler.release_hosts_async( + machine_ids=["vm-1"], + resource_id="ignored", + context=AzureReleaseContext(resource_group="test-rg"), + ) + + assert result is not None + assert result["provider_data"]["operation_status"] == "submitted" + async_compute.virtual_machines.begin_delete.assert_awaited_once() + + +def test_status_uses_direct_vm_name_lookup_without_listing_resource_group(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm = MagicMock() + vm.name = "vm-1" + vm.vm_id = "vm-guid-1" + vm.instance_view.statuses = [] + vm.hardware_profile.vm_size = "Standard_D4s_v5" + vm.location = "eastus2" + vm.zones = ["1"] + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": "10.0.0.4", + "public_ip": None, + "subnet_id": "/subscriptions/sub/.../subnets/default", + "vnet_id": "/subscriptions/sub/.../virtualNetworks/test-vnet", + "nic_id": "nic-id", + "nic_name": "nic-vm-1", + } + ) + azure_client.compute_client.virtual_machines.get.return_value = vm + + request = MagicMock() + request.resource_ids = ["vm-1"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert result[0]["instance_id"] == "vm-1" + azure_client.compute_client.virtual_machines.list.assert_not_called() + + +def test_release_returns_submitted_delete_metadata(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm_1 = MagicMock() + vm_1.name = "vm-1" + vm_1.vm_id = "guid-1" + vm_2 = MagicMock() + vm_2.name = "vm-2" + vm_2.vm_id = "guid-2" + azure_client.compute_client.virtual_machines.list.return_value = [vm_1, vm_2] + azure_client.compute_client.virtual_machines.get.side_effect = ResourceNotFoundError("NotFound") + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + result = run_operation( + handler.release_hosts_async( + machine_ids=["guid-1", "guid-2"], + resource_id="unused", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert result["provider_data"]["operation_status"] == "submitted" + assert result["provider_data"]["submitted_deletions"] == [ + {"requested_id": "guid-1", "vm_name": "vm-1"}, + {"requested_id": "guid-2", "vm_name": "vm-2"}, + ] + assert _deleted_vm_names(azure_client) == ["vm-1", "vm-2"] + + +def test_release_uses_direct_vm_name_lookup_without_listing_resource_group(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm = MagicMock() + vm.name = "vm-1" + azure_client.compute_client.virtual_machines.get.return_value = vm + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + result = run_operation( + handler.release_hosts_async( + machine_ids=["vm-1"], + resource_id="unused", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert result["provider_data"]["submitted_deletions"] == [ + {"requested_id": "vm-1", "vm_name": "vm-1"}, + ] + azure_client.compute_client.virtual_machines.list.assert_not_called() + + +def test_release_attempts_all_deletes_before_raising_aggregated_failure(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = SingleVMHandler(azure_client=azure_client, logger=logger) + + vm_1 = MagicMock() + vm_1.name = "vm-1" + vm_1.vm_id = "guid-1" + vm_2 = MagicMock() + vm_2.name = "vm-2" + vm_2.vm_id = "guid-2" + vm_3 = MagicMock() + vm_3.name = "vm-3" + vm_3.vm_id = "guid-3" + azure_client.compute_client.virtual_machines.list.return_value = [vm_1, vm_2, vm_3] + azure_client.compute_client.virtual_machines.get.side_effect = ResourceNotFoundError("NotFound") + + def _begin_delete(*, resource_group_name, vm_name): + if vm_name == "vm-2": + raise RuntimeError("delete blocked") + return MagicMock() + + azure_client.compute_client.virtual_machines.begin_delete.side_effect = _begin_delete + + with pytest.raises(TerminationError) as exc_info: + run_operation( + handler.release_hosts_async( + machine_ids=["guid-1", "guid-2", "guid-3"], + resource_id="unused", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + exc = exc_info.value + assert exc.resource_ids == ["guid-2"] + assert exc.details["submitted_deletions"] == [ + {"requested_id": "guid-1", "vm_name": "vm-1"}, + {"requested_id": "guid-3", "vm_name": "vm-3"}, + ] + assert exc.details["failed_deletions"] == [ + {"requested_id": "guid-2", "vm_name": "vm-2", "error": "delete blocked"}, + ] + assert _deleted_vm_names(azure_client) == ["vm-1", "vm-2", "vm-3"] + + +@pytest.mark.asyncio +async def test_resolve_vm_names_async_maps_vm_ids_via_resource_group_listing(): + azure_client = _make_azure_client() + handler = SingleVMHandler(azure_client=azure_client, logger=MagicMock()) + + vm_1 = MagicMock() + vm_1.name = "vm-1" + vm_1.vm_id = "11111111-1111-1111-1111-111111111111" + + async_compute = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + async_compute.virtual_machines.get = AsyncMock(side_effect=ResourceNotFoundError("NotFound")) + + async_compute.virtual_machines.list.return_value = AsyncPager([vm_1]) + + resolved = await handler._resolve_vm_names_async( + "test-rg", + ["11111111-1111-1111-1111-111111111111"], + ) + + assert resolved == ["vm-1"] + + +@pytest.mark.asyncio +async def test_resolve_vm_names_async_preserves_input_order_for_mixed_vm_names_and_ids(): + azure_client = _make_azure_client() + handler = SingleVMHandler(azure_client=azure_client, logger=MagicMock()) + + vm_1 = MagicMock() + vm_1.name = "vm-1" + vm_1.vm_id = "11111111-1111-1111-1111-111111111111" + + def _get_vm(*, resource_group_name, vm_name): + if resource_group_name == "test-rg" and vm_name == "vm-2": + vm = MagicMock() + vm.name = "vm-2" + return vm + raise ResourceNotFoundError("NotFound") + + async_compute = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + async_compute.virtual_machines.get = AsyncMock(side_effect=_get_vm) + + async_compute.virtual_machines.list.return_value = AsyncPager([vm_1]) + + resolved = await handler._resolve_vm_names_async( + "test-rg", + ["vm-2", "11111111-1111-1111-1111-111111111111"], + ) + + assert resolved == ["vm-2", "vm-1"] + + +@pytest.mark.asyncio +async def test_resolve_vm_names_async_maps_vm_ids_via_resource_group_listing_against_async_client(): + azure_client = _make_azure_client() + handler = SingleVMHandler(azure_client=azure_client, logger=MagicMock()) + + vm_1 = MagicMock() + vm_1.name = "vm-1" + vm_1.vm_id = "11111111-1111-1111-1111-111111111111" + + async_compute = MagicMock() + async_compute.virtual_machines.get = AsyncMock(side_effect=ResourceNotFoundError("NotFound")) + + async_compute.virtual_machines.list.return_value = AsyncPager([vm_1]) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + resolved = await handler._resolve_vm_names_async( + "test-rg", + ["11111111-1111-1111-1111-111111111111"], + ) + + assert resolved == ["vm-1"] diff --git a/tests/providers/azure/test_ssh_key_resolver.py b/tests/providers/azure/test_ssh_key_resolver.py new file mode 100644 index 000000000..c54ab15a7 --- /dev/null +++ b/tests/providers/azure/test_ssh_key_resolver.py @@ -0,0 +1,79 @@ +"""Focused tests for Azure SSH key resolution.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError + +from orb.providers.azure.infrastructure.services.ssh_key_resolver import resolve_ssh_keys_async + + +@pytest.mark.asyncio +async def test_resolve_ssh_keys_async_returns_inline_keys_without_sdk_lookup(): + compute_client = MagicMock() + + result = await resolve_ssh_keys_async( + ssh_key_name="orb-key", + ssh_public_keys=["ssh-rsa inline test@host"], + resource_group="test-rg", + compute_client=compute_client, + ) + + assert result == ["ssh-rsa inline test@host"] + compute_client.ssh_public_keys.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_resolve_ssh_keys_async_fetches_named_azure_ssh_key(): + compute_client = MagicMock() + ssh_resource = MagicMock() + ssh_resource.public_key = "ssh-rsa resolved test@host" + compute_client.ssh_public_keys.get = AsyncMock(return_value=ssh_resource) + + result = await resolve_ssh_keys_async( + ssh_key_name="orb-key", + ssh_public_keys=[], + resource_group="test-rg", + compute_client=compute_client, + ) + + assert result == ["ssh-rsa resolved test@host"] + + +@pytest.mark.asyncio +async def test_resolve_ssh_keys_async_requires_name_or_inline_key_data(): + with pytest.raises(ValueError, match="neither 'ssh_key_name' nor 'ssh_public_keys'"): + await resolve_ssh_keys_async( + ssh_key_name=None, + ssh_public_keys=[], + resource_group="test-rg", + compute_client=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_resolve_ssh_keys_async_rejects_missing_ssh_key_resource(): + compute_client = MagicMock() + compute_client.ssh_public_keys.get = AsyncMock(side_effect=ResourceNotFoundError("missing")) + + with pytest.raises(ValueError, match="was not found"): + await resolve_ssh_keys_async( + ssh_key_name="orb-key", + ssh_public_keys=[], + resource_group="test-rg", + compute_client=compute_client, + ) + + +@pytest.mark.asyncio +async def test_resolve_ssh_keys_async_rejects_transport_error(): + compute_client = MagicMock() + compute_client.ssh_public_keys.get = AsyncMock(side_effect=HttpResponseError("boom")) + + with pytest.raises(ValueError, match="Failed to resolve Azure SSH Public Key"): + await resolve_ssh_keys_async( + ssh_key_name="orb-key", + ssh_public_keys=[], + resource_group="test-rg", + compute_client=compute_client, + ) diff --git a/tests/providers/azure/test_vmss_cleanup.py b/tests/providers/azure/test_vmss_cleanup.py new file mode 100644 index 000000000..652342324 --- /dev/null +++ b/tests/providers/azure/test_vmss_cleanup.py @@ -0,0 +1,376 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from orb.providers.azure.infrastructure.vmss_cleanup import ( + PendingVmssCleanup, + VmssCleanupCoordinator, +) + + +def test_pending_vmss_cleanup_round_trips_metadata(): + cleanup = PendingVmssCleanup.create( + resource_group="test-rg", + vmss_name="vmss-demo", + machine_ids=["vm-a", "vm-a", "vm-b"], + delete_vmss_when_empty=True, + member_delete_submitted=True, + delete_retry_pending=True, + last_delete_error="transient delete failure", + ) + + restored = PendingVmssCleanup.from_metadata(cleanup.to_metadata()) + + assert restored.to_metadata() == { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a", "vm-b"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "last_delete_error": "transient delete failure", + } + + +def test_pending_vmss_cleanup_merge_preserves_ids_and_retry_state(): + base = PendingVmssCleanup.from_metadata( + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + } + ) + update = PendingVmssCleanup.from_metadata( + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-b"], + "delete_vmss_when_empty": True, + "delete_retry_pending": True, + "last_delete_error": "still deleting", + } + ) + + merged = base.combine_for_same_vmss(update) + + assert merged.to_status_detail() == { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a", "vm-b"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "last_delete_error": "still deleting", + } + + +@pytest.mark.asyncio +async def test_vmss_cleanup_coordinator_reconciles_delete_retry_state(): + logger = MagicMock() + coordinator = VmssCleanupCoordinator( + logger=logger, + get_vmss_member_count=AsyncMock(return_value=0), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=AsyncMock(side_effect=RuntimeError("delete blocked")), + ) + + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + } + ) + + await coordinator.reconcile( + resource_group="test-rg", + resource_ids=["vmss-demo"], + observed_ids=set(), + ) + + assert coordinator.status_metadata( + resource_group="test-rg", + resource_ids=["vmss-demo"], + ) == { + "termination_follow_up_pending": True, + "termination_follow_up_details": [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "delete_retry_count": 1, + "last_delete_error": "delete blocked", + } + ], + } + logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_vmss_cleanup_coordinator_recovers_delete_retry_state_after_record_replaces_entry(): + logger = MagicMock() + + async def begin_delete_vmss(**_: object) -> None: + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-b"], + "delete_vmss_when_empty": True, + } + } + } + ) + raise RuntimeError("delete blocked") + + coordinator = VmssCleanupCoordinator( + logger=logger, + get_vmss_member_count=AsyncMock(return_value=0), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=begin_delete_vmss, + ) + + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + } + ) + + await coordinator.reconcile( + resource_group="test-rg", + resource_ids=["vmss-demo"], + observed_ids=set(), + ) + + assert coordinator.status_metadata( + resource_group="test-rg", + resource_ids=["vmss-demo"], + ) == { + "termination_follow_up_pending": True, + "termination_follow_up_details": [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a", "vm-b"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "delete_retry_count": 1, + "last_delete_error": "delete blocked", + } + ], + } + logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_vmss_cleanup_coordinator_marks_terminal_failure_after_retry_exhaustion(): + logger = MagicMock() + begin_delete_vmss = AsyncMock(side_effect=RuntimeError("delete blocked")) + coordinator = VmssCleanupCoordinator( + logger=logger, + get_vmss_member_count=AsyncMock(return_value=0), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=begin_delete_vmss, + max_delete_retries=2, + ) + + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + } + ) + + for _ in range(2): + await coordinator.reconcile( + resource_group="test-rg", + resource_ids=["vmss-demo"], + observed_ids=set(), + ) + + await coordinator.reconcile( + resource_group="test-rg", + resource_ids=["vmss-demo"], + observed_ids=set(), + ) + + assert begin_delete_vmss.await_count == 2 + assert coordinator.status_metadata( + resource_group="test-rg", + resource_ids=["vmss-demo"], + ) == { + "termination_follow_up_pending": False, + "termination_follow_up_failed": True, + "termination_follow_up_details": [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + "delete_retry_count": 2, + "delete_retry_exhausted": True, + "last_delete_error": "delete blocked", + } + ], + } + + +def test_vmss_cleanup_coordinator_restores_pending_state_from_request_metadata(): + logger = MagicMock() + coordinator = VmssCleanupCoordinator( + logger=logger, + get_vmss_member_count=AsyncMock(return_value=1), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=AsyncMock(), + ) + + coordinator.restore_from_request_metadata( + { + "termination_requests": [ + { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + ] + } + ) + + assert coordinator.has_pending(resource_group="test-rg", resource_ids=["vmss-demo"]) is True + assert coordinator.status_metadata( + resource_group="test-rg", + resource_ids=["vmss-demo"], + )["termination_follow_up_details"] == [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": False, + } + ] + + +def test_vmss_cleanup_coordinator_clears_pending_state(): + coordinator = VmssCleanupCoordinator( + logger=MagicMock(), + get_vmss_member_count=AsyncMock(return_value=1), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=AsyncMock(), + ) + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + } + ) + + coordinator.clear() + + assert coordinator.has_pending(resource_group="test-rg", resource_ids=["vmss-demo"]) is False + + +@pytest.mark.asyncio +async def test_vmss_cleanup_coordinator_submits_delete_when_vmss_is_empty(): + begin_delete_vmss = AsyncMock() + coordinator = VmssCleanupCoordinator( + logger=MagicMock(), + get_vmss_member_count=AsyncMock(return_value=0), + vmss_exists=AsyncMock(return_value=True), + begin_delete_vmss=begin_delete_vmss, + ) + coordinator.record( + { + "provider_data": { + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + } + } + } + ) + + await coordinator.reconcile( + resource_group="test-rg", + resource_ids=["vmss-demo"], + observed_ids=set(), + ) + + begin_delete_vmss.assert_called_once_with(resource_group="test-rg", vmss_name="vmss-demo") + assert coordinator.status_metadata( + resource_group="test-rg", + resource_ids=["vmss-demo"], + ) == { + "termination_follow_up_pending": True, + "termination_follow_up_details": [ + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": True, + "delete_retry_pending": False, + } + ], + } + + +def test_pending_vmss_cleanup_defaults_member_delete_submission_for_legacy_metadata(): + restored = PendingVmssCleanup.from_metadata( + { + "resource_group": "test-rg", + "vmss_name": "vmss-demo", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "delete_retry_pending": True, + } + ) + + assert restored.member_delete_submitted is True diff --git a/tests/providers/azure/test_vmss_handler.py b/tests/providers/azure/test_vmss_handler.py new file mode 100644 index 000000000..decf78cce --- /dev/null +++ b/tests/providers/azure/test_vmss_handler.py @@ -0,0 +1,1085 @@ +"""Focused tests for VMSS handler behavior.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.mgmt.compute.models import OrchestrationMode + +from orb.providers.azure.domain.template.value_objects import AzureVMSSOrchestrationMode +from orb.providers.azure.exceptions.azure_exceptions import AzureValidationError, TerminationError +from orb.providers.azure.infrastructure.handlers.azure_handler import ( + RAISE_ON_STATUS_ERROR_METADATA_KEY, + AzureReleaseContext, +) +from orb.providers.azure.infrastructure.handlers.vmss_handler import VMSSHandler +from tests.providers.azure.strategy_test_support import ( + AsyncPager, + make_azure_template, + make_vmss_azure_client, + run_operation, +) + + +def _deleted_vm_names(azure_client: MagicMock) -> list[str]: + return [ + str(call.kwargs["vm_name"]) + for call in azure_client.compute_client.virtual_machines.begin_delete.call_args_list + ] + + +def _make_template(**overrides): + return make_azure_template( + template_id="azure-vmss-test", + provider_api="VMSS", + **overrides, + ) + + +def _make_azure_client() -> MagicMock: + return make_vmss_azure_client() + + +def test_acquire_hosts_submits_native_vmss_create_and_returns_submitted_status(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.return_value = ( + MagicMock() + ) + + request = MagicMock() + request.requested_count = 2 + request.request_id = "req-1" + request.metadata = {} + + result = run_operation(handler.acquire_hosts_async(request, _make_template())) + + assert result["success"] is True + assert result["resource_ids"] == [result["provider_data"]["vmss_name"]] + assert result["provider_data"]["provisioning_state"] == "creating" + assert result["provider_data"]["operation_status"] == "submitted" + create_call = azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.call_args.kwargs + assert create_call["resource_group_name"] == "test-rg" + assert create_call["vm_scale_set_name"] == result["provider_data"]["vmss_name"] + assert create_call["parameters"]["sku"]["capacity"] == 2 + + +@pytest.mark.asyncio +async def test_acquire_hosts_async_submits_native_vmss_create_and_returns_submitted_status(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machine_scale_sets.begin_create_or_update = AsyncMock( + return_value=MagicMock() + ) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + request = MagicMock() + request.requested_count = 2 + request.request_id = "req-async" + request.metadata = {} + + result = await handler.acquire_hosts_async(request, _make_template()) + + assert result["success"] is True + assert result["provider_data"]["operation_status"] == "submitted" + async_compute.virtual_machine_scale_sets.begin_create_or_update.assert_awaited_once() + + +def test_acquire_hosts_does_not_mutate_template_when_network_config_is_derived_from_subnet_ids(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.return_value = ( + MagicMock() + ) + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-network-copy" + request.metadata = {} + + template = _make_template( + network_config=None, + subnet_ids=["/subscriptions/.../subnets/derived"], + ) + + result = run_operation(handler.acquire_hosts_async(request, template)) + + assert result["success"] is True + assert template.network_config is None + create_call = ( + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.call_args + ) + subnet_id = create_call.kwargs["parameters"]["properties"]["virtualMachineProfile"][ + "networkProfile" + ]["networkInterfaceConfigurations"][0]["properties"]["ipConfigurations"][0]["properties"][ + "subnet" + ]["id"] + assert subnet_id == "/subscriptions/.../subnets/derived" + + +def test_acquire_hosts_rejects_multiple_subnet_ids_without_network_config(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-ambiguous-subnet" + request.metadata = {} + + template = _make_template( + network_config=None, + subnet_ids=[ + "/subscriptions/.../subnets/first", + "/subscriptions/.../subnets/second", + ], + ) + + with pytest.raises(AzureValidationError, match="support a single subnet"): + run_operation(handler.acquire_hosts_async(request, template)) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.assert_not_called() + + +def test_acquire_hosts_prefers_network_config_over_legacy_subnet_ids(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.return_value = ( + MagicMock() + ) + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-network-config-wins" + request.metadata = {} + + template = _make_template( + network_config={"subnet_id": "/subscriptions/.../subnets/explicit"}, + subnet_ids=[ + "/subscriptions/.../subnets/first", + "/subscriptions/.../subnets/second", + ], + ) + + result = run_operation(handler.acquire_hosts_async(request, template)) + + assert result["success"] is True + create_call = ( + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.call_args + ) + subnet_id = create_call.kwargs["parameters"]["properties"]["virtualMachineProfile"][ + "networkProfile" + ]["networkInterfaceConfigurations"][0]["properties"]["ipConfigurations"][0]["properties"][ + "subnet" + ]["id"] + assert subnet_id == "/subscriptions/.../subnets/explicit" + + +def test_acquire_hosts_raises_validation_error_when_no_subnet_is_available(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-no-subnet" + request.metadata = {} + + template = _make_template(network_config=None, subnet_ids=[]) + + with pytest.raises(AzureValidationError, match="No subnet specified"): + run_operation(handler.acquire_hosts_async(request, template)) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.assert_not_called() + + +def test_flexible_vmss_status_returns_only_member_vms(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + member_vm = MagicMock() + member_vm.name = "vmss-azure-test_abcd1234" + member_vm.instance_id = "vmss-azure-test_abcd1234" + member_vm.vm_id = "vm-guid-1" + member_vm.instance_view.statuses = [] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = ["1"] + + azure_client.subscription_id = "sub" + azure_client.compute_client.virtual_machines.list.return_value = [member_vm] + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert len(result) == 1 + assert result[0]["instance_id"] == "vmss-azure-test_abcd1234" + azure_client.compute_client.virtual_machines.list.assert_called_once_with( + resource_group_name="test-rg", + filter="'virtualMachineScaleSet/id' eq '/subscriptions/sub/resourceGroups/test-rg/providers/Microsoft.Compute/virtualMachineScaleSets/vmss-azure-test'", + expand="instanceView", + ) + + +@pytest.mark.asyncio +async def test_check_hosts_status_async_returns_listed_instances(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock( + return_value=[{"instance_id": "vmss-1", "provider_data": {"vmss_instance_id": "1"}}] + ) + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = await handler.check_hosts_status_async(request) + + assert result == [{"instance_id": "vmss-1", "provider_data": {"vmss_instance_id": "1"}}] + handler._list_vmss_instances_async.assert_awaited_once_with( + "test-rg", + "vmss-azure-test", + include_instance_view=True, + ) + + +@pytest.mark.asyncio +async def test_flexible_vmss_listing_uses_azure_side_filter_without_client_side_membership_scan(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.subscription_id = "sub" + async_compute = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + async_compute.virtual_machines.list.return_value = AsyncPager([]) + + result = await handler._list_vmss_instances_async( + resource_group="test-rg", + vmss_name="vmss-azure-test", + include_instance_view=False, + orchestration_mode=AzureVMSSOrchestrationMode.FLEXIBLE, + ) + + assert result == [] + async_compute.virtual_machines.list.assert_called_once_with( + resource_group_name="test-rg", + filter="'virtualMachineScaleSet/id' eq '/subscriptions/sub/resourceGroups/test-rg/providers/Microsoft.Compute/virtualMachineScaleSets/vmss-azure-test'", + ) + + +def test_flexible_vmss_status_uses_filtered_list_result_directly(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + member_vm = MagicMock() + member_vm.name = "vmss-azure-test_abcd1234" + member_vm.instance_id = "vmss-azure-test_abcd1234" + member_vm.vm_id = "vm-guid-1" + member_vm.instance_view.statuses = [] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = ["1"] + + azure_client.subscription_id = "sub" + azure_client.compute_client.virtual_machines.list.return_value = [member_vm] + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert len(result) == 1 + assert result[0]["instance_id"] == "vmss-azure-test_abcd1234" + azure_client.compute_client.virtual_machines.list.assert_called_once_with( + resource_group_name="test-rg", + filter="'virtualMachineScaleSet/id' eq '/subscriptions/sub/resourceGroups/test-rg/providers/Microsoft.Compute/virtualMachineScaleSets/vmss-azure-test'", + expand="instanceView", + ) + + +def test_vmss_instance_status_includes_structured_provisioning_errors(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + status = MagicMock() + status.code = "ProvisioningState/failed" + status.level = "Error" + status.message = "Allocation failed in zone 1" + status.display_status = "Provisioning failed" + + member_vm = MagicMock() + member_vm.instance_id = "3" + member_vm.name = "vmss-3" + member_vm.vm_id = "vm-guid-3" + member_vm.instance_view.statuses = [status] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = ["1"] + + azure_client.compute_client.virtual_machine_scale_set_vms.list.return_value = [member_vm] + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert result[0]["status"] == "failed" + assert result[0]["provider_data"]["fleet_errors"][0]["error_code"] == "ProvisioningStateFailed" + assert "Allocation failed" in result[0]["provider_data"]["fleet_errors"][0]["error_message"] + + +@pytest.mark.asyncio +async def test_release_hosts_async_submits_uniform_vmss_deletes(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machine_scale_sets.begin_delete_instances = AsyncMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._get_vmss_orchestration_mode_async = AsyncMock( + return_value=AzureVMSSOrchestrationMode.UNIFORM + ) + handler._list_vmss_instances_async = AsyncMock(return_value=[{"instance_id": "1"}]) + handler._resolve_vmss_instance_ids_async = AsyncMock(return_value=["1"]) + + result = await handler.release_hosts_async( + machine_ids=["vmss-1"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + + assert result is not None + assert result["provider_data"]["resolved_instance_ids"] == ["1"] + async_compute.virtual_machine_scale_sets.begin_delete_instances.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_release_hosts_async_marks_flexible_vmss_for_cleanup_when_last_instance_is_returned(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machines.begin_delete = AsyncMock() + async_compute.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._get_vmss_orchestration_mode_async = AsyncMock( + return_value=AzureVMSSOrchestrationMode.FLEXIBLE + ) + handler._list_vmss_instances_async = AsyncMock(return_value=[{"instance_id": "vm-a"}]) + + result = await handler.release_hosts_async( + machine_ids=["vm-a"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + + assert result == { + "provider_data": { + "resource_group": "test-rg", + "vmss_name": "vmss-azure-test", + "operation_status": "submitted", + "submitted_deletions": [{"requested_id": "vm-a", "vm_name": "vm-a"}], + "pending_resource_cleanup": { + "resource_group": "test-rg", + "vmss_name": "vmss-azure-test", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": True, + "delete_retry_pending": False, + }, + } + } + async_compute.virtual_machines.begin_delete.assert_awaited_once_with( + resource_group_name="test-rg", + vm_name="vm-a", + ) + async_compute.virtual_machine_scale_sets.begin_delete.assert_awaited_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-azure-test", + ) + + +@pytest.mark.asyncio +async def test_release_hosts_async_raises_on_partial_flexible_vmss_delete_failures(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machines.begin_delete = AsyncMock( + side_effect=[None, RuntimeError("delete failed")] + ) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._get_vmss_orchestration_mode_async = AsyncMock( + return_value=AzureVMSSOrchestrationMode.FLEXIBLE + ) + handler._list_vmss_instances_async = AsyncMock( + return_value=[{"instance_id": "vm-a"}, {"instance_id": "vm-b"}, {"instance_id": "vm-c"}] + ) + + with pytest.raises(TerminationError) as exc_info: + await handler.release_hosts_async( + machine_ids=["vm-a", "vm-b"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + + exc = exc_info.value + assert exc.resource_ids == ["vm-b"] + assert exc.details["submitted_deletions"] == [{"requested_id": "vm-a", "vm_name": "vm-a"}] + assert exc.details["failed_deletions"] == [ + {"requested_id": "vm-b", "vm_name": "vm-b", "error": "delete failed"} + ] + logger.error.assert_called() + + +@pytest.mark.asyncio +async def test_release_hosts_async_skips_empty_vmss_delete_after_flexible_member_failure(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machines.begin_delete = AsyncMock( + side_effect=[None, RuntimeError("delete failed")] + ) + async_compute.virtual_machine_scale_sets.begin_delete = AsyncMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + handler._get_vmss_orchestration_mode_async = AsyncMock( + return_value=AzureVMSSOrchestrationMode.FLEXIBLE + ) + handler._list_vmss_instances_async = AsyncMock( + return_value=[{"instance_id": "vm-a"}, {"instance_id": "vm-b"}] + ) + + with pytest.raises(TerminationError): + await handler.release_hosts_async( + machine_ids=["vm-a", "vm-b"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + + async_compute.virtual_machine_scale_sets.begin_delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_submit_vmss_delete_if_emptying_async_returns_retry_pending_when_delete_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + async_compute.virtual_machine_scale_sets.begin_delete = AsyncMock( + side_effect=RuntimeError("scale set still has deleting members") + ) + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + result = await handler._submit_vmss_delete_if_emptying_async( + resource_group="test-rg", + vmss_name="vmss-azure-test", + ) + + assert result == { + "delete_submitted": False, + "delete_retry_pending": True, + "last_delete_error": "scale set still has deleting members", + } + + +def test_vmss_status_raises_when_explicit_status_errors_are_fatal(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.get.side_effect = Exception( + "transient ARM failure" + ) + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = { + "resource_group": "test-rg", + RAISE_ON_STATUS_ERROR_METADATA_KEY: True, + } + + with pytest.raises(RuntimeError, match="Failed to list instances for VMSS 'vmss-azure-test'"): + run_operation(handler.check_hosts_status_async(request)) + + +def test_vmss_status_rejects_non_boolean_status_error_policy(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = { + "resource_group": "test-rg", + RAISE_ON_STATUS_ERROR_METADATA_KEY: "true", + } + + with pytest.raises(AzureValidationError, match="must be a boolean"): + run_operation(handler.check_hosts_status_async(request)) + + +def test_vmss_status_raises_when_best_effort_has_no_successful_observations(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.get.side_effect = Exception( + "transient ARM failure" + ) + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + with pytest.raises(RuntimeError, match="Failed to list instances for VMSS 'vmss-azure-test'"): + run_operation(handler.check_hosts_status_async(request)) + + +def test_vmss_status_raises_after_partial_success_when_status_errors_are_fatal(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + member_vm = MagicMock() + member_vm.instance_id = "1" + member_vm.vm_id = "vm-guid-1" + member_vm.name = "vmss_1" + member_vm.instance_view.statuses = [] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = [] + azure_client.compute_client.virtual_machine_scale_set_vms.list.side_effect = [ + [member_vm], + RuntimeError("transient ARM failure"), + ] + + request = MagicMock() + request.resource_ids = ["vmss-ok", "vmss-fail"] + request.metadata = { + "resource_group": "test-rg", + RAISE_ON_STATUS_ERROR_METADATA_KEY: True, + } + + with pytest.raises(RuntimeError, match="Failed to list instances for VMSS 'vmss-fail'"): + run_operation(handler.check_hosts_status_async(request)) + + +def test_vmss_status_populates_network_identity(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + nic_ref = MagicMock() + nic_ref.id = ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/networkInterfaces/nic-vmss-1" + ) + nic_ref.properties.primary = True + + member_vm = MagicMock() + member_vm.instance_id = "3" + member_vm.name = "vmss-3" + member_vm.vm_id = "vm-guid-3" + member_vm.instance_view.statuses = [] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = ["1"] + member_vm.network_profile.network_interfaces = [nic_ref] + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + return_value={ + "private_ip": "10.0.0.7", + "public_ip": None, + "subnet_id": ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet/subnets/default" + ), + "vnet_id": ( + "/subscriptions/sub/resourceGroups/test-rg/providers/" + "Microsoft.Network/virtualNetworks/test-vnet" + ), + "nic_id": nic_ref.id, + "nic_name": "nic-vmss-1", + } + ) + + azure_client.compute_client.virtual_machine_scale_set_vms.list.return_value = [member_vm] + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert result[0]["private_ip"] == "10.0.0.7" + assert result[0]["subnet_id"].endswith("/subnets/default") + assert result[0]["vpc_id"].endswith("/virtualNetworks/test-vnet") + assert result[0]["provider_data"]["nic_name"] == "nic-vmss-1" + + +def test_vmss_status_still_returns_instance_when_network_identity_resolution_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + azure_client.resolve_network_identity_from_vm_async = AsyncMock( + side_effect=AttributeError("missing network property") + ) + + member_vm = MagicMock() + member_vm.instance_id = "3" + member_vm.name = "vmss-3" + member_vm.vm_id = "vm-guid-3" + member_vm.instance_view.statuses = [] + member_vm.hardware_profile.vm_size = "Standard_D4s_v5" + member_vm.location = "eastus2" + member_vm.zones = ["1"] + + azure_client.compute_client.virtual_machine_scale_set_vms.list.return_value = [member_vm] + + request = MagicMock() + request.resource_ids = ["vmss-azure-test"] + request.metadata = {"resource_group": "test-rg"} + + result = run_operation(handler.check_hosts_status_async(request)) + + assert len(result) == 1 + assert result[0]["instance_id"] == "3" + assert result[0]["private_ip"] is None + assert result[0]["provider_data"]["nic_id"] is None + logger.warning.assert_called() + + +@pytest.mark.asyncio +async def test_vmss_resource_errors_surface_failed_scale_set_without_instances(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + vmss = MagicMock() + vmss.provisioning_state = "Failed" + vmss.statuses = [] + async_compute.virtual_machine_scale_sets.get = AsyncMock(return_value=vmss) + + errors = await handler.get_vmss_resource_errors_async("test-rg", "vmss-azure-test") + + assert errors[0]["error_code"] == "ProvisioningStateFailed" + assert errors[0]["instance_id"] == "vmss-azure-test" + + +@pytest.mark.asyncio +async def test_vmss_resource_errors_logs_and_returns_empty_list_when_vmss_lookup_fails(): + azure_client = MagicMock() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + async_compute = MagicMock() + azure_client.get_async_compute_client = AsyncMock(return_value=async_compute) + + async_compute.virtual_machine_scale_sets.get = AsyncMock(side_effect=RuntimeError("boom")) + + errors = await handler.get_vmss_resource_errors_async("test-rg", "vmss-azure-test") + + assert errors == [] + logger.warning.assert_called_once() + + +def test_vmss_release_deletes_only_requested_uniform_instances(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._resolve_vmss_instance_ids_async = AsyncMock(return_value=["3", "4"]) + handler._list_vmss_instances_async = AsyncMock( + return_value=[ + {"instance_id": "3"}, + {"instance_id": "4"}, + {"instance_id": "5"}, + ] + ) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + delete_poller = MagicMock() + azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances.return_value = ( + delete_poller + ) + + run_operation( + handler.release_hosts_async( + machine_ids=["3", "4"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + delete_call = azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances.call_args.kwargs + assert delete_call["resource_group_name"] == "test-rg" + assert delete_call["vm_scale_set_name"] == "vmss-azure-test" + delete_ids = delete_call["vm_instance_i_ds"] + if hasattr(delete_ids, "instance_ids"): + assert delete_ids.instance_ids == ["3", "4"] + else: + assert delete_ids["instance_ids"] == ["3", "4"] + + +def test_vmss_release_deletes_flexible_members_without_cleanup_metadata_when_vmss_not_empty(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock( + return_value=[ + {"instance_id": "vm-a"}, + {"instance_id": "vm-b"}, + {"instance_id": "vm-c"}, + ] + ) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + result = run_operation( + handler.release_hosts_async( + machine_ids=["vm-a", "vm-b"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert result["provider_data"]["operation_status"] == "submitted" + assert "pending_resource_cleanup" not in result["provider_data"] + assert result["provider_data"]["submitted_deletions"] == [ + {"requested_id": "vm-a", "vm_name": "vm-a"}, + {"requested_id": "vm-b", "vm_name": "vm-b"}, + ] + assert _deleted_vm_names(azure_client) == ["vm-a", "vm-b"] + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.assert_not_called() + + +def test_vmss_release_marks_flexible_vmss_for_cleanup_when_last_instance_is_returned(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock(return_value=[{"instance_id": "vm-a"}]) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + result = run_operation( + handler.release_hosts_async( + machine_ids=["vm-a"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert result["provider_data"]["operation_status"] == "submitted" + assert result["provider_data"]["submitted_deletions"] == [ + {"requested_id": "vm-a", "vm_name": "vm-a"} + ] + assert result["provider_data"]["pending_resource_cleanup"] == { + "resource_group": "test-rg", + "vmss_name": "vmss-azure-test", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": True, + "delete_retry_pending": False, + } + assert _deleted_vm_names(azure_client) == ["vm-a"] + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.assert_called_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-azure-test", + ) + + +def test_vmss_release_rejects_unresolved_flexible_member_ids(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock( + return_value=[ + {"instance_id": "vm-a"}, + {"instance_id": "vm-b"}, + {"instance_id": "vm-c"}, + ] + ) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + with pytest.raises(TerminationError) as exc_info: + run_operation( + handler.release_hosts_async( + machine_ids=["guid-a", "guid-b", "guid-c"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + exc = exc_info.value + assert exc.resource_ids == ["guid-a", "guid-b", "guid-c"] + assert exc.details["unresolved_ids"] == ["guid-a", "guid-b", "guid-c"] + azure_client.compute_client.virtual_machines.begin_delete.assert_not_called() + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.assert_not_called() + + +def test_vmss_release_rejects_unresolved_uniform_member_ids(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock( + return_value=[{"instance_id": "3"}, {"instance_id": "4"}] + ) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + with pytest.raises(TerminationError) as exc_info: + run_operation( + handler.release_hosts_async( + machine_ids=["missing"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + exc = exc_info.value + assert exc.resource_ids == ["missing"] + assert exc.details["unresolved_ids"] == ["missing"] + azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances.assert_not_called() + + +def test_vmss_release_resolves_flexible_vm_ids_to_vm_names(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock( + return_value=[ + { + "instance_id": "vm-a", + "name": "vm-a", + "provider_data": {"vm_id": "guid-a", "vm_name": "vm-a"}, + }, + { + "instance_id": "vm-b", + "name": "vm-b", + "provider_data": {"vm_id": "guid-b", "vm_name": "vm-b"}, + }, + ] + ) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + + result = run_operation( + handler.release_hosts_async( + machine_ids=["guid-a", "vm-b"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert _deleted_vm_names(azure_client) == ["vm-a", "vm-b"] + assert result["provider_data"]["submitted_deletions"] == [ + {"requested_id": "guid-a", "vm_name": "vm-a"}, + {"requested_id": "vm-b", "vm_name": "vm-b"}, + ] + assert result["provider_data"]["pending_resource_cleanup"]["machine_ids"] == [ + "guid-a", + "vm-b", + ] + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.assert_called_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-azure-test", + ) + + +def test_vmss_release_marks_uniform_vmss_for_cleanup_when_last_instance_is_returned(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._resolve_vmss_instance_ids_async = AsyncMock(return_value=["3"]) + handler._list_vmss_instances_async = AsyncMock(return_value=[{"instance_id": "3"}]) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.UNIFORM + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + + delete_instances_poller = MagicMock() + azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances.return_value = ( + delete_instances_poller + ) + result = run_operation( + handler.release_hosts_async( + machine_ids=["3"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + delete_call = azure_client.compute_client.virtual_machine_scale_sets.begin_delete_instances.call_args.kwargs + delete_ids = delete_call["vm_instance_i_ds"] + if hasattr(delete_ids, "instance_ids"): + assert delete_ids.instance_ids == ["3"] + else: + assert delete_ids["instance_ids"] == ["3"] + assert result["provider_data"]["pending_resource_cleanup"] == { + "resource_group": "test-rg", + "vmss_name": "vmss-azure-test", + "machine_ids": ["3"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": True, + "delete_retry_pending": False, + } + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.assert_called_once_with( + resource_group_name="test-rg", + vm_scale_set_name="vmss-azure-test", + ) + + +def test_vmss_release_surfaces_retry_pending_when_immediate_empty_vmss_delete_fails(): + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + handler._list_vmss_instances_async = AsyncMock(return_value=[{"instance_id": "vm-a"}]) + + vmss = MagicMock() + vmss.orchestration_mode = OrchestrationMode.FLEXIBLE + azure_client.compute_client.virtual_machine_scale_sets.get.return_value = vmss + azure_client.compute_client.virtual_machines.begin_delete.return_value = MagicMock() + azure_client.compute_client.virtual_machine_scale_sets.begin_delete.side_effect = RuntimeError( + "scale set still has deleting members" + ) + + result = run_operation( + handler.release_hosts_async( + machine_ids=["vm-a"], + resource_id="vmss-azure-test", + context=AzureReleaseContext(resource_group="test-rg"), + ) + ) + + assert result["provider_data"]["pending_resource_cleanup"] == { + "resource_group": "test-rg", + "vmss_name": "vmss-azure-test", + "machine_ids": ["vm-a"], + "delete_vmss_when_empty": True, + "member_delete_submitted": True, + "delete_submitted": False, + "delete_retry_pending": True, + "last_delete_error": "scale set still has deleting members", + } + + +def test_acquire_hosts_classifies_quota_runtime_error_as_quota_exceeded(): + """A RuntimeError whose message implies a quota fault classifies as QuotaExceededError. + + The classifier falls back to message-based string matching only when no canonical + Azure error code is available — which is exactly the case for a generic RuntimeError. + """ + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.side_effect = ( + RuntimeError("quota exceeded") + ) + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-null-message" + request.metadata = {} + + with pytest.raises(Exception) as exc_info: + run_operation(handler.acquire_hosts_async(request, _make_template())) + + assert exc_info.type.__name__ == "QuotaExceededError" + + +def test_acquire_hosts_does_not_misclassify_unrelated_error_with_quota_in_resource_name(): + """The fix the classifier guards against: tag/resource-name collisions. + + Before the canonical-code-first refactor, an exception whose canonical code was + a real, non-quota Azure error would still be misclassified as a quota fault if the + surfaced message happened to contain the substring "quota" or "exceeded" — e.g. + because a resource name, tag, or correlation id contained those tokens. + """ + from azure.core.exceptions import HttpResponseError + + azure_client = _make_azure_client() + logger = MagicMock() + handler = VMSSHandler(azure_client=azure_client, logger=logger) + + fake_error = MagicMock() + fake_error.code = "InvalidParameter" + fake_error.message = "subnet 'team-quota-subnet-1' does not exist" + + response = MagicMock() + response.status_code = 400 + + exc = HttpResponseError(response=response) + exc.error = fake_error + + azure_client.compute_client.virtual_machine_scale_sets.begin_create_or_update.side_effect = exc + + request = MagicMock() + request.requested_count = 1 + request.request_id = "req-quota-in-name" + request.metadata = {} + + with pytest.raises(Exception) as exc_info: + run_operation(handler.acquire_hosts_async(request, _make_template())) + + assert exc_info.type.__name__ == "AzureValidationError" diff --git a/uv.lock b/uv.lock index 663176142..c0a004958 100644 --- a/uv.lock +++ b/uv.lock @@ -323,6 +323,103 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ef/c3/f30a7a63e664acc7c2545ca0491b6ce8264536e0e5cad3965f1d1b91e960/aws_xray_sdk-2.15.0-py2.py3-none-any.whl", hash = "sha256:422d62ad7d52e373eebb90b642eb1bb24657afe03b22a8df4a8b2e5108e278a3", size = 103228, upload-time = "2025-10-29T21:00:24.12Z" }, ] +[[package]] +name = "azure-core" +version = "1.41.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a6/f3/b416179e408990df5db0d516283022dde0f5d0111d98c1a848e41853e81c/azure_core-1.41.0.tar.gz", hash = "sha256:f46ff5dfcd230f25cf1c19e8a34b8dc08a337b2503e268bb600a16c00db8ad5a", size = 381042, upload-time = "2026-05-07T23:30:54.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/db/325c6d7312d2200251c52323878281045aaffcb5586612296484e4280eaa/azure_core-1.41.0-py3-none-any.whl", hash = "sha256:522b4011e8180b1a3dcd2024396a4e7fe9ac37fb8597db47163d230b5efe892d", size = 220920, upload-time = "2026-05-07T23:30:56.357Z" }, +] + +[[package]] +name = "azure-identity" +version = "1.25.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "msal" }, + { name = "msal-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/0e/3a63efb48aa4a5ae2cfca61ee152fbcb668092134d3eb8bfda472dd5c617/azure_identity-1.25.3.tar.gz", hash = "sha256:ab23c0d63015f50b630ef6c6cf395e7262f439ce06e5d07a64e874c724f8d9e6", size = 286304, upload-time = "2026-03-13T01:12:20.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" }, +] + +[[package]] +name = "azure-mgmt-compute" +version = "38.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-mgmt-core" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/74/d756dbdcc8d43a53950c5d8b1ae7745223ba46205ba6037ec9d448f50062/azure_mgmt_compute-38.1.0.tar.gz", hash = "sha256:dbbe8355472bf58e9b358041d0124ca6b0fffff79d60aa6d0f094d8644ccffe5", size = 527915, upload-time = "2026-06-22T06:57:26.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/29/06d4e0f00ac41b3b55e30aa565bbc98da1d04f20ec4b66e76b67017b0edd/azure_mgmt_compute-38.1.0-py3-none-any.whl", hash = "sha256:c119b41fd1a52a742bfbf20e6570bba81d22ed665f5bc05e9552382d431044a5", size = 447701, upload-time = "2026-06-22T06:57:28.671Z" }, +] + +[[package]] +name = "azure-mgmt-core" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3e/99/fa9e7551313d8c7099c89ebf3b03cd31beb12e1b498d575aa19bb59a5d04/azure_mgmt_core-1.6.0.tar.gz", hash = "sha256:b26232af857b021e61d813d9f4ae530465255cb10b3dde945ad3743f7a58e79c", size = 30818, upload-time = "2025-07-03T02:02:24.093Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/26/c79f962fd3172b577b6f38685724de58b6b4337a51d3aad316a43a4558c6/azure_mgmt_core-1.6.0-py3-none-any.whl", hash = "sha256:0460d11e85c408b71c727ee1981f74432bc641bb25dfcf1bb4e90a49e776dbc4", size = 29310, upload-time = "2025-07-03T02:02:25.203Z" }, +] + +[[package]] +name = "azure-mgmt-network" +version = "31.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-mgmt-core" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/07/10/0c6e043eaaa21574464cd36ee15f9acd71f81c647ad89f72bd3e938d071b/azure_mgmt_network-31.0.1.tar.gz", hash = "sha256:2112f720eedca76173be4d47acb8bc58af16a90debc7b601bf11b37faaaed137", size = 902532, upload-time = "2026-07-02T05:44:23.634Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d6/60/2e473151925f893b0629e3cb913bf4fce82edbc5da884607b9a5e636b8d2/azure_mgmt_network-31.0.1-py3-none-any.whl", hash = "sha256:cb36bff1a08237211093b0c3b900cadfd344d08d4a437d526c1239fee3331d0c", size = 725755, upload-time = "2026-07-02T05:44:25.803Z" }, +] + +[[package]] +name = "azure-mgmt-resource" +version = "26.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-mgmt-core" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/e1/c6196381cc6bb3eff9e24f5bf31b7fd962651ddbd6120bc2a4c482a7af16/azure_mgmt_resource-26.0.0.tar.gz", hash = "sha256:54227c5db06d3be4dd4b77ad885b273e9eb2606d1de653c5b534066cfbb60c0c", size = 117372, upload-time = "2026-06-24T09:48:29.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/6e/3e776a0817ebe3558fbd8d3f98b1bd8fa75a46596e5cf9bb4377960c4611/azure_mgmt_resource-26.0.0-py3-none-any.whl", hash = "sha256:ef36559a4c0843b4272aa584efb455c686e4c3f6357421b38c82f405c01a9a9f", size = 102492, upload-time = "2026-06-24T09:48:30.854Z" }, +] + +[[package]] +name = "azure-mgmt-resource-subscriptions" +version = "1.0.0b2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-mgmt-core" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/d2/1d42399ff4db0b3aa1c86ffde853199e7d34ca8416967a38418a883251dd/azure_mgmt_resource_subscriptions-1.0.0b2.tar.gz", hash = "sha256:9da8a9532469d738f4a24950565994f3627272edf63d0798c9ef523435e0f4a4", size = 57164, upload-time = "2026-05-22T08:33:50.352Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/c1/a16f83c1fcde7e3c5f4e85768afdd55975e6cfa4d04b83ef9d613a3757c8/azure_mgmt_resource_subscriptions-1.0.0b2-py3-none-any.whl", hash = "sha256:bbd458f54a5f273674a1b0fe746aceee5dd5b582b28e34ddf22acc82f1e5c7a7", size = 64649, upload-time = "2026-05-22T08:33:51.646Z" }, +] + [[package]] name = "babel" version = "2.18.0" @@ -1797,6 +1894,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/56/4cc7fc9e9e3f38fd324f24f8afe0ad8bb5fa41283f37f1aaf9de0612c968/ipython-8.39.0-py3-none-any.whl", hash = "sha256:bb3c51c4fa8148ab1dea07a79584d1c854e234ea44aa1283bcb37bc75054651f", size = 831849, upload-time = "2026-03-27T10:02:07.846Z" }, ] +[[package]] +name = "isodate" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/54/4d/e940025e2ce31a8ce1202635910747e5a87cc3a6a6bb2d00973375014749/isodate-0.7.2.tar.gz", hash = "sha256:4cd1aa0f43ca76f4a6c6c0292a85f40b35ec2e43e315b59f06e6d32171a953e6", size = 29705, upload-time = "2024-10-08T23:04:11.5Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, +] + [[package]] name = "jaraco-classes" version = "3.4.0" @@ -2703,6 +2809,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] +[[package]] +name = "msal" +version = "1.37.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9a/99/d840198ecf6e8057bbc937f129ae940404485d736cda73253bbff9537f01/msal-1.37.0.tar.gz", hash = "sha256:1b1672a33ee467c1d70b341bb16cafd51bb3c817147a95b93263794b03971bec", size = 182444, upload-time = "2026-05-29T19:49:05.561Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/b0/d807279f4b55d16d1f120d5ac4344c6e39b56732e2a224d40bded7fd67ad/msal-1.37.0-py3-none-any.whl", hash = "sha256:dd17e95a7c71bce75e8108113438ba7c4a086b3bcad4f57a8c09b7af3d753c2d", size = 123725, upload-time = "2026-05-29T19:49:04.335Z" }, +] + +[[package]] +name = "msal-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/99/5d239b6156eddf761a636bded1118414d161bd6b7b37a9335549ed159396/msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4", size = 23315, upload-time = "2025-03-14T23:51:03.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, +] + [[package]] name = "msgpack" version = "1.2.1" @@ -3248,7 +3380,7 @@ wheels = [ [[package]] name = "orb-py" -version = "1.7.0" +version = "1.8.0" source = { editable = "." } dependencies = [ { name = "boto3" }, @@ -3269,9 +3401,16 @@ dependencies = [ [package.optional-dependencies] all = [ { name = "alembic" }, + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-mgmt-compute" }, + { name = "azure-mgmt-network" }, + { name = "azure-mgmt-resource" }, + { name = "azure-mgmt-resource-subscriptions" }, { name = "boto3" }, { name = "botocore" }, { name = "fastapi" }, + { name = "httpx" }, { name = "jinja2" }, { name = "kubernetes" }, { name = "opentelemetry-api" }, @@ -3290,8 +3429,15 @@ all = [ { name = "uvicorn" }, ] all-providers = [ + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-mgmt-compute" }, + { name = "azure-mgmt-network" }, + { name = "azure-mgmt-resource" }, + { name = "azure-mgmt-resource-subscriptions" }, { name = "boto3" }, { name = "botocore" }, + { name = "httpx" }, { name = "kubernetes" }, ] api = [ @@ -3304,6 +3450,15 @@ aws = [ { name = "boto3" }, { name = "botocore" }, ] +azure = [ + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-mgmt-compute" }, + { name = "azure-mgmt-network" }, + { name = "azure-mgmt-resource" }, + { name = "azure-mgmt-resource-subscriptions" }, + { name = "httpx" }, +] ci = [ { name = "alembic" }, { name = "bandit" }, @@ -3365,6 +3520,12 @@ cli = [ ] dev = [ { name = "alembic" }, + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-mgmt-compute" }, + { name = "azure-mgmt-network" }, + { name = "azure-mgmt-resource" }, + { name = "azure-mgmt-resource-subscriptions" }, { name = "bandit" }, { name = "bandit-sarif-formatter" }, { name = "boto3" }, @@ -3494,6 +3655,11 @@ test-aws = [ { name = "requests-mock" }, { name = "responses" }, ] +test-azure = [ + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-mgmt-compute" }, +] test-k8s = [ { name = "kmock" }, { name = "kubernetes" }, @@ -3653,6 +3819,15 @@ requires-dist = [ { name = "alembic", marker = "extra == 'ci'", specifier = ">=1.13" }, { name = "alembic", marker = "extra == 'k8s-legacy'" }, { name = "alembic", marker = "extra == 'sql'", specifier = ">=1.13" }, + { name = "azure-core", marker = "extra == 'azure'", specifier = ">=1.38.2" }, + { name = "azure-core", marker = "extra == 'test-azure'", specifier = ">=1.38.2" }, + { name = "azure-identity", marker = "extra == 'azure'", specifier = ">=1.25.2" }, + { name = "azure-identity", marker = "extra == 'test-azure'", specifier = ">=1.25.2" }, + { name = "azure-mgmt-compute", marker = "extra == 'azure'", specifier = ">=37.2.0" }, + { name = "azure-mgmt-compute", marker = "extra == 'test-azure'", specifier = ">=37.2.0" }, + { name = "azure-mgmt-network", marker = "extra == 'azure'", specifier = ">=30.2.0" }, + { name = "azure-mgmt-resource", marker = "extra == 'azure'", specifier = ">=25.0.0" }, + { name = "azure-mgmt-resource-subscriptions", marker = "extra == 'azure'", specifier = ">=1.0.0b1" }, { name = "bandit", marker = "extra == 'ci'", specifier = ">=1.7.5,<2.0.0" }, { name = "bandit-sarif-formatter", marker = "extra == 'ci'", specifier = ">=1.1.1,<2.0.0" }, { name = "boto3", specifier = ">=1.42.21" }, @@ -3670,6 +3845,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'k8s-legacy'", specifier = ">=0.128.0" }, { name = "filelock", specifier = ">=3.20.1" }, { name = "git-changelog", marker = "extra == 'dev'", specifier = ">=2.6.0,<3.0.0" }, + { name = "httpx", marker = "extra == 'azure'", specifier = ">=0.27.0" }, { name = "httpx", marker = "extra == 'ci'", specifier = ">=0.27.0" }, { name = "import-linter", marker = "extra == 'ci'", specifier = ">=2.0,<3.0.0" }, { name = "ipdb", marker = "extra == 'dev'", specifier = ">=0.13.13,<1.0.0" }, @@ -3710,6 +3886,7 @@ requires-dist = [ { name = "orb-py", extras = ["aws"], marker = "extra == 'all-providers'" }, { name = "orb-py", extras = ["aws"], marker = "extra == 'test-aws'" }, { name = "orb-py", extras = ["aws", "monitoring"], marker = "extra == 'monitoring-aws'" }, + { name = "orb-py", extras = ["azure"], marker = "extra == 'all-providers'" }, { name = "orb-py", extras = ["ci", "all-providers"], marker = "extra == 'dev'" }, { name = "orb-py", extras = ["cli", "api", "sql", "monitoring", "all-providers"], marker = "extra == 'all'" }, { name = "orb-py", extras = ["k8s"], marker = "extra == 'all-providers'" }, @@ -3779,7 +3956,7 @@ requires-dist = [ { name = "werkzeug", marker = "extra == 'ci'", specifier = ">=3.1.6" }, { name = "wheel", marker = "extra == 'ci'", specifier = ">=0.41.3,<1.0.0" }, ] -provides-extras = ["sql", "aws", "k8s", "all-providers", "k8s-legacy", "ui", "cli", "api", "otel", "monitoring", "monitoring-aws", "all", "test-aws", "test-k8s", "ci", "dev"] +provides-extras = ["sql", "aws", "k8s", "azure", "all-providers", "k8s-legacy", "ui", "cli", "api", "otel", "monitoring", "monitoring-aws", "all", "test-aws", "test-azure", "test-k8s", "ci", "dev"] [package.metadata.requires-dev] ci = [ @@ -4569,6 +4746,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" }, ] +[package.optional-dependencies] +crypto = [ + { name = "cryptography" }, +] + [[package]] name = "pymdown-extensions" version = "11.0.1"