Add HW 2
This commit is contained in:
@@ -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