tests: Add test_hub:test_reserve_place()

tests/test_hub.py

@@ -1,11 +1,16 @@
 import asyncio
+import contextlib
 
 import pytest
 
 import not_my_board._hub as hubmodule
 import not_my_board._util as util
 
 
+DEFAULT_EXPORTER_IP = "3.1.1.1"
+DEFAULT_AGENT_IP = "6.1.1.1"
+
+
 @pytest.fixture(scope="function")
 def hub():
     yield hubmodule.Hub()
@@ -16,39 +21,92 @@ async def test_no_places_on_startup(hub):
     assert places["places"] == []
 
 
-async def test_register_exporter(hub):
-    class FakeExporter:
-        def __init__(self, register_event):
-            self._register_event = register_event
-
-        async def io_loop(self):
-            # wait forever
-            await asyncio.Event().wait()
-
-        async def get_place(self):
-            self._register_event.set()
-            return {
-                "port": 1234,
-                "parts": [
-                    {
-                        "compatible": "test-board",
-                        "tcp": {
-                            "test-if": {
-                                "host": "localhost",
-                                "port": 8080,
-                            },
+class FakeExporter:
+    def __init__(self, register_event):
+        self._register_event = register_event
+
+    async def io_loop(self):
+        # wait forever
+        await asyncio.Event().wait()
+
+    async def get_place(self):
+        self._register_event.set()
+        return {
+            "port": 1234,
+            "parts": [
+                {
+                    "compatible": "test-board",
+                    "tcp": {
+                        "test-if": {
+                            "host": "localhost",
+                            "port": 8080,
                         },
                     },
-                ],
-            }
+                },
+            ],
+        }
+
+    async def set_allowed_ips(self, ips):
+        self._allowed_ips = ips
 
-    exporter_ip = "1.1.1.3"
+    @property
+    def allowed_ips(self):
+        return self._allowed_ips
+
+
+@contextlib.asynccontextmanager
+async def register_exporter(hub):
+    exporter_ip = DEFAULT_EXPORTER_IP
     register_event = asyncio.Event()
     fake_exporter = FakeExporter(register_event)
     coro = hub.exporter_communicate(exporter_ip, fake_exporter)
     async with util.background_task(coro):
         async with asyncio.timeout(2):
             await register_event.wait()
+        yield fake_exporter
+
 
+async def test_register_exporter(hub):
+    async with register_exporter(hub):
         places = await hub.get_places()
         assert len(places["places"]) == 1
+
+
+class FakeAgent:
+    def __init__(self, register_event):
+        self._register_event = register_event
+
+    def set_api_object(self, api_obj):
+        self._api_obj = api_obj
+        self._register_event.set()
+
+    async def serve_forever(self):
+        # wait forever
+        await asyncio.Event().wait()
+
+    def __getattr__(self, method_name):
+        if method_name.startswith("_"):
+            raise AttributeError(f"invalid attribute '{method_name}'")
+        return getattr(self._api_obj, method_name)
+
+
+@contextlib.asynccontextmanager
+async def register_agent(hub):
+    agent_ip = DEFAULT_AGENT_IP
+    register_event = asyncio.Event()
+    fake_agent = FakeAgent(register_event)
+    coro = hub.agent_communicate(agent_ip, fake_agent)
+    async with util.background_task(coro):
+        async with asyncio.timeout(2):
+            await register_event.wait()
+        yield fake_agent
+
+
+async def test_reserve_place(hub):
+    async with register_exporter(hub) as exporter:
+        async with register_agent(hub) as agent:
+            places = await hub.get_places()
+            candidate_ids = [places["places"][0]["id"]]
+            reserved_id = await agent.reserve(candidate_ids)
+            assert reserved_id == candidate_ids[0]
+            assert exporter.allowed_ips == [DEFAULT_AGENT_IP]