def test_no_keys():
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
values = get_keys_values(request, context, [])
assert values == []
def test_ratelimit_reached_for_different_groups(clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
server = FakeGRPCServer()
with freeze_time(datetime.utcfromtimestamp(14)):
assert server.Foo(request, context)
assert server.Foo(request, context)
assert server.Foo(request, context)
with pytest.raises(Exception):
server.Foo(request, context)
assert context.abort_status == grpc.StatusCode.RESOURCE_EXHAUSTED
assert context.abort_message == (
"Reached limit of 3 calls per 10 seconds. "
"Resource will be available in 6 seconds.")
with freeze_time(datetime.utcfromtimestamp(17)):
assert server.Bar(request, context)
assert server.Bar(request, context)
with pytest.raises(Exception):
server.Bar(request, context)
assert context.abort_status == grpc.StatusCode.RESOURCE_EXHAUSTED
assert context.abort_message == (
"Reached limit of 2 calls per 5 seconds. "
"Resource will be available in 3 seconds.")
def test_get_request_field():
request = helloworld_pb2.HelloRequest(name="Aleksandr")
context = FakeServicerContext()
values = get_keys_values(request, context, ["request:name"])
assert values == ["Aleksandr"]
def test_call_method_under_limit(clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
server = FakeGRPCServer()
with freeze_time(datetime.utcfromtimestamp(14)):
assert server.Foo(request, context)
assert server.Foo(request, context)
assert server.Foo(request, context)
def test_get_metadata_value():
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
context.set_invocation_metadata((("user-agent", "Python 3.7 client"), ))
values = get_keys_values(request, context, ["metadata:user-agent"])
assert values == ["Python 3.7 client"]
def test_key_is_function():
request = helloworld_pb2.HelloRequest(name="Aleksandr")
context = FakeServicerContext()
values = get_keys_values(request, context,
[lambda request, context: request.name])
assert values == ["Aleksandr"]
def test_record_new_call(clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
with freeze_time(datetime.utcfromtimestamp(14)):
count, time_left = record_call(FakeGRPCServer.Foo,
request,
context,
time_period=10)
assert count == 1
assert time_left == 6
def test_record_same_group_but_different_keys(clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
with freeze_time(datetime.utcfromtimestamp(14)):
count, time_left = record_call(FakeGRPCServer.Foo,
request,
context,
time_period=10)
assert count == 1
assert time_left == 6
with freeze_time(datetime.utcfromtimestamp(16)):
count, time_left = record_call(
FakeGRPCServer.Bar,
helloworld_pb2.HelloRequest(name="Aleksandr"),
context,
time_period=10,
group="FakeGRPCServer.Foo",
keys=["request:name"],
)
assert count == 1
assert time_left == 4
def test_multiple_rpc_with_the_same_group_with_multiple_decorators(
clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
server = FakeGRPCServer()
with freeze_time(datetime.utcfromtimestamp(14)):
assert server.Foo(request, context)
assert server.Foo(request, context)
assert server.Foo(request, context)
with freeze_time(datetime.utcfromtimestamp(17)):
server.Baz(request, context)
server.Baz(request, context)
with pytest.raises(Exception):
server.Baz(request, context)
assert context.abort_status == grpc.StatusCode.RESOURCE_EXHAUSTED
assert context.abort_message == (
"Reached limit of 5 calls per 10 seconds. "
"Resource will be available in 3 seconds.")
def test_rpc_with_multiple_decorators(clear_cache):
request = helloworld_pb2.HelloRequest()
context = FakeServicerContext()
context.set_invocation_metadata((("user-agent", "Python 3.9 client"), ))
server = FakeGRPCServer()
with freeze_time(datetime.utcfromtimestamp(17)):
assert server.Baz(request, context)
assert server.Baz(request, context)
with pytest.raises(Exception):
server.Baz(request, context)
assert context.abort_status == grpc.StatusCode.RESOURCE_EXHAUSTED
assert context.abort_message == (
"Reached limit of 2 calls per 5 seconds. "
"Resource will be available in 3 seconds.")
context1 = FakeServicerContext()
context1.set_invocation_metadata((("user-agent", "Python 3.10 client"), ))
with freeze_time(datetime.utcfromtimestamp(17)):
assert server.Baz(request, context1)
assert server.Baz(request, context1)
def call_hello_method(addr, name):
with grpc.insecure_channel(addr) as channel:
stub = helloworld_pb2_grpc.GreeterStub(channel)
response, call = stub.SayHello.with_call(
helloworld_pb2.HelloRequest(name=name))
return response.message