Files
2026-07-13 13:17:40 +08:00

401 lines
16 KiB
Python

"""Tests for Azure autoscaler availability zone functionality."""
import copy
import unittest
from unittest.mock import Mock, patch
from ray.autoscaler._private._azure.node_provider import AzureNodeProvider
class TestAzureAvailabilityZones(unittest.TestCase):
"""Test cases for Azure autoscaler availability zone support."""
def setUp(self):
"""Set up test fixtures."""
self.provider_config = {
"resource_group": "test-rg",
"location": "westus2",
"subscription_id": "test-sub-id",
}
self.cluster_name = "test-cluster"
# Create a mock provider that doesn't initialize Azure clients
with patch.object(
AzureNodeProvider,
"__init__",
lambda self, provider_config, cluster_name: None,
):
self.provider = AzureNodeProvider(self.provider_config, self.cluster_name)
self.provider.provider_config = self.provider_config
self.provider.cluster_name = self.cluster_name
def test_parse_availability_zones_none_input(self):
"""Test _parse_availability_zones with None input returns empty list."""
result = self.provider._parse_availability_zones(None)
self.assertEqual(result, [])
def test_parse_availability_zones_empty_string(self):
"""Test _parse_availability_zones with empty string returns empty list."""
result = self.provider._parse_availability_zones("")
self.assertEqual(result, [])
def test_parse_availability_zones_auto(self):
"""Test _parse_availability_zones with 'auto' returns empty list."""
result = self.provider._parse_availability_zones("auto")
self.assertEqual(result, [])
def test_parse_availability_zones_whitespace_only(self):
"""Test _parse_availability_zones with whitespace-only string returns empty list."""
result = self.provider._parse_availability_zones(" ")
self.assertEqual(result, [])
def test_parse_availability_zones_single_zone(self):
"""Test _parse_availability_zones with single zone string."""
result = self.provider._parse_availability_zones("1")
self.assertEqual(result, ["1"])
def test_parse_availability_zones_multiple_zones(self):
"""Test _parse_availability_zones with comma-separated zones."""
result = self.provider._parse_availability_zones("1,2,3")
self.assertEqual(result, ["1", "2", "3"])
def test_parse_availability_zones_zones_with_spaces(self):
"""Test _parse_availability_zones with spaces around zones."""
result = self.provider._parse_availability_zones("1, 2, 3")
self.assertEqual(result, ["1", "2", "3"])
def test_parse_availability_zones_zones_with_extra_spaces(self):
"""Test _parse_availability_zones with extra spaces and tabs."""
result = self.provider._parse_availability_zones(" 1 , 2 , 3 ")
self.assertEqual(result, ["1", "2", "3"])
def test_parse_availability_zones_none_disable_case_insensitive(self):
"""Test _parse_availability_zones with 'none' variations disables zones."""
test_cases = ["none", "None", "NONE"]
for case in test_cases:
with self.subTest(case=case):
result = self.provider._parse_availability_zones(case)
self.assertIsNone(result)
def test_parse_availability_zones_null_disable_case_insensitive(self):
"""Test _parse_availability_zones with 'null' variations disables zones."""
test_cases = ["null", "Null", "NULL"]
for case in test_cases:
with self.subTest(case=case):
result = self.provider._parse_availability_zones(case)
self.assertIsNone(result)
def test_parse_availability_zones_invalid_type(self):
"""Test _parse_availability_zones with invalid input type raises ValueError."""
with self.assertRaises(ValueError) as context:
self.provider._parse_availability_zones(123)
self.assertIn("availability_zone must be a string", str(context.exception))
self.assertIn("got int: 123", str(context.exception))
def test_parse_availability_zones_list_input_invalid(self):
"""Test _parse_availability_zones with list input raises ValueError."""
with self.assertRaises(ValueError) as context:
self.provider._parse_availability_zones(["1", "2", "3"])
self.assertIn("availability_zone must be a string", str(context.exception))
def test_parse_availability_zones_dict_input_invalid(self):
"""Test _parse_availability_zones with dict input raises ValueError."""
with self.assertRaises(ValueError) as context:
self.provider._parse_availability_zones({"zones": ["1", "2"]})
self.assertIn("availability_zone must be a string", str(context.exception))
def test_parse_availability_zones_numeric_zones(self):
"""Test _parse_availability_zones with numeric zone strings."""
result = self.provider._parse_availability_zones("1,2,3")
self.assertEqual(result, ["1", "2", "3"])
def test_parse_availability_zones_alpha_zones(self):
"""Test _parse_availability_zones with alphabetic zone strings."""
result = self.provider._parse_availability_zones("east,west,central")
self.assertEqual(result, ["east", "west", "central"])
def test_parse_availability_zones_mixed_zones(self):
"""Test _parse_availability_zones with mixed numeric and alpha zones."""
result = self.provider._parse_availability_zones("1,zone-b,3")
self.assertEqual(result, ["1", "zone-b", "3"])
class TestAzureAvailabilityZonePrecedence(unittest.TestCase):
"""Test cases for Azure availability zone precedence logic."""
def setUp(self):
"""Set up test fixtures."""
self.base_provider_config = {
"resource_group": "test-rg",
"location": "westus2",
"subscription_id": "test-sub-id",
}
self.cluster_name = "test-cluster"
def _create_mock_provider(self, provider_config=None):
"""Create a mock Azure provider for testing."""
config = copy.deepcopy(self.base_provider_config)
if provider_config:
config.update(provider_config)
with patch.object(
AzureNodeProvider,
"__init__",
lambda self, provider_config, cluster_name: None,
):
provider = AzureNodeProvider(config, self.cluster_name)
provider.provider_config = config
provider.cluster_name = self.cluster_name
# Mock the validation method to avoid Azure API calls
provider._validate_zones_for_node_pool = Mock(
side_effect=lambda zones, location, vm_size: zones
)
return provider
def _extract_zone_logic(self, provider, node_config):
"""Extract zone determination logic similar to _create_node method."""
node_availability_zone = node_config.get("azure_arm_parameters", {}).get(
"availability_zone"
)
provider_availability_zone = provider.provider_config.get("availability_zone")
if node_availability_zone is not None:
return (
provider._parse_availability_zones(node_availability_zone),
"node config availability_zone",
)
elif provider_availability_zone is not None:
return (
provider._parse_availability_zones(provider_availability_zone),
"provider availability_zone",
)
else:
return ([], "default")
def test_node_availability_zone_overrides_provider(self):
"""Test that node-level availability_zone overrides provider-level."""
provider = self._create_mock_provider({"availability_zone": "1,2"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "3",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, ["3"])
self.assertEqual(source, "node config availability_zone")
def test_provider_availability_zone_used_when_no_node_override(self):
"""Test that provider-level availability_zone is used when no node override."""
provider = self._create_mock_provider({"availability_zone": "1,2"})
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, ["1", "2"])
self.assertEqual(source, "provider availability_zone")
def test_none_disables_zones_at_node_level(self):
"""Test that 'none' at node level disables zones even with provider zones."""
provider = self._create_mock_provider({"availability_zone": "1,2"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "none",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertIsNone(zones)
self.assertEqual(source, "node config availability_zone")
def test_no_zones_when_neither_provider_nor_node_specify(self):
"""Test default behavior when neither provider nor node specify zones."""
provider = self._create_mock_provider()
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, [])
self.assertEqual(source, "default")
def test_node_empty_string_overrides_provider_zones(self):
"""Test that node empty string overrides provider zones (auto-selection)."""
provider = self._create_mock_provider({"availability_zone": "1,2"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, [])
self.assertEqual(source, "node config availability_zone")
def test_node_auto_overrides_provider_zones(self):
"""Test that node 'auto' overrides provider zones (auto-selection)."""
provider = self._create_mock_provider({"availability_zone": "1,2"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "auto",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, [])
self.assertEqual(source, "node config availability_zone")
def test_provider_none_disables_zones(self):
"""Test that provider-level 'none' disables zones."""
provider = self._create_mock_provider({"availability_zone": "none"})
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertIsNone(zones)
self.assertEqual(source, "provider availability_zone")
def test_provider_empty_string_allows_auto_selection(self):
"""Test that provider-level empty string allows auto-selection."""
provider = self._create_mock_provider({"availability_zone": ""})
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, [])
self.assertEqual(source, "provider availability_zone")
def test_provider_auto_allows_auto_selection(self):
"""Test that provider-level 'auto' allows auto-selection."""
provider = self._create_mock_provider({"availability_zone": "auto"})
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, [])
self.assertEqual(source, "provider availability_zone")
def test_node_null_overrides_provider_zones(self):
"""Test that node-level 'null' overrides provider zones."""
provider = self._create_mock_provider({"availability_zone": "1,2,3"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "null",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertIsNone(zones)
self.assertEqual(source, "node config availability_zone")
def test_provider_null_disables_zones(self):
"""Test that provider-level 'null' disables zones."""
provider = self._create_mock_provider({"availability_zone": "NULL"})
node_config = {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertIsNone(zones)
self.assertEqual(source, "provider availability_zone")
def test_complex_override_scenario(self):
"""Test complex scenario with multiple node types and different overrides."""
provider = self._create_mock_provider({"availability_zone": "1,2,3"})
# Test different node configurations
test_cases = [
# Node with specific zone override
{
"config": {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "2",
}
},
"expected_zones": ["2"],
"expected_source": "node config availability_zone",
},
# Node with disabled zones
{
"config": {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "none",
}
},
"expected_zones": None,
"expected_source": "node config availability_zone",
},
# Node with auto-selection
{
"config": {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "",
}
},
"expected_zones": [],
"expected_source": "node config availability_zone",
},
# Node using provider default
{
"config": {"azure_arm_parameters": {"vmSize": "Standard_D2s_v3"}},
"expected_zones": ["1", "2", "3"],
"expected_source": "provider availability_zone",
},
]
for i, test_case in enumerate(test_cases):
with self.subTest(case=i):
zones, source = self._extract_zone_logic(provider, test_case["config"])
self.assertEqual(zones, test_case["expected_zones"])
self.assertEqual(source, test_case["expected_source"])
def test_mixed_case_precedence(self):
"""Test precedence with mixed case 'none' values."""
provider = self._create_mock_provider({"availability_zone": "None"})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": "NONE",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
# Both should be None (disabled), but node should take precedence
self.assertIsNone(zones)
self.assertEqual(source, "node config availability_zone")
def test_whitespace_handling_in_precedence(self):
"""Test that whitespace is properly handled in precedence logic."""
provider = self._create_mock_provider({"availability_zone": " 1, 2, 3 "})
node_config = {
"azure_arm_parameters": {
"vmSize": "Standard_D2s_v3",
"availability_zone": " 2 ",
}
}
zones, source = self._extract_zone_logic(provider, node_config)
self.assertEqual(zones, ["2"])
self.assertEqual(source, "node config availability_zone")
if __name__ == "__main__":
unittest.main()