refactoring connection into a package

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

Author: jhjaggars

receptor/connection.py receptor/connection/__init__.pyreceptor/connection.py renamed to receptor/connection/__init__.py

@@ -1,11 +1,8 @@
 import logging
 
-import functools
 import asyncio
-import aiohttp
-import aiohttp.web
 
-from .messages.envelope import FramedBuffer
+from ..messages.envelope import FramedBuffer
 
 logger = logging.getLogger(__name__)
 
@@ -28,25 +25,6 @@ async def send(self, bytes_):
         raise NotImplementedError("subclasses should implement this")
 
 
-class WebSocket(Transport):
-    def __init__(self, ws):
-        self.ws = ws
-
-    async def __anext__(self):
-        msg = await self.ws.__anext__()
-        return msg.data
-
-    async def close(self):
-        return await self.ws.close()
-
-    @property
-    def closed(self):
-        return self.ws.closed
-
-    async def send(self, bytes_):
-        await self.ws.send_bytes(bytes_)
-
-
 async def watch_queue(conn, buf):
     while not conn.closed:
         try:
@@ -145,35 +123,3 @@ async def server(self, transport):
             await self.start_processing()
         finally:
             self.unregister()
-
-
-async def connect(uri, factory, loop=None):
-    if not loop:
-        loop = asyncio.get_event_loop()
-
-    worker = factory()
-    try:
-        async with aiohttp.ClientSession().ws_connect(uri) as ws:
-            t = WebSocket(ws)
-            await worker.client(t)
-    except Exception:
-        logger.exception("connect")
-    finally:
-        await asyncio.sleep(5)
-        logger.debug("reconnecting")
-        loop.create_task(connect(uri, factory=factory, loop=loop))
-
-
-async def serve(request, factory):
-    ws = aiohttp.web.WebSocketResponse()
-    await ws.prepare(request)
-
-    t = WebSocket(ws)
-    await factory().server(t)
-
-
-def app(factory):
-    handler = functools.partial(serve, factory=factory)
-    app = aiohttp.web.Application()
-    app.add_routes([aiohttp.web.get("/", handler)])
-    return app

receptor/connection/sock.py

@@ -0,0 +1,23 @@
+from . import Transport
+
+
+class RawSocket(Transport):
+    def __init__(self, reader, writer):
+        self.reader = reader
+        self.writer = writer
+        self._closed = False
+
+    async def __anext__(self):
+        return await self.reader.read()
+
+    @property
+    def closed(self):
+        return self._closed
+
+    async def close(self):
+        self._closed = True
+        await self.writer.close()
+
+    async def send(self, bytes_):
+        self.writer.write(bytes_)
+        await self.writer.drain()

receptor/connection/ws.py

@@ -0,0 +1,60 @@
+import logging
+import functools
+import aiohttp
+import aiohttp.web
+import asyncio
+
+from . import Transport
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocket(Transport):
+    def __init__(self, ws):
+        self.ws = ws
+
+    async def __anext__(self):
+        msg = await self.ws.__anext__()
+        return msg.data
+
+    async def close(self):
+        return await self.ws.close()
+
+    @property
+    def closed(self):
+        return self.ws.closed
+
+    async def send(self, bytes_):
+        await self.ws.send_bytes(bytes_)
+
+
+async def connect(uri, factory, loop=None):
+    if not loop:
+        loop = asyncio.get_event_loop()
+
+    worker = factory()
+    try:
+        async with aiohttp.ClientSession().ws_connect(uri) as ws:
+            t = WebSocket(ws)
+            await worker.client(t)
+    except Exception:
+        logger.exception("connect")
+    finally:
+        await asyncio.sleep(5)
+        logger.debug("reconnecting")
+        loop.create_task(connect(uri, factory=factory, loop=loop))
+
+
+async def serve(request, factory):
+    ws = aiohttp.web.WebSocketResponse()
+    await ws.prepare(request)
+
+    t = WebSocket(ws)
+    await factory().server(t)
+
+
+def app(factory):
+    handler = functools.partial(serve, factory=factory)
+    app = aiohttp.web.Application()
+    app.add_routes([aiohttp.web.get("/", handler)])
+    return app

receptor/controller.py

@@ -8,6 +8,8 @@
 from .receptor import Receptor
 from .messages import envelope
 from . import connection
+from .connection import ws, Worker
+from . import protocol
 
 logger = logging.getLogger(__name__)
 
@@ -38,9 +40,9 @@ def enable_server(self, listen_url):
 
     def enable_websocket_server(self, listen_url):
         service = urlparse(listen_url)
-        factory = lambda: connection.Worker(receptor, loop)
+        factory = lambda: Worker(receptor, loop)
         listener = self.loop.create_server(
-            connection.app(factory).make_handler(),
+            ws.app(factory).make_handler(),
             service.hostname, service.port,
             ssl=self.receptor.config.get_server_ssl_context())
         logger.info("Serving websockets on {}:{}".format(service.hostname, service.port))