def test_retry_policy_does_not_mark_null_locations_unavailable(self):
self.original_get_database_account = cosmos_client_connection.CosmosClientConnection.GetDatabaseAccount
cosmos_client_connection.CosmosClientConnection.GetDatabaseAccount = self.mock_get_database_account
client = cosmos_client_connection.CosmosClientConnection(
self.DEFAULT_ENDPOINT, {'masterKey': self.MASTER_KEY}, None,
documents.ConsistencyLevel.Eventual)
endpoint_manager = global_endpoint_manager._GlobalEndpointManager(
client)
self.original_mark_endpoint_unavailable_for_read_function = endpoint_manager.mark_endpoint_unavailable_for_read
endpoint_manager.mark_endpoint_unavailable_for_read = self._mock_mark_endpoint_unavailable_for_read
self.original_mark_endpoint_unavailable_for_write_function = endpoint_manager.mark_endpoint_unavailable_for_write
endpoint_manager.mark_endpoint_unavailable_for_write = self._mock_mark_endpoint_unavailable_for_write
self.original_resolve_service_endpoint = endpoint_manager.resolve_service_endpoint
endpoint_manager.resolve_service_endpoint = self._mock_resolve_service_endpoint
# Read and write counters count the number of times the endpoint manager's
# mark_endpoint_unavailable_for_read() and mark_endpoint_unavailable_for_read()
# functions were called. When a 'None' location is returned by resolve_service_endpoint(),
# these functions should not be called
self._read_counter = 0
self._write_counter = 0
request = RequestObject(http_constants.ResourceType.Document,
documents._OperationType.Read)
endpointDiscovery_retry_policy = _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy(
documents.ConnectionPolicy(), endpoint_manager, request)
endpointDiscovery_retry_policy.ShouldRetry(
errors.CosmosHttpResponseError(
status_code=http_constants.StatusCodes.FORBIDDEN))
self.assertEqual(self._read_counter, 0)
self.assertEqual(self._write_counter, 0)
self._read_counter = 0
self._write_counter = 0
request = RequestObject(http_constants.ResourceType.Document,
documents._OperationType.Create)
endpointDiscovery_retry_policy = _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy(
documents.ConnectionPolicy(), endpoint_manager, request)
endpointDiscovery_retry_policy.ShouldRetry(
errors.CosmosHttpResponseError(
status_code=http_constants.StatusCodes.FORBIDDEN))
self.assertEqual(self._read_counter, 0)
self.assertEqual(self._write_counter, 0)
endpoint_manager.mark_endpoint_unavailable_for_read = self.original_mark_endpoint_unavailable_for_read_function
endpoint_manager.mark_endpoint_unavailable_for_write = self.original_mark_endpoint_unavailable_for_write_function
cosmos_client_connection.CosmosClientConnection.GetDatabaseAccount = self.original_get_database_account
def _MockExecuteFunctionSessionReadFailureOnce(self, function, *args,
**kwargs):
response = test_config.FakeResponse(
{HttpHeaders.SubStatus: SubStatusCodes.READ_SESSION_NOTAVAILABLE})
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.NOT_FOUND,
message="Read Session not available",
response=response)
def __call__(self, func, *args, **kwargs):
self.counter = self.counter + 1
if self.counter % 3 == 0:
return self.org_func(func, *args, **kwargs)
else:
raise errors.CosmosHttpResponseError(
status_code=10054,
message="Connection was reset",
response=test_config.FakeResponse({}))
def _MockExecuteFunction(self, function, *args, **kwargs):
response = test_config.FakeResponse({
HttpHeaders.RetryAfterInMilliseconds:
self.retry_after_in_milliseconds
})
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.TOO_MANY_REQUESTS,
message="Request rate is too large",
response=response)
def MockExecuteFunction(self, function, *args, **kwargs):
global location_changed
if self.endpoint_discovery_retry_count == 2:
_retry_utility.ExecuteFunction = self.OriginalExecuteFunction
return (json.dumps([{ 'id': 'mock database' }]), None)
else:
self.endpoint_discovery_retry_count += 1
location_changed = True
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.FORBIDDEN,
message="Forbidden",
response=test_config.FakeResponse({'x-ms-substatus' : 3}))
def _MockExecuteFunctionEndpointDiscover(self, function, *args, **kwargs):
self.counter += 1
if self.counter >= 10 or (len(args) > 0 and args[1].operation_type
== documents._OperationType.Read):
return ({}, {})
else:
self.endpoint_sequence.append(args[1].location_endpoint_to_route)
response = test_config.FakeResponse(
{HttpHeaders.SubStatus: SubStatusCodes.WRITE_FORBIDDEN})
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.FORBIDDEN,
message="Request is not permitted in this region",
response=response)
def _MockExecuteFunctionSessionReadFailureTwice(self, function, *args,
**kwargs):
request = args[1]
if self.counter == 0:
if not self.use_multiple_write_locations:
expected_endpoint = self.database_account.WritableLocations[0][
'databaseAccountEndpoint'] if self.is_preferred_locations_list_empty else self.preferred_locations[
0]
else:
expected_endpoint = self.endpoint_by_location[
self.preferred_locations[0]]
self.assertFalse(
request.should_clear_session_token_on_session_read_failure)
elif self.counter == 1:
expected_endpoint = self.database_account.WritableLocations[0][
'databaseAccountEndpoint']
if not self.use_multiple_write_locations:
self.assertTrue(
request.should_clear_session_token_on_session_read_failure)
else:
self.assertFalse(
request.should_clear_session_token_on_session_read_failure)
elif self.counter == 2:
expected_endpoint = self.endpoint_by_location[
self.preferred_locations[1]]
self.assertFalse(
request.should_clear_session_token_on_session_read_failure)
elif self.counter == 3:
expected_endpoint = self.database_account.ReadableLocations[2][
'databaseAccountEndpoint']
self.assertTrue(
request.should_clear_session_token_on_session_read_failure)
self.assertEqual(expected_endpoint, request.location_endpoint_to_route)
self.counter += 1
response = test_config.FakeResponse(
{HttpHeaders.SubStatus: SubStatusCodes.READ_SESSION_NOTAVAILABLE})
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.NOT_FOUND,
message="Read Session not available",
response=response)
def MockGetDatabaseAccountStub(self, endpoint):
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.SERVICE_UNAVAILABLE, message="Service unavailable")
def _MockExecuteFunction(self, function, *args, **kwargs):
response = test_config.FakeResponse({'x-ms-substatus' : SubStatusCodes.WRITE_FORBIDDEN})
raise errors.CosmosHttpResponseError(
status_code=StatusCodes.FORBIDDEN,
message="Write Forbidden",
response=response)