Add HW 2
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
FROM python:3.12-alpine
|
||||
|
||||
COPY tests/requirements.txt .
|
||||
RUN --mount=type=cache,id=distsys-course-pip,target=/root/.cache/pip,sharing=locked \
|
||||
pip install -r requirements.txt
|
||||
|
||||
ARG GRPCURL_VERSION=1.9.3
|
||||
ARG GRPCURL_SHA256=a926b62a85787ccf73ef8736b3ae554f1242e39d92bb8767a79d6dd23b11d1d5
|
||||
RUN apk add --no-cache curl tini \
|
||||
&& curl --fail --silent --show-error --location \
|
||||
"https://github.com/fullstorydev/grpcurl/releases/download/v${GRPCURL_VERSION}/grpcurl_${GRPCURL_VERSION}_linux_x86_64.tar.gz" \
|
||||
--output /tmp/grpcurl.tar.gz \
|
||||
&& echo "${GRPCURL_SHA256} /tmp/grpcurl.tar.gz" | sha256sum -c \
|
||||
&& tar -xzf /tmp/grpcurl.tar.gz -C /usr/local/bin \
|
||||
&& rm -f /tmp/grpcurl.tar.gz
|
||||
|
||||
COPY solution/proto/ solution/proto/
|
||||
COPY tests/*.py tests/
|
||||
|
||||
ENTRYPOINT ["/sbin/tini", "--", "python3", "-u", "tests/main.py"]
|
||||
@@ -0,0 +1,109 @@
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
SCRIPT_DIR = pathlib.Path(__file__).parent.resolve()
|
||||
SUITE_TIMEOUT_S = 120
|
||||
|
||||
|
||||
class PassedCounter:
|
||||
def __init__(self):
|
||||
self.passed = 0
|
||||
|
||||
def pytest_report_teststatus(self, report, config):
|
||||
if report.when == 'call' and report.passed:
|
||||
self.passed += 1
|
||||
|
||||
|
||||
def suite_score(passed, expected, maximum, exit_code=pytest.ExitCode.OK):
|
||||
return maximum if exit_code == pytest.ExitCode.OK and passed == expected else 0
|
||||
|
||||
|
||||
def run_suite(filename, expected, maximum):
|
||||
# Each component has its own deadline so a timeout cannot discard other scores.
|
||||
process = subprocess.Popen(
|
||||
[sys.executable, '-u', str(pathlib.Path(__file__).resolve()),
|
||||
'--suite', str(SCRIPT_DIR / filename), str(expected)],
|
||||
start_new_session=os.name == 'posix',
|
||||
)
|
||||
try:
|
||||
exit_code = process.wait(timeout=SUITE_TIMEOUT_S)
|
||||
except subprocess.TimeoutExpired:
|
||||
if os.name == 'posix':
|
||||
# Include grpcurl and other subprocesses started by this test suite.
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass # The suite may have exited just after wait() timed out.
|
||||
else:
|
||||
process.kill()
|
||||
process.wait()
|
||||
print(f'{filename} exceeded {SUITE_TIMEOUT_S} seconds; component score is zero.', flush=True)
|
||||
return 0
|
||||
return maximum if exit_code == 0 else 0
|
||||
|
||||
|
||||
def suite_exit_code(filename, expected):
|
||||
counter = PassedCounter()
|
||||
exit_code = pytest.main(['-vs', filename], plugins=[counter])
|
||||
return 0 if suite_score(counter.passed, expected, 1, exit_code) else 1
|
||||
|
||||
|
||||
def component_enabled(environment_name):
|
||||
return os.environ.get(environment_name, '1') == '1'
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--component',
|
||||
choices=('all', 'proto', 'server', 'client'),
|
||||
default='all',
|
||||
)
|
||||
parser.add_argument('--suite', nargs=2, metavar=('FILE', 'EXPECTED'), help=argparse.SUPPRESS)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
if args.suite is not None:
|
||||
return suite_exit_code(args.suite[0], int(args.suite[1]))
|
||||
component = args.component
|
||||
score = 0
|
||||
|
||||
if component in ('all', 'proto'):
|
||||
proto_score = run_suite('test_proto.py', expected=1, maximum=2)
|
||||
score += proto_score
|
||||
print(f'Proto: {proto_score}/2')
|
||||
print()
|
||||
|
||||
if component in ('all', 'server'):
|
||||
if component_enabled('MESSENGER_SERVER_TESTS_ENABLED'):
|
||||
server_score = run_suite('test_server.py', expected=4, maximum=4)
|
||||
else:
|
||||
server_score = 0
|
||||
print('Server tests were not run because the server image did not build.')
|
||||
score += server_score
|
||||
print(f'Server: {server_score}/4')
|
||||
print()
|
||||
|
||||
if component in ('all', 'client'):
|
||||
if component_enabled('MESSENGER_CLIENT_TESTS_ENABLED'):
|
||||
client_score = run_suite('test_client.py', expected=3, maximum=4)
|
||||
else:
|
||||
client_score = 0
|
||||
print('Client tests were not run because the client image did not build.')
|
||||
score += client_score
|
||||
print(f'Client: {client_score}/4')
|
||||
|
||||
print(f'\nSCORE: {score}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,5 @@
|
||||
grpcio==1.75.0
|
||||
grpcio-tools==1.75.0
|
||||
protobuf==6.32.1
|
||||
pytest==8.4.2
|
||||
requests==2.32.5
|
||||
@@ -0,0 +1,179 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
HTTP_CONNECT_TIMEOUT_S = 1
|
||||
HTTP_READ_TIMEOUT_S = 5
|
||||
HTTP_READY_TIMEOUT_S = 20
|
||||
HTTP_RETRY_INTERVAL_S = 0.5
|
||||
HTTP_TIMEOUT = (HTTP_CONNECT_TIMEOUT_S, HTTP_READ_TIMEOUT_S)
|
||||
MESSAGE_TIMEOUT_S = 10
|
||||
MESSAGE_POLL_INTERVAL_S = 0.05
|
||||
|
||||
|
||||
def wait_for_http(url):
|
||||
deadline = time.monotonic() + HTTP_READY_TIMEOUT_S
|
||||
last_exception = None
|
||||
while True:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
pytest.fail(
|
||||
f'timed out waiting for HTTP endpoint {url}: {last_exception}',
|
||||
pytrace=False,
|
||||
)
|
||||
readiness_timeout = (
|
||||
min(HTTP_CONNECT_TIMEOUT_S, remaining),
|
||||
min(HTTP_READ_TIMEOUT_S, remaining),
|
||||
)
|
||||
try:
|
||||
response = requests.get(url, timeout=readiness_timeout)
|
||||
response.close()
|
||||
return
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as exc:
|
||||
last_exception = exc
|
||||
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
pytest.fail(
|
||||
f'timed out waiting for HTTP endpoint {url}: {last_exception}',
|
||||
pytrace=False,
|
||||
)
|
||||
time.sleep(min(HTTP_RETRY_INTERVAL_S, remaining))
|
||||
|
||||
|
||||
def post_json(url, path, timeout=HTTP_TIMEOUT, **kwargs):
|
||||
endpoint = url + path
|
||||
try:
|
||||
response = requests.post(endpoint, timeout=timeout, **kwargs)
|
||||
except requests.exceptions.Timeout as exc:
|
||||
pytest.fail(f'timed out waiting for HTTP response from {endpoint}: {exc}', pytrace=False)
|
||||
except requests.exceptions.ConnectionError as exc:
|
||||
pytest.fail(f'could not connect to HTTP endpoint {endpoint}: {exc}', pytrace=False)
|
||||
|
||||
try:
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()
|
||||
finally:
|
||||
response.close()
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def client1_ready_url():
|
||||
url = 'http://' + os.environ.get('MESSENGER_TEST_CLIENT1_ADDR', '127.0.0.1:8080')
|
||||
wait_for_http(url)
|
||||
return url
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def client2_ready_url():
|
||||
url = 'http://' + os.environ.get('MESSENGER_TEST_CLIENT2_ADDR', '127.0.0.1:8081')
|
||||
wait_for_http(url)
|
||||
return url
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client1_url(client1_ready_url):
|
||||
url = client1_ready_url
|
||||
get_messages(url) # we need to flush pending messages before and after each tests
|
||||
yield url
|
||||
get_messages(url)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client2_url(client2_ready_url):
|
||||
url = client2_ready_url
|
||||
get_messages(url)
|
||||
yield url
|
||||
get_messages(url)
|
||||
|
||||
|
||||
def send_message(url, mes):
|
||||
return post_json(url, '/sendMessage', json=mes)
|
||||
|
||||
|
||||
def get_messages(url, **kwargs):
|
||||
return post_json(url, '/getAndFlushMessages', **kwargs)
|
||||
|
||||
|
||||
def wait_for_messages(url, expected):
|
||||
deadline = time.monotonic() + MESSAGE_TIMEOUT_S
|
||||
messages = []
|
||||
while len(messages) < len(expected):
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
pytest.fail(
|
||||
f'timed out waiting for messages from {url}: '
|
||||
f'expected {expected!r}, received {messages!r}',
|
||||
pytrace=False,
|
||||
)
|
||||
# Divide the remaining budget between connecting and reading the response.
|
||||
timeout = (
|
||||
min(HTTP_CONNECT_TIMEOUT_S, remaining / 2),
|
||||
min(HTTP_READ_TIMEOUT_S, remaining / 2),
|
||||
)
|
||||
batch = get_messages(url, timeout=timeout)
|
||||
assert isinstance(batch, list), f'expected a message array, received {batch!r}'
|
||||
messages.extend(batch)
|
||||
assert messages == expected[:len(messages)], \
|
||||
f'expected {expected!r}, received {messages!r}'
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
pytest.fail(f'timed out waiting for messages from {url}', pytrace=False)
|
||||
if len(messages) < len(expected):
|
||||
time.sleep(min(MESSAGE_POLL_INTERVAL_S, remaining))
|
||||
return messages
|
||||
|
||||
|
||||
def test_single_client_single_message(client1_url, client2_url):
|
||||
mes = {
|
||||
'author': 'TestSingleClient',
|
||||
'text': 'This is test text'
|
||||
}
|
||||
resp = send_message(client1_url, mes)
|
||||
mes['sendTime'] = resp['sendTime']
|
||||
assert wait_for_messages(client1_url, [mes]) == [mes]
|
||||
# Drain both subscriptions before the next test sends more messages.
|
||||
assert wait_for_messages(client2_url, [mes]) == [mes]
|
||||
|
||||
|
||||
def test_single_client_multiple_messages(client1_url, client2_url):
|
||||
mes = [{
|
||||
'author': 'TestSingleClient1',
|
||||
'text': 'This is test text'
|
||||
}, {
|
||||
'author': 'TestSingleClient2',
|
||||
'text': 'This is test text'
|
||||
}]
|
||||
for m in mes:
|
||||
resp = send_message(client1_url, m)
|
||||
m['sendTime'] = resp['sendTime']
|
||||
assert wait_for_messages(client1_url, mes) == mes
|
||||
assert wait_for_messages(client2_url, mes) == mes
|
||||
|
||||
|
||||
def test_two_clients_multiple_messages(client1_url, client2_url):
|
||||
client1_name = 'TestMultiClient1'
|
||||
client2_name = 'TestMultiClient2'
|
||||
mes = [{
|
||||
'author': client1_name,
|
||||
'text': 'This is test text #1'
|
||||
}, {
|
||||
'author': client1_name,
|
||||
'text': 'This is test text #2'
|
||||
}, {
|
||||
'author': client2_name,
|
||||
'text': 'This is test text #3'
|
||||
}, {
|
||||
'author': client2_name,
|
||||
'text': 'This is test text #4'
|
||||
}]
|
||||
times = set()
|
||||
for m in mes:
|
||||
resp = send_message(client1_url if m['author'] == client1_name else client2_url, m)
|
||||
m['sendTime'] = resp['sendTime']
|
||||
times.add(m['sendTime'])
|
||||
assert len(times) == len(mes)
|
||||
assert wait_for_messages(client1_url, mes) == mes
|
||||
assert wait_for_messages(client2_url, mes) == mes
|
||||
@@ -0,0 +1,140 @@
|
||||
import pathlib
|
||||
import tempfile
|
||||
|
||||
import grpc_tools
|
||||
from google.protobuf import descriptor
|
||||
from google.protobuf import descriptor_pb2
|
||||
from google.protobuf import descriptor_pool
|
||||
from grpc_tools import protoc
|
||||
|
||||
|
||||
SCRIPT_DIR = pathlib.Path(__file__).parent.resolve()
|
||||
PROTO_DIR = SCRIPT_DIR.parent / 'solution' / 'proto'
|
||||
PROTO_FILE = PROTO_DIR / 'messenger.proto'
|
||||
WELL_KNOWN_PROTO_DIR = pathlib.Path(grpc_tools.__file__).parent / '_proto'
|
||||
|
||||
|
||||
def compile_descriptor_set(output_path):
|
||||
result = protoc.main([
|
||||
'grpc_tools.protoc',
|
||||
f'-I{PROTO_DIR}',
|
||||
f'-I{WELL_KNOWN_PROTO_DIR}',
|
||||
f'--descriptor_set_out={output_path}',
|
||||
'--include_imports',
|
||||
str(PROTO_FILE),
|
||||
])
|
||||
assert result == 0, 'messenger.proto must compile successfully'
|
||||
|
||||
descriptor_set = descriptor_pb2.FileDescriptorSet()
|
||||
descriptor_set.ParseFromString(output_path.read_bytes())
|
||||
return descriptor_set
|
||||
|
||||
|
||||
def build_descriptor_pool(descriptor_set):
|
||||
pool = descriptor_pool.DescriptorPool()
|
||||
remaining = list(descriptor_set.file)
|
||||
while remaining:
|
||||
deferred = []
|
||||
for file_descriptor in remaining:
|
||||
try:
|
||||
pool.Add(file_descriptor)
|
||||
except TypeError:
|
||||
deferred.append(file_descriptor)
|
||||
assert len(deferred) < len(remaining), 'messenger.proto imports could not be resolved'
|
||||
remaining = deferred
|
||||
return pool
|
||||
|
||||
|
||||
def require_singular_field(message_type, field_name, field_type, message_type_name=None):
|
||||
assert field_name in message_type.fields_by_name, \
|
||||
f'{message_type.full_name} must contain field {field_name}'
|
||||
field = message_type.fields_by_name[field_name]
|
||||
assert not field.is_repeated, f'{field.full_name} must be a singular field'
|
||||
assert field.type == field_type, f'{field.full_name} has an invalid type'
|
||||
if message_type_name is not None:
|
||||
assert field.message_type is not None
|
||||
assert field.message_type.full_name == message_type_name, \
|
||||
f'{field.full_name} has an invalid message type'
|
||||
|
||||
|
||||
def require_fields_can_coexist(message_type, field_names):
|
||||
fields_by_oneof = {}
|
||||
for field_name in field_names:
|
||||
field = message_type.fields_by_name[field_name]
|
||||
if field.containing_oneof is None:
|
||||
continue
|
||||
previous_field = fields_by_oneof.setdefault(field.containing_oneof.full_name, field_name)
|
||||
assert previous_field == field_name, \
|
||||
f'{message_type.full_name} fields must allow simultaneous values'
|
||||
|
||||
|
||||
def test_proto_contract():
|
||||
with tempfile.TemporaryDirectory() as temporary_directory:
|
||||
descriptor_path = pathlib.Path(temporary_directory) / 'messenger.pb'
|
||||
descriptor_set = compile_descriptor_set(descriptor_path)
|
||||
|
||||
submitted_file = next(
|
||||
(file_descriptor for file_descriptor in descriptor_set.file
|
||||
if pathlib.PurePosixPath(file_descriptor.name).name == PROTO_FILE.name),
|
||||
None,
|
||||
)
|
||||
assert submitted_file is not None, 'messenger.proto descriptor is missing'
|
||||
assert submitted_file.syntax == 'proto3', 'messenger.proto must use proto3 syntax'
|
||||
assert submitted_file.package == 'mes_grpc', 'messenger.proto must use package mes_grpc'
|
||||
|
||||
pool = build_descriptor_pool(descriptor_set)
|
||||
try:
|
||||
messenger = pool.FindServiceByName('mes_grpc.MessengerServer')
|
||||
except KeyError:
|
||||
raise AssertionError('gRPC service must be named mes_grpc.MessengerServer') from None
|
||||
|
||||
assert 'SendMessage' in messenger.methods_by_name, \
|
||||
'MessengerServer must contain method SendMessage'
|
||||
send_message = messenger.methods_by_name['SendMessage']
|
||||
assert not send_message.client_streaming and not send_message.server_streaming, \
|
||||
'SendMessage must be unary'
|
||||
require_singular_field(
|
||||
send_message.input_type,
|
||||
'author',
|
||||
descriptor.FieldDescriptor.TYPE_STRING,
|
||||
)
|
||||
require_singular_field(
|
||||
send_message.input_type,
|
||||
'text',
|
||||
descriptor.FieldDescriptor.TYPE_STRING,
|
||||
)
|
||||
require_fields_can_coexist(send_message.input_type, ('author', 'text'))
|
||||
require_singular_field(
|
||||
send_message.output_type,
|
||||
'sendTime',
|
||||
descriptor.FieldDescriptor.TYPE_MESSAGE,
|
||||
'google.protobuf.Timestamp',
|
||||
)
|
||||
|
||||
assert 'ReadMessages' in messenger.methods_by_name, \
|
||||
'MessengerServer must contain method ReadMessages'
|
||||
read_messages = messenger.methods_by_name['ReadMessages']
|
||||
assert not read_messages.client_streaming and read_messages.server_streaming, \
|
||||
'ReadMessages must be a unary request with a server stream response'
|
||||
assert not read_messages.input_type.fields, \
|
||||
'ReadMessages request must not contain fields'
|
||||
require_singular_field(
|
||||
read_messages.output_type,
|
||||
'author',
|
||||
descriptor.FieldDescriptor.TYPE_STRING,
|
||||
)
|
||||
require_singular_field(
|
||||
read_messages.output_type,
|
||||
'text',
|
||||
descriptor.FieldDescriptor.TYPE_STRING,
|
||||
)
|
||||
require_singular_field(
|
||||
read_messages.output_type,
|
||||
'sendTime',
|
||||
descriptor.FieldDescriptor.TYPE_MESSAGE,
|
||||
'google.protobuf.Timestamp',
|
||||
)
|
||||
require_fields_can_coexist(
|
||||
read_messages.output_type,
|
||||
('author', 'text', 'sendTime'),
|
||||
)
|
||||
@@ -0,0 +1,285 @@
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import queue
|
||||
import re
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict
|
||||
|
||||
import pytest
|
||||
|
||||
test_message = {'author': 'alice', 'text': 'hello'}
|
||||
|
||||
PROTO_DIR = pathlib.Path(__file__).resolve().parent.parent / 'solution' / 'proto'
|
||||
|
||||
SOCKET_CONNECT_TIMEOUT_S = 1
|
||||
SERVICE_READY_TIMEOUT_S = 20
|
||||
SERVICE_RETRY_INTERVAL_S = 0.5
|
||||
GRPC_CALL_TIMEOUT_S = 5
|
||||
GRPC_PROCESS_TIMEOUT_S = 10
|
||||
GRPC_STREAM_TIMEOUT_S = 60
|
||||
PROCESS_STOP_TIMEOUT_S = 5
|
||||
MESSAGE_TIMEOUT_S = 10
|
||||
|
||||
|
||||
def wait_for_socket(host, port):
|
||||
deadline = time.monotonic() + SERVICE_READY_TIMEOUT_S
|
||||
last_exception = None
|
||||
while True:
|
||||
try:
|
||||
with socket.create_connection((host, port), timeout=SOCKET_CONNECT_TIMEOUT_S):
|
||||
pass
|
||||
return
|
||||
except OSError as exc:
|
||||
last_exception = exc
|
||||
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
pytest.fail(
|
||||
f'timed out waiting for TCP endpoint {host}:{port}: {last_exception}',
|
||||
pytrace=False,
|
||||
)
|
||||
time.sleep(min(SERVICE_RETRY_INTERVAL_S, remaining))
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def server_addr():
|
||||
addr = os.environ.get('MESSENGER_TEST_SERVER_ADDR', '127.0.0.1:51075')
|
||||
host = addr.split(':')[0]
|
||||
port = int(addr.split(':')[1])
|
||||
wait_for_socket(host, port)
|
||||
yield addr
|
||||
|
||||
|
||||
def send_message(server_address, message: Dict[str, str]) -> Dict[str, str]:
|
||||
grpcurl_cmd = ['grpcurl',
|
||||
'-max-time', str(GRPC_CALL_TIMEOUT_S),
|
||||
'-import-path', str(PROTO_DIR),
|
||||
'-proto', 'messenger.proto',
|
||||
'-d',
|
||||
json.dumps(message),
|
||||
'-plaintext',
|
||||
server_address,
|
||||
'mes_grpc.MessengerServer/SendMessage']
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
grpcurl_cmd,
|
||||
capture_output=True,
|
||||
check=False,
|
||||
timeout=GRPC_PROCESS_TIMEOUT_S,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
pytest.fail(
|
||||
f'grpcurl did not finish within {GRPC_PROCESS_TIMEOUT_S} seconds',
|
||||
pytrace=False,
|
||||
)
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
assert len(completed.stderr) == 0, completed.stderr
|
||||
output_str = completed.stdout.decode('ascii')
|
||||
output = json.loads(output_str)
|
||||
|
||||
message_with_timestamp = copy.deepcopy(message)
|
||||
message_with_timestamp['sendTime'] = output['sendTime']
|
||||
return message_with_timestamp
|
||||
|
||||
|
||||
class MessageStream:
|
||||
def __init__(self, server_address):
|
||||
grpcurl_cmd = ['grpcurl',
|
||||
'-max-time', str(GRPC_STREAM_TIMEOUT_S),
|
||||
'-import-path', str(PROTO_DIR),
|
||||
'-proto', 'messenger.proto',
|
||||
'-plaintext',
|
||||
server_address,
|
||||
'mes_grpc.MessengerServer/ReadMessages']
|
||||
self._process = subprocess.Popen(
|
||||
grpcurl_cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
self._messages = queue.Queue()
|
||||
self._reader_error = None
|
||||
self._reader_finished = threading.Event()
|
||||
self._closed = False
|
||||
self._reader = threading.Thread(target=self._read_messages, daemon=True)
|
||||
self._reader.start()
|
||||
|
||||
def _read_messages(self):
|
||||
try:
|
||||
message_lines = []
|
||||
for line in self._process.stdout:
|
||||
message_lines.append(line)
|
||||
if line.rstrip() == '}':
|
||||
self._messages.put(json.loads(''.join(message_lines)))
|
||||
message_lines = []
|
||||
except Exception as exc:
|
||||
self._reader_error = exc
|
||||
finally:
|
||||
self._reader_finished.set()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.close()
|
||||
|
||||
def _failure_detail(self):
|
||||
if self._reader_error is not None:
|
||||
return f'reader failed: {self._reader_error}'
|
||||
returncode = self._process.poll()
|
||||
if returncode is not None:
|
||||
stderr = self._process.stderr.read().strip()
|
||||
return f'grpcurl exited with status {returncode}: {stderr}'
|
||||
return 'grpcurl is still running but produced no matching message'
|
||||
|
||||
def _get_message(self, deadline, timeout_message):
|
||||
while True:
|
||||
if self._reader_error is not None:
|
||||
raise AssertionError(self._failure_detail())
|
||||
if self._reader_finished.is_set() and self._messages.empty():
|
||||
raise AssertionError(self._failure_detail())
|
||||
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise AssertionError(f'{timeout_message}: {self._failure_detail()}')
|
||||
try:
|
||||
return self._messages.get(timeout=min(remaining, 0.1))
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
def wait_for_message(self, expected_message, timeout, preceding_messages=None):
|
||||
deadline = time.monotonic() + timeout
|
||||
while True:
|
||||
message = self._get_message(
|
||||
deadline,
|
||||
f'timed out waiting for message {expected_message}',
|
||||
)
|
||||
if message == expected_message:
|
||||
return
|
||||
if preceding_messages is not None:
|
||||
preceding_messages.append(message)
|
||||
|
||||
def read_messages(self, count, timeout):
|
||||
deadline = time.monotonic() + timeout
|
||||
messages = []
|
||||
while len(messages) < count:
|
||||
messages.append(self._get_message(
|
||||
deadline,
|
||||
f'timed out after receiving {len(messages)} of {count} messages',
|
||||
))
|
||||
return messages
|
||||
|
||||
def close(self):
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
if self._process.poll() is None:
|
||||
self._process.terminate()
|
||||
try:
|
||||
self._process.wait(timeout=PROCESS_STOP_TIMEOUT_S)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._process.kill()
|
||||
try:
|
||||
self._process.wait(timeout=PROCESS_STOP_TIMEOUT_S)
|
||||
except subprocess.TimeoutExpired:
|
||||
pytest.fail('grpcurl did not exit after SIGKILL', pytrace=False)
|
||||
self._reader.join(timeout=PROCESS_STOP_TIMEOUT_S)
|
||||
assert not self._reader.is_alive()
|
||||
|
||||
|
||||
def wait_for_streams(server_address, streams):
|
||||
preceding_messages = [[] for _ in streams]
|
||||
probes = []
|
||||
for attempt in range(10):
|
||||
probe = send_message(
|
||||
server_address,
|
||||
{'author': 'StreamProbe', 'text': f'probe #{attempt}'},
|
||||
)
|
||||
probes.append(probe)
|
||||
streams_ready = True
|
||||
for index, stream in enumerate(streams):
|
||||
try:
|
||||
stream.wait_for_message(
|
||||
probe,
|
||||
timeout=1,
|
||||
preceding_messages=preceding_messages[index],
|
||||
)
|
||||
except AssertionError:
|
||||
streams_ready = False
|
||||
if streams_ready:
|
||||
return [
|
||||
[message for message in messages if message not in probes]
|
||||
for messages in preceding_messages
|
||||
]
|
||||
raise AssertionError('ReadMessages streams did not become ready')
|
||||
|
||||
|
||||
def timestamp_key(timestamp):
|
||||
match = re.fullmatch(r'(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(?:\.(\d{1,9}))?Z', timestamp)
|
||||
assert match is not None, f'invalid protobuf timestamp: {timestamp}'
|
||||
seconds = int(datetime.strptime(match.group(1), '%Y-%m-%dT%H:%M:%S')
|
||||
.replace(tzinfo=timezone.utc).timestamp())
|
||||
nanos = int((match.group(2) or '').ljust(9, '0'))
|
||||
return seconds, nanos
|
||||
|
||||
|
||||
def test_send_smoke(server_addr):
|
||||
send_message(server_addr, test_message)
|
||||
|
||||
|
||||
def test_send_returns_ascending_time(server_addr):
|
||||
outputs = []
|
||||
for _ in range(10):
|
||||
outputs.append(send_message(server_addr, test_message))
|
||||
|
||||
for output1, output2 in zip(outputs, outputs[1:]):
|
||||
assert timestamp_key(output1['sendTime']) < timestamp_key(output2['sendTime'])
|
||||
|
||||
|
||||
def test_get_messages_smoke(server_addr):
|
||||
with MessageStream(server_addr) as stream:
|
||||
wait_for_streams(server_addr, [stream])
|
||||
test_message_with_timestamp = send_message(server_addr, test_message)
|
||||
messages = stream.read_messages(1, timeout=MESSAGE_TIMEOUT_S)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0] == test_message_with_timestamp
|
||||
|
||||
|
||||
def test_get_only_sends_new(server_addr):
|
||||
messages1 = []
|
||||
messages3 = []
|
||||
n1, n2, n3 = 2, 3, 4
|
||||
|
||||
with MessageStream(server_addr) as stream:
|
||||
messages_before_ready = wait_for_streams(server_addr, [stream])
|
||||
assert messages_before_ready == [[]]
|
||||
for _ in range(n1):
|
||||
message = send_message(server_addr, test_message)
|
||||
messages1.append(message)
|
||||
messages = stream.read_messages(n1, timeout=MESSAGE_TIMEOUT_S)
|
||||
|
||||
assert len(messages1) == len(messages)
|
||||
for m1, m2 in zip(messages1, messages):
|
||||
assert m1 == m2
|
||||
|
||||
for _ in range(n2):
|
||||
send_message(server_addr, test_message)
|
||||
|
||||
with MessageStream(server_addr) as stream:
|
||||
messages_before_ready = wait_for_streams(server_addr, [stream])
|
||||
assert messages_before_ready == [[]]
|
||||
for _ in range(n3):
|
||||
message = send_message(server_addr, test_message)
|
||||
messages3.append(message)
|
||||
messages = stream.read_messages(n3, timeout=MESSAGE_TIMEOUT_S)
|
||||
|
||||
assert len(messages3) == len(messages)
|
||||
for m1, m2 in zip(messages3, messages):
|
||||
assert m1 == m2
|
||||
Reference in New Issue
Block a user