Expand integration tests to use stable/edge versions of juju (#155)

[osm/N2VC.git] / tests / base.py

diff --git a/tests/base.py b/tests/base.py
index af386ea..e1ec452 100644 (file)
--- a/tests/base.py
+++ b/tests/base.py
@@ -19,6 +19,19 @@ bootstrapped = pytest.mark.skipif(
     reason='bootstrapped Juju environment required')
 
 
+class CleanController():
+    def __init__(self):
+        self.controller = None
+
+    async def __aenter__(self):
+        self.controller = Controller()
+        await self.controller.connect_current()
+        return self.controller
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.controller.disconnect()
+
+
 class CleanModel():
     def __init__(self):
         self.controller = None
@@ -31,6 +44,9 @@ class CleanModel():
         model_name = 'model-{}'.format(uuid.uuid4())
         self.model = await self.controller.add_model(model_name)
 
+        # save the model UUID in case test closes model
+        self.model_uuid = self.model.info.uuid
+
         # Ensure that we connect to the new model by default.  This also
         # prevents failures if test was started with no current model.
         self._patch_cm = mock.patch.object(JujuData, 'current_model',
@@ -42,5 +58,10 @@ class CleanModel():
     async def __aexit__(self, exc_type, exc, tb):
         self._patch_cm.stop()
         await self.model.disconnect()
-        await self.controller.destroy_model(self.model.info.uuid)
+        await self.controller.destroy_model(self.model_uuid)
         await self.controller.disconnect()
+
+
+class AsyncMock(mock.MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)