Flock tests: BIRDInstance class for MinimalistMachines running BIRD

flock/bgp-secondary/test.py

@@ -12,22 +12,16 @@ name = selfpath.parent.stem
 
 sys.path.insert(0, str(selfpath.parent.parent / "lib"))
 
-from BIRD.CLI import CLI, Transport
+from BIRD.Test import Test, BIRDInstance
-from BIRD.Test import Test
 
 os.chdir(pathlib.Path(__file__).parent)
 
-class MinimalistTransport(Transport):
-    def __init__(self, socket, machine):
-        self.sock = socket
-        self.machine = machine
-
-    async def send_cmd(self, *args):
-        return await self.sock.send_cmd("run_in", self.machine, "./birdc", "-l", *args)
-
 class ThisTest(Test):
     async def start(self):
-        self.src, self.dest = await self.machines("src", "dest")
+        self.src, self.dest = await self.machines(
+                "src", "dest",
+                t=BIRDInstance,
+                )
 
 async def main():
     t = ThisTest(name)
@@ -89,23 +83,20 @@ async def main():
 
     await asyncio.sleep(5)
 
-    src_cli = CLI(MinimalistTransport(h.control_socket, "src"))
-    dest_cli = CLI(MinimalistTransport(h.control_socket, "dest"))
-
     print(await asyncio.gather(*[
         where.show_route()
-        for where in (src_cli, dest_cli)
+        for where in (t.src, t.dest)
         ]))
 
     await asyncio.sleep(1)
 
     for p in ("p170", "p180", "p190", "p200"):
-        await src_cli.enable(p)
+        await t.src.enable(p)
         await asyncio.sleep(1)
 
         shr = await asyncio.gather(*[
             where.show_route()
-            for where in (src_cli, dest_cli)
+            for where in (t.src, t.dest)
             ])
 
         print(shr[0]["out"].decode(), shr[1]["out"].decode())
@@ -114,12 +105,12 @@ async def main():
 
     print(await asyncio.gather(*[
         where.show_route()
-        for where in (src_cli, dest_cli)
+        for where in (t.src, t.dest)
         ]))
 
     print(await asyncio.gather(*[
         c.down()
-        for c in (src_cli, dest_cli)
+        for c in (t.src, t.dest)
         ]))
 
     await asyncio.sleep(5)

flock/lib/BIRD/Test.py

@@ -7,6 +7,27 @@ sys.path.insert(0, "/home/maria/flock")
 
 from flock.Hypervisor import Hypervisor
 from flock.Machine import Machine
+from .CLI import CLI, Transport
+
+class MinimalistTransport(Transport):
+    def __init__(self, socket, machine):
+        self.sock = socket
+        self.machine = machine
+
+    async def send_cmd(self, *args):
+        return await self.sock.send_cmd("run_in", self.machine, "./birdc", "-l", *args)
+
+class BIRDInstance(CLI):
+    def __init__(self, mach: Machine):
+        self.mach = mach
+        self.workdir = self.mach.workdir
+
+        super().__init__(
+                transport=MinimalistTransport(
+                    socket=mach.hypervisor.control_socket,
+                    machine=self.mach.name
+                    )
+                )
 
 class Test:
     machines_start = {}
@@ -34,17 +55,17 @@ class Test:
 
         return await self.hypervisor.control_socket.send_cmd_early(*args)
 
-    async def machines(self, *names):
+    async def machines(self, *names, t: type):
         info = await asyncio.gather(*[
             self.hcom("machine", name, { "type": "minimalist" })
             for name in names
             ])
 
         return [
-                Machine.new(
+                t(mach=Machine.new(
                     name=n,
                     hypervisor=self.hypervisor,
                     **i
-                    ) for n,i in zip(names, info)
+                    )) for n,i in zip(names, info)
                 ]