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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Tests for mlos_bench.services.remote.azure.azure_network_services 

7""" 

8 

9from unittest.mock import MagicMock, patch 

10 

11import pytest 

12import requests.exceptions as requests_ex 

13 

14from mlos_bench.environments.status import Status 

15 

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

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

18 

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

20 

21 

22@pytest.mark.parametrize( 

23 ("total_retries", "operation_status"), [ 

24 (2, Status.SUCCEEDED), 

25 (1, Status.FAILED), 

26 (0, Status.FAILED), 

27 ]) 

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

29def test_wait_network_deployment_retry(mock_getconn: MagicMock, 

30 total_retries: int, 

31 operation_status: Status, 

32 azure_network_service: AzureNetworkService) -> None: 

33 """ 

34 Test retries of the network deployment operation. 

35 """ 

36 # Simulate intermittent connection issues with multiple connection errors 

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

38 mock_getconn.return_value.getresponse.side_effect = [ 

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

40 requests_ex.ConnectionError("Connection aborted", OSError(107, "Transport endpoint is not connected")), 

41 requests_ex.ConnectionError("Connection aborted", OSError(107, "Transport endpoint is not connected")), 

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

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

44 ] 

45 

46 (status, _) = azure_network_service.wait_network_deployment( 

47 params={ 

48 "pollInterval": 0.1, 

49 "requestTotalRetries": total_retries, 

50 "deploymentName": "TEST_DEPLOYMENT1", 

51 "subscription": "TEST_SUB1", 

52 "resourceGroup": "TEST_RG1", 

53 }, 

54 is_setup=True) 

55 assert status == operation_status 

56 

57 

58@pytest.mark.parametrize( 

59 ("operation_name", "accepts_params"), [ 

60 ("deprovision_network", True), 

61 ]) 

62@pytest.mark.parametrize( 

63 ("http_status_code", "operation_status"), [ 

64 (200, Status.SUCCEEDED), 

65 (202, Status.PENDING), 

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

67 (401, Status.SUCCEEDED), 

68 (404, Status.SUCCEEDED), 

69 ]) 

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

71# pylint: disable=too-many-arguments 

72def test_network_operation_status(mock_requests: MagicMock, 

73 azure_network_service: AzureNetworkService, 

74 operation_name: str, 

75 accepts_params: bool, 

76 http_status_code: int, 

77 operation_status: Status) -> None: 

78 """ 

79 Test network operation status. 

80 """ 

81 mock_response = MagicMock() 

82 mock_response.status_code = http_status_code 

83 mock_requests.post.return_value = mock_response 

84 

85 operation = getattr(azure_network_service, operation_name) 

86 with pytest.raises(ValueError): 

87 # Missing vnetName should raise ValueError 

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

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

90 assert status == operation_status 

91 

92 

93@pytest.fixture 

94def test_azure_network_service_no_deployment_template(azure_auth_service: AzureAuthService) -> None: 

95 """ 

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

97 """ 

98 with pytest.raises(ValueError): 

99 _ = AzureNetworkService(config={ 

100 "deploymentTemplatePath": None, 

101 "deploymentTemplateParameters": { 

102 "location": "westus2", 

103 }, 

104 }, parent=azure_auth_service) 

105 

106 with pytest.raises(ValueError): 

107 _ = AzureNetworkService(config={ 

108 # "deploymentTemplatePath": None, 

109 "deploymentTemplateParameters": { 

110 "location": "westus2", 

111 }, 

112 }, parent=azure_auth_service)