Exemple #1
0
 def test_stream_call_with_malformed_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", ":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", []))
     )
Exemple #2
0
 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)
Exemple #3
0
 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)
Exemple #4
0
 def transform(*args, **kwargs):
     response = function(*args, **kwargs)
     try:
         message = next(response)
     except StopIteration:
         pass
     else:
         yield beretta.encode(message)
         if message == (":info", ":stream", []):
             yield beretta.encode(next(response))
             for message in response:
                 yield message
             yield b""
Exemple #5
0
 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"")
Exemple #6
0
    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)
Exemple #7
0
 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"")
Exemple #8
0
 def test_encode_regex(self):
   regex = re.compile('^(kitty)$', re.I|re.X)
   self.assertEqual(regex.pattern, '^(kitty)$')
   self.assertIn(regex.flags, (66, 98)) # python 2.x / 66, python 3.x / 98
   bytes = beretta.encode(regex)
   self.assertEqual(bytes, b'\x83h\x04d\x00\x04bertd\x00\x05regexm\x00\x00\x00\t'
                           b'^(kitty)$h\x02d\x00\x08extendedd\x00\x08caseless')
Exemple #9
0
 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"))
Exemple #10
0
 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)
Exemple #11
0
 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'", []))
     )
Exemple #12
0
 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):")
Exemple #13
0
 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"")
Exemple #14
0
 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")
Exemple #15
0
 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):")
Exemple #16
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)
Exemple #17
0
 def handle(self, connection, address):
     self.logger.info("{0}:{1} connected".format(*address))
     agent = Agent(self.modules, address)
     stream = kyoto.network.stream.receive(connection)
     try:
         for request in stream:
             for response in agent.handle(request):
                 try:
                     message = kyoto.utils.berp.pack(response)
                 except kyoto.utils.berp.MaxBERPSizeError as exception:
                     name = exception.__class__.__name__
                     description = str(exception)
                     trace = traceback.format_exc().splitlines()
                     message = (":error", (":user", 500, name, description, trace))
                     message = kyoto.utils.berp.pack(beretta.encode(message))
                 connection.sendall(message)
     except Exception as exception:
         self.logger.exception(exception)
     finally:
         connection.close()
     self.logger.info("{0}:{1} disconnected".format(*address))
Exemple #18
0
def receive(connection, server=True):
    receive_buffer = b""
    while connection:
        message = connection.recv(kyoto.conf.settings.READ_CHUNK_SIZE)
        if message:
            receive_buffer += message
            while len(receive_buffer) >= 4:
                try:
                    _, message, tail = kyoto.utils.berp.unpack(receive_buffer)
                except ValueError as exception:
                    break  # received incomplete packet, continue loop
                except kyoto.utils.berp.MaxBERPSizeError as exception:
                    if server:
                        exception = (":error", (":protocol", 3, "MaxBERPSizeError", str(exception), []))
                        exception = kyoto.utils.berp.pack(beretta.encode(exception))
                        connection.sendall(exception)
                    raise
                else:
                    receive_buffer = tail
                    yield message
        else:
            break
Exemple #19
0
 def send_message(self, connection, message):
     message = kyoto.utils.berp.pack(beretta.encode(message))
     return connection.sendall(message)
Exemple #20
0
 def test_encode_empty_list(self):
   bytes = beretta.encode([])
   self.assertEqual(bytes, b'\x83h\x02d\x00\x04bertd\x00\x03nil')
Exemple #21
0
 def test_encode_compressed(self):
   bytes = beretta.encode((':call', 'Module', 'function', []) * 4, compressed=6)
   self.assertTrue(bytes.startswith(b'\x83P'))
Exemple #22
0
 def test_encode_1(self):
   bytes = beretta.encode(1)
   self.assertEqual(bytes, b'\x83a\x01')
Exemple #23
0
 def test_encode_0_float(self):
   bytes = beretta.encode(0.0)
   self.assertEqual(bytes, b'\x83c0.00000000000000000000e+00\x00\x00\x00\x00\x00')
Exemple #24
0
 def test_encode_datetime(self):
   bytes = beretta.encode(datetime.datetime(2014, 2, 10, 6, 2, 51, 36215))
   self.assertEqual(bytes, b'\x83h\x05d\x00\x04bertd\x00\x04timeb'
                           b'\x00\x00\x05pb\x00\x00/\x8bb\x00\x00\x8dw')
Exemple #25
0
 def test_encode_0(self):
   bytes = beretta.encode(0)
   self.assertEqual(bytes, b'\x83a\x00')
Exemple #26
0
 def test_encode_empty_dict(self):
   bytes = beretta.encode({})
   self.assertEqual(bytes, b'\x83h\x03d\x00\x04bertd\x00\x04dictj')
Exemple #27
0
 def test_encode_dict(self):
   bytes = beretta.encode({'key': 'value'})
   self.assertEqual(bytes, b'\x83h\x03d\x00\x04bertd\x00\x04dictl'
                           b'\x00\x00\x00\x01h\x02m\x00\x00\x00\x03'
                           b'keym\x00\x00\x00\x05valuej')
Exemple #28
0
 def test_encode_none(self):
   bytes = beretta.encode(None)
   self.assertEqual(bytes, b'\x83h\x02d\x00\x04bertd\x00\tundefined')
Exemple #29
0
 def test_encode_false(self):
   bytes = beretta.encode(False)
   self.assertEqual(bytes, b'\x83h\x02d\x00\x04bertd\x00\x05false')
Exemple #30
0
 def test_encode_true(self):
   bytes = beretta.encode(True)
   self.assertEqual(bytes, b'\x83h\x02d\x00\x04bertd\x00\x04true')