def test_sync_call_with_broken_streaming_echo(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":broken_streaming_echo", ["hello"])))
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(beretta.decode(next(response)), (":info", ":stream", []))
self.assertEqual(beretta.decode(next(response)), (":reply", {"count": 10}))
self.assertEqual(next(response), b"")
def test_sync_call_with_streaming_response(self):
request = beretta.encode((":call", ":dummy", ":streaming_echo_response", ["hello"]))
response = self.agent.handle(request)
self.assertEqual(beretta.decode(next(response)), (":info", ":stream", []))
self.assertEqual(beretta.decode(next(response)), (":reply", {"count": 10}))
for _ in range(10):
self.assertEqual(next(response), "hello?")
self.assertEqual(next(response), b"")
def test_sync_with_streaming_request(self):
message = kyoto.utils.berp.pack(beretta.encode((":info", ":stream", [])))
status = self.connection.sendall(message)
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":streaming_echo_request", [])))
status = self.connection.sendall(message)
for message in [b"hello" for _ in range(10)]:
status = self.connection.sendall(kyoto.utils.berp.pack(message))
status = self.connection.sendall(kyoto.utils.berp.pack(b""))
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(beretta.decode(next(response)), (":info", ":stream", []))
self.assertEqual(beretta.decode(next(response)), (":noreply",))
for _ in range(10):
self.assertEqual(next(response), b"hello")
self.assertEqual(next(response), b"")
def test_reopen_connection(self):
connection = self.connections.acquire()
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":echo", ["hello"])))
connection.sendall(message)
response = kyoto.network.stream.receive(connection)
self.assertEqual(beretta.decode(next(response)), (":reply", "hello?"))
self.connections.release(connection)
connection.close()
self.assertTrue(connection.closed)
connection = self.connections.acquire()
self.assertFalse(connection.closed)
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":echo", ["hello"])))
connection.sendall(message)
response = kyoto.network.stream.receive(connection)
self.assertEqual(beretta.decode(next(response)), (":reply", "hello?"))
self.connections.release(connection)
def test_call_with_malformed_request(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":echo", ["hello"]))[2:])
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(
beretta.decode(next(response)), (":error", (":server", 3, "ValueError", "Corrupt request data", []))
)
def test_invalid_mfa(self):
response = self.agent.handle(beretta.encode((":call", ":dummy", ":kittens")))
response = beretta.decode(next(response))
self.assertEqual(response[1][0], ":server")
self.assertEqual(response[1][1], 4)
self.assertEqual(response[1][2], "ValueError")
self.assertTrue(response[1][3].startswith("Invalid MFA"))
def test_unknown_function(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":kittens", ["hello"])))
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(
beretta.decode(next(response)), (":error", (":server", 2, "NameError", "No such function: ':kittens'", []))
)
def test_sync_call_exception(self):
request = beretta.encode((":call", ":dummy", ":echo_with_exception", ["hello"]))
response = beretta.decode(next(self.agent.handle(request)))
self.assertEqual(response[0], ":error")
self.assertEqual(response[1][0], ":user")
self.assertEqual(response[1][2], "ValueError")
self.assertEqual(response[1][3], "This is exception with your text: hello")
self.assertEqual(response[1][4][0], "Traceback (most recent call last):")
def test_receive_large_stream(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":large_echo", ["hello"])))
self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
message = beretta.decode(next(response))
self.assertEqual(message[0], ":error")
self.assertEqual(message[1][0], ":user")
self.assertEqual(message[1][2], "MaxBERPSizeError")
def test_sync_call_exception(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":echo_with_exception", ["hello"])))
status = self.connection.sendall(message)
response = beretta.decode(next(kyoto.network.stream.receive(self.connection)))
self.assertEqual(response[0], ":error")
self.assertEqual(response[1][0], ":user")
self.assertEqual(response[1][2], "ValueError")
self.assertEqual(response[1][3], "This is exception with your text: hello")
self.assertEqual(response[1][4][0], "Traceback (most recent call last):")
def test_large_berp(self):
packet_size = kyoto.conf.settings.MAX_BERP_SIZE + 1024
message = struct.pack(">I", packet_size) + b"message"
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(
beretta.decode(next(response)),
(":error", (":protocol", 3, "MaxBERPSizeError", "Invalid BERP length: 33554432/33555456", [])),
)
def handle(self, message):
if not self.state["stream"]["on"]:
try:
request = beretta.decode(message)
except ValueError:
yield (":error", (":server", 3, "ValueError", "Corrupt request data", []))
else:
self.logger.info("{1}:{2} ~> {0}".format(request, *self.address))
if kyoto.utils.validation.is_valid_request(request):
for message in self.dispatcher.handle(request):
yield message
elif kyoto.utils.validation.is_valid_info(request):
if request[1] == ":stream":
self.state["stream"]["on"] = True
self.state["stream"]["queue"] = gevent.queue.Queue()
else:
raise NotImplementedError
else:
yield (":error", (":server", 4, "ValueError", "Invalid MFA: {0}".format(request), []))
else:
if not self.state["stream"]["request"]:
try:
request = beretta.decode(message)
except ValueError:
self.state["stream"]["on"] = False
yield (":error", (":server", 3, "ValueError", "Corrupt request data", []))
else:
if kyoto.utils.validation.is_valid_request(request):
self.state["stream"]["request"] = request
self.state["stream"]["worker"] = gevent.spawn(self.dispatcher.handle,
self.state["stream"]["request"],
stream=self.state["stream"]["queue"])
else:
raise NotImplementedError
else:
if message:
self.state["stream"]["queue"].put(message)
else:
self.state["stream"]["on"] = False
self.state["stream"]["request"] = None
self.state["stream"]["queue"].put(StopIteration)
response = self.state["stream"]["worker"].get()
for message in response:
yield message
def test_stream_call_with_malformed_request(self):
self.assertEqual(self.agent.state["stream"]["on"], False)
self.assertEqual(self.agent.state["stream"]["request"], None)
info = beretta.encode((":info", ":stream", []))
request = beretta.encode((":call", ":dummy", ":echo", ["hello"]))[2:]
with self.assertRaises(StopIteration):
response = next(self.agent.handle(info))
response = next(self.agent.handle(request))
self.assertEqual(beretta.decode(response), (":error", (":server", 3, "ValueError", "Corrupt request data", [])))
self.assertEqual(self.agent.state["stream"]["on"], False)
self.assertEqual(self.agent.state["stream"]["request"], None)
def handle_response(self, stream):
response = beretta.decode(next(stream))
rtype = response[0]
if rtype == ":reply":
return response[1]
elif rtype == ":noreply":
return None
elif rtype == ":error":
raise ValueError(response)
else:
raise NotImplementedError
def test_send_file_stream(self):
info = kyoto.utils.berp.pack(beretta.encode((":info", ":stream", [])))
self.connection.sendall(info)
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":streaming_echo_length", [])))
self.connection.sendall(message)
with open("/etc/passwd", "rb") as source:
stream = kyoto.network.stream.send(source)
for chunk in stream:
self.connection.sendall(chunk)
response = kyoto.network.stream.receive(self.connection)
reply, length = beretta.decode(next(response))
self.assertTrue(length > 0)
def test_sync_with_streaming_request(self):
self.assertEqual(self.agent.state["stream"]["on"], False)
self.assertEqual(self.agent.state["stream"]["request"], None)
info = beretta.encode((":info", ":stream", []))
request = beretta.encode((":call", ":dummy", ":streaming_echo_request", []))
with self.assertRaises(StopIteration):
response = next(self.agent.handle(info))
self.assertEqual(self.agent.state["stream"]["on"], True)
self.assertEqual(self.agent.state["stream"]["request"], None)
with self.assertRaises(StopIteration):
response = next(self.agent.handle(request))
self.assertEqual(self.agent.state["stream"]["on"], True)
self.assertEqual(self.agent.state["stream"]["request"], (":call", ":dummy", ":streaming_echo_request", []))
for message in ["hello" for _ in range(10)]:
with self.assertRaises(StopIteration):
response = next(self.agent.handle(message))
response = self.agent.handle(b"")
self.assertEqual(beretta.decode(next(response)), (":info", ":stream", []))
self.assertEqual(beretta.decode(next(response)), (":noreply",))
for _ in range(10):
self.assertEqual(next(response), "hello")
self.assertEqual(next(response), b"")
self.assertEqual(self.agent.state["stream"]["on"], False)
self.assertEqual(self.agent.state["stream"]["request"], None)
def test_send_generator_stream(self):
def payload(message):
for x in range(10):
yield message
info = kyoto.utils.berp.pack(beretta.encode((":info", ":stream", [])))
self.connection.sendall(info)
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":streaming_echo_length", [])))
self.connection.sendall(message)
stream = kyoto.network.stream.send(payload("hello"))
for chunk in stream:
self.connection.sendall(chunk)
response = kyoto.network.stream.receive(self.connection)
reply, length = beretta.decode(next(response))
self.assertEqual(length, 50)
def test_async_call_with_streaming_response(self):
request = beretta.encode((":cast", ":dummy", ":streaming_echo_response", ["hello"]))
response = self.agent.handle(request)
self.assertEqual(beretta.decode(next(response)), (":noreply",))
def test_decode_unknown_bert_type(self):
with self.assertRaises(ValueError):
result = beretta.decode(b'\x83h\x03d\x00\x04bertd\x00\x07'
b'kittensh\x02d\x00\x04bertd\x00\x03nil')
def test_decode_zero_float(self): result = beretta.decode(b'\x83c0.00000000000000000000e+00\x00\x00\x00\x00\x00') self.assertEqual(result, 0.0)
def test_unknown_function(self):
response = self.agent.handle(beretta.encode((":call", ":dummy", ":kittens", ["hello"])))
self.assertEqual(
beretta.decode(next(response)), (":error", (":server", 2, "NameError", "No such function: ':kittens'", []))
)
def test_decode_regex(self):
result = beretta.decode(b'\x83h\x04d\x00\x04bertd\x00\x05regexm\x00\x00\x00\t'
b'^(kitty)$h\x02d\x00\x08extendedd\x00\x08caseless')
self.assertEqual(result.pattern, '^(kitty)$')
self.assertIn(result.flags, (66, 98)) # python 2.x / 66, python 3.x / 98
def test_decode_tuple(self):
result = beretta.decode(b'\x83h\x04d\x00\x04callm\x00\x00\x00'
b'\x06Modulem\x00\x00\x00\x08functionl'
b'\x00\x00\x00\x01h\x02d\x00\x04bertd'
b'\x00\x04truej')
self.assertEqual(result, (':call', 'Module', 'function', [True]))
def test_async_call_exception(self):
message = kyoto.utils.berp.pack(beretta.encode((":cast", ":dummy", ":echo_with_exception", ["hello"])))
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(beretta.decode(next(response)), (":noreply",))
def test_unknown_module(self):
response = self.agent.handle(beretta.encode((":call", ":Kittens", ":echo", ["hello"])))
self.assertEqual(
beretta.decode(next(response)), (":error", (":server", 1, "NameError", "No such module: ':Kittens'", []))
)
def test_async_call_exception(self):
request = beretta.encode((":cast", ":dummy", ":echo_with_exception", ["hello"]))
response = self.agent.handle(request)
self.assertEqual(beretta.decode(next(response)), (":noreply",))
def test_decode_datetime(self):
result = beretta.decode(b'\x83h\x05d\x00\x04bertd\x00\x04timeb'
b'\x00\x00\x05pb\x00\x00/\x8bb\x00\x00\x8dw')
self.assertEqual(result, datetime.datetime(2014, 2, 10, 6, 2, 51, 36215))
def test_call_with_malformed_request(self):
request = beretta.encode((":call", ":dummy", ":echo", ["hello"]))[2:]
response = next(self.agent.handle(request))
self.assertEqual(beretta.decode(response), (":error", (":server", 3, "ValueError", "Corrupt request data", [])))
def test_sync_request(self):
message = kyoto.utils.berp.pack(beretta.encode((":call", ":dummy", ":echo", ["hello"])))
status = self.connection.sendall(message)
response = kyoto.network.stream.receive(self.connection)
self.assertEqual(beretta.decode(next(response)), (":reply", "hello?"))
def test_decode_empty_dict(self):
result = beretta.decode(b'\x83h\x03d\x00\x04bertd\x00\x04dictj')
self.assertEqual(result, {})