wehub-resource-sync
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
import sys
|
|
from typing import List, Optional
|
|
|
|
import pytest
|
|
|
|
from ray.autoscaler._private.gcp.node_provider import _retry
|
|
|
|
|
|
class MockGCPNodeProvider:
|
|
def __init__(self, errors: List[type]):
|
|
# List off errors to raise while retrying self.mock_method
|
|
self.errors = errors
|
|
# Incremented during each retry via self._construct_client
|
|
self.error_index = -1
|
|
# Mirrors the __init__ of GCPNodeProvider
|
|
# Also called during each retry in _retry
|
|
self._construct_clients()
|
|
|
|
def _construct_clients(self):
|
|
# In real life, called during each retry to reinitializes api clients.
|
|
# Here, increments index in list of errors passed into test.
|
|
self.error_index += 1
|
|
|
|
@_retry
|
|
def mock_method(self, *args, **kwargs):
|
|
error = self.errors[self.error_index]
|
|
if error:
|
|
raise error
|
|
return (args, kwargs)
|
|
|
|
|
|
# Short names for two types of errors
|
|
B, V = BrokenPipeError, ValueError
|
|
# BrokenPipeError is supposed to caught with up to 5 tries.
|
|
# ValueError is an arbitrarily chosen exception which should not be caught.
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"error_input,expected_error_raised",
|
|
[
|
|
([None], None),
|
|
([B, B, B, B, None], None),
|
|
([B, B, V, B, None], V),
|
|
([B, B, B, B, B, None], B),
|
|
([B, B, B, B, B, B, None], B),
|
|
],
|
|
)
|
|
def test_gcp_broken_pipe_retry(
|
|
error_input: list, expected_error_raised: Optional[type]
|
|
):
|
|
"""Tests retries of BrokenPipeError in GCPNodeProvider.
|
|
|
|
Args:
|
|
error_input: List of exceptions hit during retries of test mock_method.
|
|
None means no exception.
|
|
expected_error_raised: Expected exception raised.
|
|
None means no exception.
|
|
"""
|
|
provider = MockGCPNodeProvider(error_input)
|
|
if expected_error_raised:
|
|
with pytest.raises(expected_error_raised):
|
|
provider.mock_method(1, 2, a=4, b=5)
|
|
else:
|
|
ret = provider.mock_method(1, 2, a=4, b=5)
|
|
assert ret == ((1, 2), {"a": 4, "b": 5})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|