fix: small issues

Author: ikemilian-lewis

src/orb/providers/azure/cli/azure_cli_spec.py

@@ -116,14 +116,10 @@ def validate_add(self, args: argparse.Namespace) -> list[str]:
 
     def generate_name(self, args: argparse.Namespace) -> str:
         """Generate a provider instance name from Azure subscription and location."""
-        try:
-            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}"
-        except Exception:
-            pass
-        return "azure_default"
+        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."""

src/orb/providers/azure/configuration/template_extension.py

@@ -194,9 +194,13 @@ def to_template_defaults(self) -> dict[str, Any]:
             defaults["placement_regions"] = self.placement_regions
         if self.placement_zones:
             defaults["placement_zones"] = self.placement_zones
-        if self.os_disk_size_gb is not None:
-            defaults["os_disk"] = {
+        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",
-                "disk_size_gb": self.os_disk_size_gb,
             }
+            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

src/orb/providers/azure/infrastructure/adapters/azure_validation_adapter.py

@@ -53,19 +53,11 @@ def validate_provider_api(self, api: str) -> bool:
         Returns:
             True if the API is supported by azure configuration
         """
-        try:
-            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
-
-        except Exception as e:
-            self._logger.error("Error validating azure provider API %s: %s", api, e)
-            # Fall back to canonical API names so validation still works without helper wiring.
-            return api in {"VMSS", "CycleCloud", "SingleVM", "VMSSUniform"}
+        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]:
         """
@@ -74,12 +66,7 @@ def get_supported_provider_apis(self) -> list[str]:
         Returns:
             List of supported azure provider API identifiers
         """
-        try:
-            return sorted(get_supported_apis())
-        except Exception as e:
-            self._logger.error("Error getting supported azure APIs: %s", e)
-            # Fallback to hardcoded list for safety
-            return ["VMSS", "CycleCloud", "SingleVM", "VMSSUniform"]
+        return sorted(get_supported_apis())
 
     @staticmethod
     def get_api_capabilities(api: str) -> dict[str, Any]:

src/orb/providers/azure/infrastructure/handlers/single_vm_handler.py

@@ -160,14 +160,15 @@ def _build_status_result(
             "resource_id": str(vm.name),
             "cloud_host_id": vm.vm_id or vm.name,
             "vm_name": str(vm.name),
-            "vm_id": str(vm.vm_id),
             "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),

src/orb/providers/azure/strategy/azure_provider_strategy.py

@@ -213,11 +213,8 @@ def deployment_service(self) -> Optional[AzureDeploymentService]:
     @property
     def handlers(self) -> dict[str, AzureHandler]:
         """Get handler mapping, including any explicit test overrides."""
-        handlers = dict(self._handlers)
         handler_factory = self._get_handler_factory()
-        if handler_factory is None:
-            return handlers
-        handlers.update(handler_factory.get_all_handlers())
+        handlers = handler_factory.get_all_handlers() if handler_factory is not None else {}
         handlers.update(self._handlers)
         return handlers
 
@@ -442,10 +439,6 @@ async def execute_operation_async(self, operation: ProviderOperation) -> Provide
 
     def get_capabilities(self) -> ProviderCapabilities:
         """Get Azure provider capabilities."""
-        # TODO: Keep Azure and AWS capability metadata dynamic together.
-        # These example regions / instance types and operational heuristics are
-        # still hard-coded in both providers, we should evaluate if they can be made
-        # dynamic
         return ProviderCapabilities(
             provider_type="azure",
             supported_operations=[
@@ -457,8 +450,8 @@ def get_capabilities(self) -> ProviderCapabilities:
                 ProviderOperationType.GET_AVAILABLE_TEMPLATES,
                 ProviderOperationType.HEALTH_CHECK,
             ],
+            supported_apis=get_supported_apis(),
             features={
-                "supported_apis": get_supported_apis(),
                 "api_capabilities": get_supported_api_capabilities(),
                 "instance_management": True,
                 "spot_instances": True,
@@ -470,29 +463,6 @@ def get_capabilities(self) -> ProviderCapabilities:
                 "key_pairs": True,
                 "tags_support": True,
                 "monitoring": True,
-                "regions": ["eastus", "eastus2", "westus2", "westeurope", "northeurope"],
-                "instance_types": [
-                    "Standard_D2s_v5",
-                    "Standard_D4s_v5",
-                    "Standard_D8s_v5",
-                    "Standard_E4s_v5",
-                    "Standard_F4s_v2",
-                ],
-                "max_instances_per_request": 1000,
-                "supports_windows": False,
-                "supports_linux": True,
-            },
-            limitations={
-                "max_concurrent_requests": 100,
-                "rate_limit_per_second": 20,
-                "max_instance_lifetime_hours": 8760,
-                "requires_vpc": True,
-                "requires_key_pair": True,
-            },
-            performance_metrics={
-                "typical_create_time_seconds": 120,
-                "typical_terminate_time_seconds": 60,
-                "health_check_timeout_seconds": 15,
             },
         )
 

tests/providers/azure/test_azure_config_and_infra.py

@@ -422,6 +422,31 @@ def test_defaults_include_explicit_vm_size(self):
         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,

tests/providers/azure/test_single_vm_handler.py

@@ -14,6 +14,10 @@
     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,
@@ -72,6 +76,36 @@ 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()