moving tests to a more reasonable place

Signed-off-by: Jesse Jaggars <jjaggars@redhat.com>

Author: jhjaggars

Pipfile

@@ -9,6 +9,7 @@ ipython = "*"
 pytest = "*"
 pytest-asyncio = "*"
 flake8 = "*"
+rope = "*"
 
 [packages]
 python-dateutil = "*"

Pipfile.lock

@@ -36,21 +36,22 @@ async def put(self, data):
 
     async def handle_frame(self, data):
         self.current_frame, rest = Frame.from_data(data)
-        if self.current_frame.type in (Frame.START_MSG, Frame.PAYLOAD):
-            self.to_read = self.current_frame.length
-            await self.consume(rest)
-        else:
+
+        if self.current_frame.type not in (Frame.HEADER, Frame.PAYLOAD):
             raise Exception("Unknown Frame Type")
 
+        self.to_read = self.current_frame.length
+        await self.consume(rest)
+
     async def consume(self, data):
         self.to_read -= len(data)
         self.bb += data
         if self.to_read == 0:
             await self.finish()
 
     async def finish(self):
-        if self.current_frame.type == Frame.START_MSG:
-            self.header = Header(**json.loads(self.bb))
+        if self.current_frame.type == Frame.HEADER:
+            self.header = Header.from_bytes(self.bb)
         elif self.current_frame.type == Frame.PAYLOAD:
             await self.q.put((self.header, self.bb))
             self.header = None
@@ -62,9 +63,11 @@ async def get(self):
 
 
 class Frame:
-    START_MSG = 0
+    HEADER = 0
     PAYLOAD = 1
 
+    __slots__ = ('type', 'version', 'length', 'msg_id', 'id')
+
     def __init__(self, type_, version, length, msg_id, id_):
         self.type = type_
         self.version = version
@@ -103,8 +106,12 @@ def __init__(self, sender, recipient, route_list):
     def serialize(self):
         return json.dumps({"sender": self.sender, "recipient": self.recipient, "route_list": self.route_list}).encode("utf-8")
 
+    @classmethod
+    def from_bytes(cls, data):
+        return cls(**json.loads(data))
+
     def __repr__(self):
-        return f"Header: {self.sender}, {self.recipient}, {self.route_list}"
+        return f"Header({self.sender}, {self.recipient}, {self.route_list})"
 
     def __eq__(self, other):
         return (self.sender, self.recipient, self.route_list) == (other.sender, other.recipient, other.route_list)
@@ -117,7 +124,7 @@ def gen_chunks(data, header, msg_id=None, chunksize=2 ** 8):
     buf = bytearray(chunksize)
     bv = memoryview(buf)
     header = header.serialize()
-    yield Frame(Frame.START_MSG, 1, len(header), msg_id, next(seq)).serialize() + header
+    yield Frame(Frame.HEADER, 1, len(header), msg_id, next(seq)).serialize() + header
     yield Frame(Frame.PAYLOAD, 1, len(data), msg_id, next(seq)).serialize()
     buffer = io.BytesIO(data)
     bytes_read = buffer.readinto(buf)

receptor/messages/envelope.py

@@ -0,0 +1,83 @@
+import json
+import uuid
+
+import pytest
+
+from receptor.messages.envelope import Frame, FramedBuffer, Header, gen_chunks
+
+
+@pytest.yield_fixture
+def msg_id():
+    return uuid.uuid4().int
+
+
+@pytest.yield_fixture
+def framed_buffer():
+    return FramedBuffer()
+
+
+@pytest.mark.asyncio
+async def test_framedbuffer(framed_buffer, msg_id):
+    header = Header("node1", "node2", [])
+    header_bytes = header.serialize()
+    f1 = Frame(Frame.HEADER, 1, len(header_bytes), msg_id, 1)
+
+    await framed_buffer.put(f1.serialize() + header_bytes)
+
+    payload = b"tina loves butts"
+    payload2 = b"yep yep yep"
+    f2 = Frame(Frame.PAYLOAD, 1, len(payload) + len(payload2), msg_id, 2)
+
+    await framed_buffer.put(f2.serialize() + payload)
+    await framed_buffer.put(payload2)
+
+    h, p = await framed_buffer.get()
+
+    assert h == header
+    assert p == payload + payload2
+
+
+@pytest.mark.asyncio
+async def test_gen_chunks():
+
+    b = FramedBuffer()
+
+    header = Header("node1", "node2", [])
+    payload = b"this is a test with a buffer"
+    for chunk in gen_chunks(payload, header):
+        await b.put(chunk)
+
+    h, p = await b.get()
+    assert h == header
+    assert p == payload
+
+
+@pytest.mark.asyncio
+async def test_hi(msg_id, framed_buffer):
+    hi = json.dumps({"cmd": "hi"}).encode("utf-8")
+    f1 = Frame(Frame.PAYLOAD, 1, len(hi), msg_id, 1)
+
+    await framed_buffer.put(f1.serialize())
+    await framed_buffer.put(hi)
+
+    h, p = await framed_buffer.get()
+
+    assert h is None
+    assert p == hi
+
+
+@pytest.mark.asyncio
+async def test_extra_header(framed_buffer, msg_id):
+    h1 = Header("node1", "node2", [])
+    payload = h1.serialize()
+    f1 = Frame(Frame.HEADER, 1, len(payload), msg_id, 1)
+    await framed_buffer.put(f1.serialize())
+    await framed_buffer.put(payload)
+
+    h2 = Header("node3", "node4", [])
+    payload = h2.serialize()
+    f2 = Frame(Frame.HEADER, 1, len(payload), msg_id, 2)
+    await framed_buffer.put(f2.serialize())
+    await framed_buffer.put(payload)
+
+    assert framed_buffer.header == h2

receptor/tests/test_framedbuffer.py

@@ -1,9 +1,6 @@
-import uuid
-
 import pytest
 
 from receptor import protocol
-from receptor.messages.envelope import Frame, FramedBuffer, Header, gen_chunks
 
 
 def deser(x):
@@ -35,42 +32,3 @@ async def test_databuffer_many_msgs():
     assert msg[0] == await b.get()
     assert msg[1] == await b.get()
     assert b.q.empty()
-
-
-@pytest.mark.asyncio
-async def test_framedbuffer():
-    b = FramedBuffer()
-
-    msg_id = uuid.uuid4().int
-    header = Header("node1", "node2", [])
-    header_bytes = header.serialize()
-    f1 = Frame(Frame.START_MSG, 1, len(header_bytes), msg_id, 1)
-
-    await b.put(f1.serialize() + header_bytes)
-
-    payload = b"tina loves butts"
-    payload2 = b"yep yep yep"
-    f2 = Frame(Frame.PAYLOAD, 1, len(payload) + len(payload2), msg_id, 2)
-
-    await b.put(f2.serialize() + payload)    
-    await b.put(payload2)
-
-    h, p = await b.get()
-
-    assert h == header
-    assert p == payload + payload2
-
-
-@pytest.mark.asyncio
-async def test_gen_chunks():
-
-    b = FramedBuffer()
-
-    header = Header("node1", "node2", [])
-    payload = b"this is a test with a buffer"
-    for chunk in gen_chunks(payload, header):
-        await b.put(chunk)
-    
-    h, p = await b.get()
-    assert h == header
-    assert p == payload