dwilson-request-auth-client 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- dwilson_request_auth_client-0.1.0/PKG-INFO +56 -0
- dwilson_request_auth_client-0.1.0/dwilson_request_auth_client.egg-info/PKG-INFO +56 -0
- dwilson_request_auth_client-0.1.0/dwilson_request_auth_client.egg-info/SOURCES.txt +12 -0
- dwilson_request_auth_client-0.1.0/dwilson_request_auth_client.egg-info/dependency_links.txt +1 -0
- dwilson_request_auth_client-0.1.0/dwilson_request_auth_client.egg-info/requires.txt +2 -0
- dwilson_request_auth_client-0.1.0/dwilson_request_auth_client.egg-info/top_level.txt +2 -0
- dwilson_request_auth_client-0.1.0/proto/__init__.py +0 -0
- dwilson_request_auth_client-0.1.0/proto/permit_pb2.py +49 -0
- dwilson_request_auth_client-0.1.0/proto/permit_pb2_grpc.py +97 -0
- dwilson_request_auth_client-0.1.0/pyproject.toml +20 -0
- dwilson_request_auth_client-0.1.0/readme.md +47 -0
- dwilson_request_auth_client-0.1.0/request_auth_client.py +164 -0
- dwilson_request_auth_client-0.1.0/setup.cfg +4 -0
- dwilson_request_auth_client-0.1.0/tests/test_request_auth_client.py +137 -0
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dwilson-request-auth-client
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Python client for the Request Authorization gRPC service
|
|
5
|
+
Requires-Python: >=3.11
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: grpcio>=1.64.0
|
|
8
|
+
Requires-Dist: protobuf>=5.0.0
|
|
9
|
+
|
|
10
|
+
# request-auth-client
|
|
11
|
+
|
|
12
|
+
Python client for the Request Authorization gRPC permit service.
|
|
13
|
+
|
|
14
|
+
## Installation
|
|
15
|
+
|
|
16
|
+
Install from the local package directory:
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
pip install ./client
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
Or from PyPI:
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install dwilson-request-auth-client
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
**Requires:** Python 3.11+, `grpcio>=1.64.0`, `protobuf>=5.0.0`
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
## Quick start
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from request_auth_client import RequestAuthClient
|
|
36
|
+
|
|
37
|
+
client = RequestAuthClient("localhost:9000")
|
|
38
|
+
|
|
39
|
+
try:
|
|
40
|
+
with client.acquire("example.com") as permit:
|
|
41
|
+
# run request here
|
|
42
|
+
permit.set_status(200)
|
|
43
|
+
finally:
|
|
44
|
+
client.close()
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
### Explicit release
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
permit = client.acquire("example.com")
|
|
51
|
+
try:
|
|
52
|
+
# run request here
|
|
53
|
+
permit.release(200)
|
|
54
|
+
except Exception:
|
|
55
|
+
permit.release(0)
|
|
56
|
+
```
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dwilson-request-auth-client
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Python client for the Request Authorization gRPC service
|
|
5
|
+
Requires-Python: >=3.11
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: grpcio>=1.64.0
|
|
8
|
+
Requires-Dist: protobuf>=5.0.0
|
|
9
|
+
|
|
10
|
+
# request-auth-client
|
|
11
|
+
|
|
12
|
+
Python client for the Request Authorization gRPC permit service.
|
|
13
|
+
|
|
14
|
+
## Installation
|
|
15
|
+
|
|
16
|
+
Install from the local package directory:
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
pip install ./client
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
Or from PyPI:
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install dwilson-request-auth-client
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
**Requires:** Python 3.11+, `grpcio>=1.64.0`, `protobuf>=5.0.0`
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
## Quick start
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from request_auth_client import RequestAuthClient
|
|
36
|
+
|
|
37
|
+
client = RequestAuthClient("localhost:9000")
|
|
38
|
+
|
|
39
|
+
try:
|
|
40
|
+
with client.acquire("example.com") as permit:
|
|
41
|
+
# run request here
|
|
42
|
+
permit.set_status(200)
|
|
43
|
+
finally:
|
|
44
|
+
client.close()
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
### Explicit release
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
permit = client.acquire("example.com")
|
|
51
|
+
try:
|
|
52
|
+
# run request here
|
|
53
|
+
permit.release(200)
|
|
54
|
+
except Exception:
|
|
55
|
+
permit.release(0)
|
|
56
|
+
```
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
pyproject.toml
|
|
2
|
+
readme.md
|
|
3
|
+
request_auth_client.py
|
|
4
|
+
dwilson_request_auth_client.egg-info/PKG-INFO
|
|
5
|
+
dwilson_request_auth_client.egg-info/SOURCES.txt
|
|
6
|
+
dwilson_request_auth_client.egg-info/dependency_links.txt
|
|
7
|
+
dwilson_request_auth_client.egg-info/requires.txt
|
|
8
|
+
dwilson_request_auth_client.egg-info/top_level.txt
|
|
9
|
+
proto/__init__.py
|
|
10
|
+
proto/permit_pb2.py
|
|
11
|
+
proto/permit_pb2_grpc.py
|
|
12
|
+
tests/test_request_auth_client.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
File without changes
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
# -*- coding: utf-8 -*-
|
|
2
|
+
# Generated by the protocol buffer compiler. DO NOT EDIT!
|
|
3
|
+
# NO CHECKED-IN PROTOBUF GENCODE
|
|
4
|
+
# source: permit.proto
|
|
5
|
+
# Protobuf Python Version: 6.31.1
|
|
6
|
+
"""Generated protocol buffer code."""
|
|
7
|
+
from google.protobuf import descriptor as _descriptor
|
|
8
|
+
from google.protobuf import descriptor_pool as _descriptor_pool
|
|
9
|
+
from google.protobuf import runtime_version as _runtime_version
|
|
10
|
+
from google.protobuf import symbol_database as _symbol_database
|
|
11
|
+
from google.protobuf.internal import builder as _builder
|
|
12
|
+
_runtime_version.ValidateProtobufRuntimeVersion(
|
|
13
|
+
_runtime_version.Domain.PUBLIC,
|
|
14
|
+
6,
|
|
15
|
+
31,
|
|
16
|
+
1,
|
|
17
|
+
'',
|
|
18
|
+
'permit.proto'
|
|
19
|
+
)
|
|
20
|
+
# @@protoc_insertion_point(imports)
|
|
21
|
+
|
|
22
|
+
_sym_db = _symbol_database.Default()
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cpermit.proto\x12\x06permit\"i\n\rClientMessage\x12(\n\x07request\x18\x01 \x01(\x0b\x32\x15.permit.PermitRequestH\x00\x12#\n\x03ret\x18\x02 \x01(\x0b\x32\x14.permit.PermitReturnH\x00\x42\t\n\x07payload\"f\n\rServerMessage\x12$\n\x05grant\x18\x01 \x01(\x0b\x32\x13.permit.PermitGrantH\x00\x12$\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x13.permit.ServerErrorH\x00\x42\t\n\x07payload\"/\n\rPermitRequest\x12\x0e\n\x06\x64omain\x18\x01 \x01(\t\x12\x0e\n\x06req_id\x18\x02 \x01(\x03\"@\n\x0bPermitGrant\x12\x11\n\tpermit_id\x18\x01 \x01(\t\x12\x0e\n\x06req_id\x18\x02 \x01(\x03\x12\x0e\n\x06ttl_ms\x18\x03 \x01(\x05\"6\n\x0cPermitReturn\x12\x11\n\tpermit_id\x18\x01 \x01(\t\x12\x13\n\x0bstatus_code\x18\x02 \x01(\x05\".\n\x0bServerError\x12\x0e\n\x06req_id\x18\x01 \x01(\x03\x12\x0f\n\x07message\x18\x02 \x01(\t2Q\n\rPermitService\x12@\n\x0cPermitStream\x12\x15.permit.ClientMessage\x1a\x15.permit.ServerMessage(\x01\x30\x01\x42\'Z%github.com/dwilson/request-auth/protob\x06proto3')
|
|
28
|
+
|
|
29
|
+
_globals = globals()
|
|
30
|
+
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
|
|
31
|
+
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'permit_pb2', _globals)
|
|
32
|
+
if not _descriptor._USE_C_DESCRIPTORS:
|
|
33
|
+
_globals['DESCRIPTOR']._loaded_options = None
|
|
34
|
+
_globals['DESCRIPTOR']._serialized_options = b'Z%github.com/dwilson/request-auth/proto'
|
|
35
|
+
_globals['_CLIENTMESSAGE']._serialized_start=24
|
|
36
|
+
_globals['_CLIENTMESSAGE']._serialized_end=129
|
|
37
|
+
_globals['_SERVERMESSAGE']._serialized_start=131
|
|
38
|
+
_globals['_SERVERMESSAGE']._serialized_end=233
|
|
39
|
+
_globals['_PERMITREQUEST']._serialized_start=235
|
|
40
|
+
_globals['_PERMITREQUEST']._serialized_end=282
|
|
41
|
+
_globals['_PERMITGRANT']._serialized_start=284
|
|
42
|
+
_globals['_PERMITGRANT']._serialized_end=348
|
|
43
|
+
_globals['_PERMITRETURN']._serialized_start=350
|
|
44
|
+
_globals['_PERMITRETURN']._serialized_end=404
|
|
45
|
+
_globals['_SERVERERROR']._serialized_start=406
|
|
46
|
+
_globals['_SERVERERROR']._serialized_end=452
|
|
47
|
+
_globals['_PERMITSERVICE']._serialized_start=454
|
|
48
|
+
_globals['_PERMITSERVICE']._serialized_end=535
|
|
49
|
+
# @@protoc_insertion_point(module_scope)
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
|
|
2
|
+
"""Client and server classes corresponding to protobuf-defined services."""
|
|
3
|
+
import grpc
|
|
4
|
+
import warnings
|
|
5
|
+
|
|
6
|
+
import permit_pb2 as permit__pb2
|
|
7
|
+
|
|
8
|
+
GRPC_GENERATED_VERSION = '1.80.0'
|
|
9
|
+
GRPC_VERSION = grpc.__version__
|
|
10
|
+
_version_not_supported = False
|
|
11
|
+
|
|
12
|
+
try:
|
|
13
|
+
from grpc._utilities import first_version_is_lower
|
|
14
|
+
_version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
|
|
15
|
+
except ImportError:
|
|
16
|
+
_version_not_supported = True
|
|
17
|
+
|
|
18
|
+
if _version_not_supported:
|
|
19
|
+
raise RuntimeError(
|
|
20
|
+
f'The grpc package installed is at version {GRPC_VERSION},'
|
|
21
|
+
+ ' but the generated code in permit_pb2_grpc.py depends on'
|
|
22
|
+
+ f' grpcio>={GRPC_GENERATED_VERSION}.'
|
|
23
|
+
+ f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
|
|
24
|
+
+ f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class PermitServiceStub(object):
|
|
29
|
+
"""Missing associated documentation comment in .proto file."""
|
|
30
|
+
|
|
31
|
+
def __init__(self, channel):
|
|
32
|
+
"""Constructor.
|
|
33
|
+
|
|
34
|
+
Args:
|
|
35
|
+
channel: A grpc.Channel.
|
|
36
|
+
"""
|
|
37
|
+
self.PermitStream = channel.stream_stream(
|
|
38
|
+
'/permit.PermitService/PermitStream',
|
|
39
|
+
request_serializer=permit__pb2.ClientMessage.SerializeToString,
|
|
40
|
+
response_deserializer=permit__pb2.ServerMessage.FromString,
|
|
41
|
+
_registered_method=True)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class PermitServiceServicer(object):
|
|
45
|
+
"""Missing associated documentation comment in .proto file."""
|
|
46
|
+
|
|
47
|
+
def PermitStream(self, request_iterator, context):
|
|
48
|
+
"""Missing associated documentation comment in .proto file."""
|
|
49
|
+
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
|
50
|
+
context.set_details('Method not implemented!')
|
|
51
|
+
raise NotImplementedError('Method not implemented!')
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def add_PermitServiceServicer_to_server(servicer, server):
|
|
55
|
+
rpc_method_handlers = {
|
|
56
|
+
'PermitStream': grpc.stream_stream_rpc_method_handler(
|
|
57
|
+
servicer.PermitStream,
|
|
58
|
+
request_deserializer=permit__pb2.ClientMessage.FromString,
|
|
59
|
+
response_serializer=permit__pb2.ServerMessage.SerializeToString,
|
|
60
|
+
),
|
|
61
|
+
}
|
|
62
|
+
generic_handler = grpc.method_handlers_generic_handler(
|
|
63
|
+
'permit.PermitService', rpc_method_handlers)
|
|
64
|
+
server.add_generic_rpc_handlers((generic_handler,))
|
|
65
|
+
server.add_registered_method_handlers('permit.PermitService', rpc_method_handlers)
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
# This class is part of an EXPERIMENTAL API.
|
|
69
|
+
class PermitService(object):
|
|
70
|
+
"""Missing associated documentation comment in .proto file."""
|
|
71
|
+
|
|
72
|
+
@staticmethod
|
|
73
|
+
def PermitStream(request_iterator,
|
|
74
|
+
target,
|
|
75
|
+
options=(),
|
|
76
|
+
channel_credentials=None,
|
|
77
|
+
call_credentials=None,
|
|
78
|
+
insecure=False,
|
|
79
|
+
compression=None,
|
|
80
|
+
wait_for_ready=None,
|
|
81
|
+
timeout=None,
|
|
82
|
+
metadata=None):
|
|
83
|
+
return grpc.experimental.stream_stream(
|
|
84
|
+
request_iterator,
|
|
85
|
+
target,
|
|
86
|
+
'/permit.PermitService/PermitStream',
|
|
87
|
+
permit__pb2.ClientMessage.SerializeToString,
|
|
88
|
+
permit__pb2.ServerMessage.FromString,
|
|
89
|
+
options,
|
|
90
|
+
channel_credentials,
|
|
91
|
+
insecure,
|
|
92
|
+
call_credentials,
|
|
93
|
+
compression,
|
|
94
|
+
wait_for_ready,
|
|
95
|
+
timeout,
|
|
96
|
+
metadata,
|
|
97
|
+
_registered_method=True)
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "dwilson-request-auth-client"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Python client for the Request Authorization gRPC service"
|
|
9
|
+
readme = "readme.md"
|
|
10
|
+
requires-python = ">=3.11"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"grpcio>=1.64.0",
|
|
13
|
+
"protobuf>=5.0.0",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[tool.setuptools]
|
|
17
|
+
py-modules = ["request_auth_client"]
|
|
18
|
+
|
|
19
|
+
[tool.setuptools.packages.find]
|
|
20
|
+
include = ["proto", "proto.*"]
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
# request-auth-client
|
|
2
|
+
|
|
3
|
+
Python client for the Request Authorization gRPC permit service.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
Install from the local package directory:
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install ./client
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
Or from PyPI:
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install dwilson-request-auth-client
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
**Requires:** Python 3.11+, `grpcio>=1.64.0`, `protobuf>=5.0.0`
|
|
20
|
+
|
|
21
|
+
---
|
|
22
|
+
|
|
23
|
+
## Quick start
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
from request_auth_client import RequestAuthClient
|
|
27
|
+
|
|
28
|
+
client = RequestAuthClient("localhost:9000")
|
|
29
|
+
|
|
30
|
+
try:
|
|
31
|
+
with client.acquire("example.com") as permit:
|
|
32
|
+
# run request here
|
|
33
|
+
permit.set_status(200)
|
|
34
|
+
finally:
|
|
35
|
+
client.close()
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
### Explicit release
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
permit = client.acquire("example.com")
|
|
42
|
+
try:
|
|
43
|
+
# run request here
|
|
44
|
+
permit.release(200)
|
|
45
|
+
except Exception:
|
|
46
|
+
permit.release(0)
|
|
47
|
+
```
|
|
@@ -0,0 +1,164 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Sync Python client for the Request Authorization gRPC service.
|
|
3
|
+
|
|
4
|
+
Usage::
|
|
5
|
+
|
|
6
|
+
from request_auth_client import RequestAuthClient
|
|
7
|
+
|
|
8
|
+
client = RequestAuthClient("localhost:9000")
|
|
9
|
+
|
|
10
|
+
# Explicit release
|
|
11
|
+
permit = client.acquire("rockauto.com")
|
|
12
|
+
try:
|
|
13
|
+
response = requests.get(url)
|
|
14
|
+
permit.release(response.status_code)
|
|
15
|
+
except Exception:
|
|
16
|
+
permit.release(0)
|
|
17
|
+
|
|
18
|
+
# Context manager (preferred — auto-releases with status 0 on exception)
|
|
19
|
+
with client.acquire("rockauto.com") as permit:
|
|
20
|
+
response = requests.get(url)
|
|
21
|
+
permit.set_status(response.status_code)
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
import queue
|
|
25
|
+
import threading
|
|
26
|
+
from typing import Optional
|
|
27
|
+
|
|
28
|
+
import grpc
|
|
29
|
+
|
|
30
|
+
import sys
|
|
31
|
+
import os
|
|
32
|
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "proto"))
|
|
33
|
+
|
|
34
|
+
from proto import permit_pb2, permit_pb2_grpc # noqa: E402
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class RequestAuthClient:
|
|
38
|
+
"""Manages one gRPC bidirectional stream to the permit server.
|
|
39
|
+
|
|
40
|
+
Thread-safe. Multiple threads can call acquire() concurrently;
|
|
41
|
+
each blocks independently until the server issues a grant.
|
|
42
|
+
"""
|
|
43
|
+
|
|
44
|
+
def __init__(self, address: str, insecure: bool = True) -> None:
|
|
45
|
+
self._address = address
|
|
46
|
+
if insecure:
|
|
47
|
+
self._channel = grpc.insecure_channel(address)
|
|
48
|
+
else:
|
|
49
|
+
self._channel = grpc.secure_channel(address, grpc.ssl_channel_credentials())
|
|
50
|
+
self._stub = permit_pb2_grpc.PermitServiceStub(self._channel)
|
|
51
|
+
|
|
52
|
+
self._lock = threading.Lock()
|
|
53
|
+
self._req_id = 0
|
|
54
|
+
self._pending: dict[int, queue.Queue] = {}
|
|
55
|
+
self._send_q: queue.Queue = queue.Queue()
|
|
56
|
+
self._stopped = threading.Event()
|
|
57
|
+
self._connect()
|
|
58
|
+
|
|
59
|
+
# ------------------------------------------------------------------
|
|
60
|
+
# Public API
|
|
61
|
+
# ------------------------------------------------------------------
|
|
62
|
+
|
|
63
|
+
def acquire(self, domain: str) -> "Permit":
|
|
64
|
+
"""Block until the server grants a permit for domain."""
|
|
65
|
+
with self._lock:
|
|
66
|
+
self._req_id += 1
|
|
67
|
+
req_id = self._req_id
|
|
68
|
+
result_q: queue.Queue = queue.Queue()
|
|
69
|
+
self._pending[req_id] = result_q
|
|
70
|
+
|
|
71
|
+
self._send_q.put(permit_pb2.ClientMessage(
|
|
72
|
+
request=permit_pb2.PermitRequest(domain=domain, req_id=req_id)
|
|
73
|
+
))
|
|
74
|
+
|
|
75
|
+
kind, payload = result_q.get()
|
|
76
|
+
if kind == "error":
|
|
77
|
+
raise RuntimeError(f"permit acquisition failed for {domain!r}: {payload}")
|
|
78
|
+
return Permit(permit_id=payload.permit_id, client=self)
|
|
79
|
+
|
|
80
|
+
def close(self) -> None:
|
|
81
|
+
self._stopped.set()
|
|
82
|
+
self._send_q.put(None) # unblock the request iterator
|
|
83
|
+
self._channel.close()
|
|
84
|
+
|
|
85
|
+
# ------------------------------------------------------------------
|
|
86
|
+
# Internal
|
|
87
|
+
# ------------------------------------------------------------------
|
|
88
|
+
|
|
89
|
+
def _release(self, permit_id: str, status_code: int) -> None:
|
|
90
|
+
self._send_q.put(permit_pb2.ClientMessage(
|
|
91
|
+
ret=permit_pb2.PermitReturn(permit_id=permit_id, status_code=status_code)
|
|
92
|
+
))
|
|
93
|
+
|
|
94
|
+
def _connect(self) -> None:
|
|
95
|
+
self._stream = self._stub.PermitStream(self._request_iter())
|
|
96
|
+
t = threading.Thread(target=self._recv_loop, daemon=True)
|
|
97
|
+
t.start()
|
|
98
|
+
|
|
99
|
+
def _request_iter(self):
|
|
100
|
+
while not self._stopped.is_set():
|
|
101
|
+
try:
|
|
102
|
+
msg = self._send_q.get(timeout=1.0)
|
|
103
|
+
if msg is None:
|
|
104
|
+
return
|
|
105
|
+
yield msg
|
|
106
|
+
except queue.Empty:
|
|
107
|
+
continue
|
|
108
|
+
|
|
109
|
+
def _recv_loop(self) -> None:
|
|
110
|
+
try:
|
|
111
|
+
for msg in self._stream:
|
|
112
|
+
self._recv_loop_step(msg)
|
|
113
|
+
except grpc.RpcError as e:
|
|
114
|
+
self._wake_all_pending(str(e))
|
|
115
|
+
|
|
116
|
+
def _recv_loop_step(self, msg) -> None:
|
|
117
|
+
if msg.HasField("grant"):
|
|
118
|
+
with self._lock:
|
|
119
|
+
q = self._pending.pop(msg.grant.req_id, None)
|
|
120
|
+
if q:
|
|
121
|
+
q.put(("grant", msg.grant))
|
|
122
|
+
elif msg.HasField("error"):
|
|
123
|
+
with self._lock:
|
|
124
|
+
q = self._pending.pop(msg.error.req_id, None)
|
|
125
|
+
if q:
|
|
126
|
+
q.put(("error", msg.error.message))
|
|
127
|
+
|
|
128
|
+
def _wake_all_pending(self, error_msg: str) -> None:
|
|
129
|
+
with self._lock:
|
|
130
|
+
pending = dict(self._pending)
|
|
131
|
+
self._pending.clear()
|
|
132
|
+
for q in pending.values():
|
|
133
|
+
q.put(("error", error_msg))
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
class Permit:
|
|
137
|
+
"""Handle returned by RequestAuthClient.acquire().
|
|
138
|
+
|
|
139
|
+
Call release(status_code) when done, or use as a context manager.
|
|
140
|
+
"""
|
|
141
|
+
|
|
142
|
+
def __init__(self, permit_id: str, client: RequestAuthClient) -> None:
|
|
143
|
+
self._permit_id = permit_id
|
|
144
|
+
self._client = client
|
|
145
|
+
self._status_code: int = 0
|
|
146
|
+
self._released = False
|
|
147
|
+
|
|
148
|
+
def set_status(self, status_code: int) -> None:
|
|
149
|
+
"""Set the HTTP status code to report when the permit is released."""
|
|
150
|
+
self._status_code = status_code
|
|
151
|
+
|
|
152
|
+
def release(self, status_code: Optional[int] = None) -> None:
|
|
153
|
+
"""Release the permit back to the server."""
|
|
154
|
+
if self._released:
|
|
155
|
+
return
|
|
156
|
+
self._released = True
|
|
157
|
+
code = status_code if status_code is not None else self._status_code
|
|
158
|
+
self._client._release(self._permit_id, code)
|
|
159
|
+
|
|
160
|
+
def __enter__(self) -> "Permit":
|
|
161
|
+
return self
|
|
162
|
+
|
|
163
|
+
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
164
|
+
self.release()
|
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
"""Unit tests for RequestAuthClient using a mock gRPC stream."""
|
|
2
|
+
import queue
|
|
3
|
+
import threading
|
|
4
|
+
import unittest
|
|
5
|
+
from unittest.mock import MagicMock, patch
|
|
6
|
+
|
|
7
|
+
import sys, os
|
|
8
|
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
9
|
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "proto"))
|
|
10
|
+
|
|
11
|
+
from proto import permit_pb2
|
|
12
|
+
from request_auth_client import Permit, RequestAuthClient
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
# ---------------------------------------------------------------------------
|
|
16
|
+
# Helpers
|
|
17
|
+
# ---------------------------------------------------------------------------
|
|
18
|
+
|
|
19
|
+
def make_grant(req_id: int, permit_id: str = "test-permit") -> permit_pb2.ServerMessage:
|
|
20
|
+
return permit_pb2.ServerMessage(
|
|
21
|
+
grant=permit_pb2.PermitGrant(permit_id=permit_id, req_id=req_id, ttl_ms=1000)
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def make_error(req_id: int, msg: str = "boom") -> permit_pb2.ServerMessage:
|
|
26
|
+
return permit_pb2.ServerMessage(
|
|
27
|
+
error=permit_pb2.ServerError(req_id=req_id, message=msg)
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class FakeStream:
|
|
32
|
+
"""Simulates the server-side of the bidirectional stream."""
|
|
33
|
+
|
|
34
|
+
def __init__(self, responses: list):
|
|
35
|
+
self._responses = iter(responses)
|
|
36
|
+
self.sent: list[permit_pb2.ClientMessage] = []
|
|
37
|
+
|
|
38
|
+
def __iter__(self):
|
|
39
|
+
return self._responses
|
|
40
|
+
|
|
41
|
+
def write(self, msg):
|
|
42
|
+
self.sent.append(msg)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
# ---------------------------------------------------------------------------
|
|
46
|
+
# Tests
|
|
47
|
+
# ---------------------------------------------------------------------------
|
|
48
|
+
|
|
49
|
+
class TestPermit(unittest.TestCase):
|
|
50
|
+
def _make_permit(self, permit_id="abc"):
|
|
51
|
+
client = MagicMock()
|
|
52
|
+
return Permit(permit_id=permit_id, client=client), client
|
|
53
|
+
|
|
54
|
+
def test_release_calls_client(self):
|
|
55
|
+
permit, client = self._make_permit()
|
|
56
|
+
permit.release(200)
|
|
57
|
+
client._release.assert_called_once_with("abc", 200)
|
|
58
|
+
|
|
59
|
+
def test_release_only_once(self):
|
|
60
|
+
permit, client = self._make_permit()
|
|
61
|
+
permit.release(200)
|
|
62
|
+
permit.release(200)
|
|
63
|
+
client._release.assert_called_once()
|
|
64
|
+
|
|
65
|
+
def test_set_status_used_on_exit(self):
|
|
66
|
+
permit, client = self._make_permit()
|
|
67
|
+
permit.set_status(404)
|
|
68
|
+
permit.__exit__(None, None, None)
|
|
69
|
+
client._release.assert_called_once_with("abc", 404)
|
|
70
|
+
|
|
71
|
+
def test_context_manager_releases_on_exception(self):
|
|
72
|
+
permit, client = self._make_permit()
|
|
73
|
+
try:
|
|
74
|
+
with permit:
|
|
75
|
+
raise ValueError("oops")
|
|
76
|
+
except ValueError:
|
|
77
|
+
pass
|
|
78
|
+
client._release.assert_called_once_with("abc", 0)
|
|
79
|
+
|
|
80
|
+
def test_context_manager_releases_status(self):
|
|
81
|
+
permit, client = self._make_permit()
|
|
82
|
+
with permit:
|
|
83
|
+
permit.set_status(200)
|
|
84
|
+
client._release.assert_called_once_with("abc", 200)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class TestRequestAuthClientRecvLoop(unittest.TestCase):
|
|
88
|
+
"""Test _recv_loop by feeding messages directly into _pending."""
|
|
89
|
+
|
|
90
|
+
def _make_client_no_grpc(self):
|
|
91
|
+
"""Create a client without starting a real gRPC connection."""
|
|
92
|
+
with patch("grpc.insecure_channel"), \
|
|
93
|
+
patch.object(RequestAuthClient, "_connect"):
|
|
94
|
+
client = RequestAuthClient("localhost:9000")
|
|
95
|
+
return client
|
|
96
|
+
|
|
97
|
+
def test_grant_unblocks_pending(self):
|
|
98
|
+
client = self._make_client_no_grpc()
|
|
99
|
+
result_q: queue.Queue = queue.Queue()
|
|
100
|
+
client._pending[1] = result_q
|
|
101
|
+
|
|
102
|
+
grant_msg = make_grant(req_id=1, permit_id="permit-xyz")
|
|
103
|
+
client._recv_loop_step(grant_msg)
|
|
104
|
+
|
|
105
|
+
kind, payload = result_q.get_nowait()
|
|
106
|
+
self.assertEqual(kind, "grant")
|
|
107
|
+
self.assertEqual(payload.permit_id, "permit-xyz")
|
|
108
|
+
|
|
109
|
+
def test_error_unblocks_pending(self):
|
|
110
|
+
client = self._make_client_no_grpc()
|
|
111
|
+
result_q: queue.Queue = queue.Queue()
|
|
112
|
+
client._pending[2] = result_q
|
|
113
|
+
|
|
114
|
+
err_msg = make_error(req_id=2, msg="unknown domain")
|
|
115
|
+
client._recv_loop_step(err_msg)
|
|
116
|
+
|
|
117
|
+
kind, payload = result_q.get_nowait()
|
|
118
|
+
self.assertEqual(kind, "error")
|
|
119
|
+
self.assertIn("unknown domain", payload)
|
|
120
|
+
|
|
121
|
+
def test_wake_all_pending_on_disconnect(self):
|
|
122
|
+
client = self._make_client_no_grpc()
|
|
123
|
+
q1: queue.Queue = queue.Queue()
|
|
124
|
+
q2: queue.Queue = queue.Queue()
|
|
125
|
+
client._pending[1] = q1
|
|
126
|
+
client._pending[2] = q2
|
|
127
|
+
|
|
128
|
+
client._wake_all_pending("stream disconnected")
|
|
129
|
+
|
|
130
|
+
for q in (q1, q2):
|
|
131
|
+
kind, _ = q.get_nowait()
|
|
132
|
+
self.assertEqual(kind, "error")
|
|
133
|
+
self.assertEqual(len(client._pending), 0)
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
if __name__ == "__main__":
|
|
137
|
+
unittest.main()
|