Coverage for mlos_bench/mlos_bench/tests/services/remote/azure/azure_network_services_test.py: 87%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Tests for mlos_bench.services.remote.azure.azure_network_services.""" 

6 

7from unittest.mock import MagicMock, patch 

8 

9import pytest 

10import requests.exceptions as requests_ex 

11 

12from mlos_bench.environments.status import Status 

13from mlos_bench.services.remote.azure.azure_auth import AzureAuthService 

14from mlos_bench.services.remote.azure.azure_network_services import AzureNetworkService 

15from mlos_bench.tests.services.remote.azure import make_httplib_json_response 

16 

17 

18@pytest.mark.parametrize( 

19 ("total_retries", "operation_status"), 

20 [ 

21 (2, Status.SUCCEEDED), 

22 (1, Status.FAILED), 

23 (0, Status.FAILED), 

24 ], 

25) 

26@patch("urllib3.connectionpool.HTTPConnectionPool._get_conn") 

27def test_wait_network_deployment_retry( 

28 mock_getconn: MagicMock, 

29 total_retries: int, 

30 operation_status: Status, 

31 azure_network_service: AzureNetworkService, 

32) -> None: 

33 """Test retries of the network deployment operation.""" 

34 # Simulate intermittent connection issues with multiple connection errors 

35 # Sufficient retry attempts should result in success, otherwise a graceful failure state 

36 mock_getconn.return_value.getresponse.side_effect = [ 

37 make_httplib_json_response(200, {"properties": {"provisioningState": "Running"}}), 

38 requests_ex.ConnectionError( 

39 "Connection aborted", OSError(107, "Transport endpoint is not connected") 

40 ), 

41 requests_ex.ConnectionError( 

42 "Connection aborted", OSError(107, "Transport endpoint is not connected") 

43 ), 

44 make_httplib_json_response(200, {"properties": {"provisioningState": "Running"}}), 

45 make_httplib_json_response(200, {"properties": {"provisioningState": "Succeeded"}}), 

46 ] 

47 

48 (status, _) = azure_network_service.wait_network_deployment( 

49 params={ 

50 "pollInterval": 0.1, 

51 "requestTotalRetries": total_retries, 

52 "deploymentName": "TEST_DEPLOYMENT1", 

53 "subscription": "TEST_SUB1", 

54 "resourceGroup": "TEST_RG1", 

55 }, 

56 is_setup=True, 

57 ) 

58 assert status == operation_status 

59 

60 

61@pytest.mark.parametrize( 

62 ("operation_name", "accepts_params"), 

63 [ 

64 ("deprovision_network", True), 

65 ], 

66) 

67@pytest.mark.parametrize( 

68 ("http_status_code", "operation_status"), 

69 [ 

70 (200, Status.SUCCEEDED), 

71 (202, Status.PENDING), 

72 # These should succeed since we set ignore_errors=True by default 

73 (401, Status.SUCCEEDED), 

74 (404, Status.SUCCEEDED), 

75 ], 

76) 

77@patch("mlos_bench.services.remote.azure.azure_deployment_services.requests") 

78# pylint: disable=too-many-arguments 

79def test_network_operation_status( 

80 mock_requests: MagicMock, 

81 azure_network_service: AzureNetworkService, 

82 operation_name: str, 

83 accepts_params: bool, 

84 http_status_code: int, 

85 operation_status: Status, 

86) -> None: 

87 """Test network operation status.""" 

88 mock_response = MagicMock() 

89 mock_response.status_code = http_status_code 

90 mock_requests.post.return_value = mock_response 

91 

92 operation = getattr(azure_network_service, operation_name) 

93 with pytest.raises(ValueError): 

94 # Missing vnetName should raise ValueError 

95 (status, _) = operation({}) if accepts_params else operation() 

96 (status, _) = operation({"vnetName": "test-vnet"}) if accepts_params else operation() 

97 assert status == operation_status 

98 

99 

100@pytest.fixture 

101def test_azure_network_service_no_deployment_template( 

102 azure_auth_service: AzureAuthService, 

103) -> None: 

104 """Tests creating a network services without a deployment template (should fail).""" 

105 with pytest.raises(ValueError): 

106 _ = AzureNetworkService( 

107 config={ 

108 "deploymentTemplatePath": None, 

109 "deploymentTemplateParameters": { 

110 "location": "westus2", 

111 }, 

112 }, 

113 parent=azure_auth_service, 

114 ) 

115 

116 with pytest.raises(ValueError): 

117 _ = AzureNetworkService( 

118 config={ 

119 # "deploymentTemplatePath": None, 

120 "deploymentTemplateParameters": { 

121 "location": "westus2", 

122 }, 

123 }, 

124 parent=azure_auth_service, 

125 )