X-Git-Url: https://osm.etsi.org/gitweb/?p=osm%2FN2VC.git;a=blobdiff_plain;f=modules%2Flibjuju%2Ftests%2Fbase.py;h=97eea530955ea07ec39f4e64a3b55a0a54e6521d;hp=e1ec45238228e284ba41294b6f08161d89cae0f9;hb=refs%2Fchanges%2F94%2F6394%2F1;hpb=68858c1915122c2dbc8999a5cd3229694abf5f3a

diff --git a/modules/libjuju/tests/base.py b/modules/libjuju/tests/base.py
index e1ec452..97eea53 100644
--- a/modules/libjuju/tests/base.py
+++ b/modules/libjuju/tests/base.py
@@ -1,67 +1,149 @@
-import mock
+import inspect
 import subprocess
 import uuid
+from contextlib import contextmanager
+from pathlib import Path
 
-import pytest
-
+import mock
+from juju.client.jujudata import FileJujuData
 from juju.controller import Controller
-from juju.client.connection import JujuData
+
+import pytest
 
 
 def is_bootstrapped():
-    result = subprocess.run(['juju', 'switch'], stdout=subprocess.PIPE)
-    return (
-        result.returncode == 0 and
-        len(result.stdout.decode().strip()) > 0)
+    try:
+        result = subprocess.run(['juju', 'switch'], stdout=subprocess.PIPE)
+        return (
+            result.returncode == 0 and
+            len(result.stdout.decode().strip()) > 0)
+    except FileNotFoundError:
+        return False
+
+
 
 bootstrapped = pytest.mark.skipif(
     not is_bootstrapped(),
     reason='bootstrapped Juju environment required')
 
+test_run_nonce = uuid.uuid4().hex[-4:]
+
 
 class CleanController():
+    """
+    Context manager that automatically connects and disconnects from
+    the currently active controller.
+
+    Note: Unlike CleanModel, this will not create a new controller for you,
+    and an active controller must already be available.
+    """
     def __init__(self):
-        self.controller = None
+        self._controller = None
 
     async def __aenter__(self):
-        self.controller = Controller()
-        await self.controller.connect_current()
-        return self.controller
+        self._controller = Controller()
+        await self._controller.connect()
+        return self._controller
 
     async def __aexit__(self, exc_type, exc, tb):
-        await self.controller.disconnect()
+        await self._controller.disconnect()
 
 
 class CleanModel():
-    def __init__(self):
-        self.controller = None
-        self.model = None
+    """
+    Context manager that automatically connects to the currently active
+    controller, adds a fresh model, returns the connection to that model,
+    and automatically disconnects and cleans up the model.
+
+    The new model is also set as the current default for the controller
+    connection.
+    """
+    def __init__(self, bakery_client=None):
+        self._controller = None
+        self._model = None
+        self._model_uuid = None
+        self._bakery_client = bakery_client
 
     async def __aenter__(self):
-        self.controller = Controller()
-        await self.controller.connect_current()
-
-        model_name = 'model-{}'.format(uuid.uuid4())
-        self.model = await self.controller.add_model(model_name)
+        model_nonce = uuid.uuid4().hex[-4:]
+        frame = inspect.stack()[1]
+        test_name = frame.function.replace('_', '-')
+        jujudata = TestJujuData()
+        self._controller = Controller(
+            jujudata=jujudata,
+            bakery_client=self._bakery_client,
+        )
+        controller_name = jujudata.current_controller()
+        user_name = jujudata.accounts()[controller_name]['user']
+        await self._controller.connect(controller_name)
+
+        model_name = 'test-{}-{}-{}'.format(
+            test_run_nonce,
+            test_name,
+            model_nonce,
+        )
+        self._model = await self._controller.add_model(model_name)
+
+        # Change the JujuData instance so that it will return the new
+        # model as the current model name, so that we'll connect
+        # to it by default.
+        jujudata.set_model(
+            controller_name,
+            user_name + "/" + model_name,
+            self._model.info.uuid,
+        )
 
         # save the model UUID in case test closes model
-        self.model_uuid = self.model.info.uuid
+        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',
-                                           return_value=model_name)
-        self._patch_cm.start()
-
-        return self.model
+        return self._model
 
     async def __aexit__(self, exc_type, exc, tb):
-        self._patch_cm.stop()
-        await self.model.disconnect()
-        await self.controller.destroy_model(self.model_uuid)
-        await self.controller.disconnect()
+        await self._model.disconnect()
+        await self._controller.destroy_model(self._model_uuid)
+        await self._controller.disconnect()
+
+
+class TestJujuData(FileJujuData):
+    def __init__(self):
+        self.__controller_name = None
+        self.__model_name = None
+        self.__model_uuid = None
+        super().__init__()
+
+    def set_model(self, controller_name, model_name, model_uuid):
+        self.__controller_name = controller_name
+        self.__model_name = model_name
+        self.__model_uuid = model_uuid
+
+    def current_model(self, *args, **kwargs):
+        return self.__model_name or super().current_model(*args, **kwargs)
+
+    def models(self):
+        all_models = super().models()
+        if self.__model_name is None:
+            return all_models
+        all_models.setdefault(self.__controller_name, {})
+        all_models[self.__controller_name].setdefault('models', {})
+        cmodels = all_models[self.__controller_name]['models']
+        cmodels[self.__model_name] = {'uuid': self.__model_uuid}
+        return all_models
 
 
 class AsyncMock(mock.MagicMock):
     async def __call__(self, *args, **kwargs):
         return super().__call__(*args, **kwargs)
+
+
+@contextmanager
+def patch_file(filename):
+    """
+    "Patch" a file so that its current contents are automatically restored
+    when the context is exited.
+    """
+    filepath = Path(filename).expanduser()
+    data = filepath.read_bytes()
+    try:
+        yield
+    finally:
+        filepath.write_bytes(data)