401 lines
16 KiB
Python
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()
|