def test_write(self):
data = io.BytesIO()
stream = messaging.JsonIOStream(data, data, "data")
for message in self.MESSAGES:
stream.write_json(message)
data = data.getvalue()
assert data == self.SERIALIZED_MESSAGES
def test_read(self):
data = io.BytesIO(self.SERIALIZED_MESSAGES)
stream = messaging.JsonIOStream(data, data, "data")
for expected_message in self.MESSAGES:
message = stream.read_json()
assert message == expected_message
with pytest.raises(messaging.NoMoreMessages) as exc_info:
stream.read_json()
assert exc_info.value.stream is stream
def start():
global _stream, _process, _worker_log_filename
if _stream is not None:
return
args = [sys.executable, worker.__file__, str(os.getpid())]
log.info(
"Spawning {0} for tests-{1}:\n\n{2}",
_name,
os.getpid(),
"\n".join(repr(s) for s in args),
)
_process = psutil.Popen(args,
bufsize=0,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE)
_stream = messaging.JsonIOStream(_process.stdout, _process.stdin, _name)
event, _worker_log_filename = _stream.read_json()
assert event == "watchdog"
atexit.register(stop)
def test_fuzz(self):
# Set up two channels over the same stream that send messages to each other
# asynchronously, and record everything that they send and receive.
# All records should match at the end.
class Fuzzer(object):
def __init__(self, name):
self.name = name
self.lock = threading.Lock()
self.sent = []
self.received = []
self.responses_sent = []
self.responses_received = []
self.done = False
def start(self, channel):
self._worker = threading.Thread(
name=self.name,
target=lambda: self._send_requests_and_events(channel),
)
self._worker.daemon = True
self._worker.start()
def wait(self):
self._worker.join()
def done_event(self, event):
with self.lock:
self.done = True
def fizz_event(self, event):
assert event.event == "fizz"
with self.lock:
self.received.append(("event", "fizz", event.body))
def buzz_event(self, event):
assert event.event == "buzz"
with self.lock:
self.received.append(("event", "buzz", event.body))
def event(self, event):
with self.lock:
self.received.append(("event", event.event, event.body))
def make_and_log_response(self, request):
x = random.randint(-100, 100)
if x < 0:
exc_type = (messaging.InvalidMessageError if x %
2 else messaging.MessageHandlingError)
x = exc_type(str(x), request)
with self.lock:
self.responses_sent.append((request.seq, x))
return x
def fizz_request(self, request):
assert request.command == "fizz"
with self.lock:
self.received.append(
("request", "fizz", request.arguments))
return self.make_and_log_response(request)
def buzz_request(self, request):
assert request.command == "buzz"
with self.lock:
self.received.append(
("request", "buzz", request.arguments))
return self.make_and_log_response(request)
def request(self, request):
with self.lock:
self.received.append(
("request", request.command, request.arguments))
return self.make_and_log_response(request)
def _got_response(self, response):
with self.lock:
self.responses_received.append(
(response.request.seq, response.body))
def _send_requests_and_events(self, channel):
types = [
random.choice(("event", "request")) for _ in range(0, 100)
]
for typ in types:
name = random.choice(("fizz", "buzz", "fizzbuzz"))
body = random.randint(0, 100)
with self.lock:
self.sent.append((typ, name, body))
if typ == "event":
channel.send_event(name, body)
elif typ == "request":
req = channel.send_request(name, body)
req.on_response(self._got_response)
channel.send_event("done")
# Spin until we receive "done", and also get responses to all requests.
requests_sent = types.count("request")
log.info("{0} waiting for {1} responses...", self.name,
requests_sent)
while True:
with self.lock:
if self.done:
if requests_sent == len(self.responses_received):
break
time.sleep(0.1)
fuzzer1 = Fuzzer("fuzzer1")
fuzzer2 = Fuzzer("fuzzer2")
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind(("localhost", 0))
_, port = server_socket.getsockname()
server_socket.listen(0)
socket1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
socket1_thread = threading.Thread(
target=lambda: socket1.connect(("localhost", port)))
socket1_thread.start()
socket2, _ = server_socket.accept()
socket1_thread.join()
try:
io1 = socket1.makefile("rwb", 0)
io2 = socket2.makefile("rwb", 0)
stream1 = messaging.JsonIOStream(io1, io1, "socket1")
channel1 = messaging.JsonMessageChannel(stream1, fuzzer1)
channel1.start()
fuzzer1.start(channel1)
stream2 = messaging.JsonIOStream(io2, io2, "socket2")
channel2 = messaging.JsonMessageChannel(stream2, fuzzer2)
channel2.start()
fuzzer2.start(channel2)
fuzzer1.wait()
fuzzer2.wait()
finally:
socket1.close()
socket2.close()
assert fuzzer1.sent == fuzzer2.received
assert fuzzer2.sent == fuzzer1.received
assert fuzzer1.responses_sent == fuzzer2.responses_received
assert fuzzer2.responses_sent == fuzzer1.responses_received