python-sdk/tests/test_retry_policy.py at e1c9fa3825f5417464e835e88abfae2ea998f877 · yandex-cloud/python-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from concurrent import futures
from unittest.mock import Mock, patch

import grpc
import pytest

from yandex.cloud.vpc.v1.network_pb2 import Network
from yandex.cloud.vpc.v1.network_service_pb2 import GetNetworkRequest
from yandex.cloud.vpc.v1.network_service_pb2_grpc import (
    NetworkServiceStub,
    add_NetworkServiceServicer_to_server,
)
from yandexcloud import SDK, RetryPolicy
from yandexcloud._channels import Channels

INSECURE_SERVICE_PORT = "50051"
SERVICE_ADDR = "localhost"


def side_effect_internal(_, context):
    context.set_code(grpc.StatusCode.INTERNAL)


def side_effect_unavailable(_, context):
    context.set_code(grpc.StatusCode.UNAVAILABLE)


class VPCServiceMock:
    def __init__(self, fn):
        self.Get = Mock(return_value=Network(id="12342314"))
        self.Create = Mock()
        self.Update = Mock()
        self.Delete = Mock()
        self.ListSubnets = Mock()
        self.ListSecurityGroups = Mock()
        self.ListRouteTables = Mock()
        self.ListOperations = Mock()
        self.Move = Mock()
        self.List = Mock()


@pytest.fixture
def mock_channel():
    with patch.multiple(
        Channels,
        _get_creds=lambda self, endpoint: grpc.local_channel_credentials(),
        _get_endpoints=lambda self: {
            "vpc": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
            "iam": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
        },
    ) as channel_patch:
        yield channel_patch


def grpc_server(side_effect):
    service = VPCServiceMock(side_effect)
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    server.add_insecure_port("localhost:" + INSECURE_SERVICE_PORT)
    add_NetworkServiceServicer_to_server(service, server)
    server.start()
    return server, service


def test_default_retries(mock_channel):
    server, service = grpc_server(side_effect_unavailable)

    sdk = SDK(
        retry_policy=RetryPolicy(),
        endpoint=f"localhost:{INSECURE_SERVICE_PORT}",
        endpoints={
            "vpc": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
            "iam": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
        },
    )
    network_client = sdk.client(NetworkServiceStub, insecure=True)
    try:
        request = GetNetworkRequest(network_id="asdf")
        network_client.Get(request)
    except grpc.RpcError:
        assert service.Get.call_count == 4

    server.stop(0)


def test_custom_retries(mock_channel):
    server, service = grpc_server(side_effect_internal)

    sdk = SDK(
        retry_policy=RetryPolicy(status_codes=(grpc.StatusCode.INTERNAL,), max_attempts=4),
        endpoint=f"localhost:{INSECURE_SERVICE_PORT}",
        endpoints={
            "vpc": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
            "iam": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
        },
    )
    network_client = sdk.client(NetworkServiceStub, insecure=True)
    try:
        request = GetNetworkRequest(network_id="asdf")
        network_client.Get(request)
    except grpc.RpcError:
        assert service.Get.call_count == 4

    server.stop(0)


def test_no_retries(mock_channel):
    server, service = grpc_server(side_effect_internal)

    sdk = SDK(
        endpoint=f"localhost:{INSECURE_SERVICE_PORT}",
        endpoints={
            "vpc": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
            "iam": SERVICE_ADDR + ":" + INSECURE_SERVICE_PORT,
        },
    )
    network_client = sdk.client(NetworkServiceStub, insecure=True)
    try:
        request = GetNetworkRequest(network_id="asdf")
        network_client.Get(request)
    except grpc.RpcError:
        assert service.Get.call_count == 1

    server.stop(0)