modules/libjuju/tests/base.py

   1 import inspect
   2 import subprocess
   3 import uuid
   4 from contextlib import contextmanager
   5 from pathlib import Path
   6
   7 import mock
   8 from juju.client.jujudata import FileJujuData
   9 from juju.controller import Controller
  10
  11 import pytest
  12
  13
  14 def is_bootstrapped():
  15     try:
  16         result = subprocess.run(['juju', 'switch'], stdout=subprocess.PIPE)
  17         return (
  18             result.returncode == 0 and
  19             len(result.stdout.decode().strip()) > 0)
  20     except FileNotFoundError:
  21         return False
  22
  23
  24 bootstrapped = pytest.mark.skipif(
  25     not is_bootstrapped(),
  26     reason='bootstrapped Juju environment required')
  27
  28 test_run_nonce = uuid.uuid4().hex[-4:]
  29
  30
  31 class CleanController():
  32     """
  33     Context manager that automatically connects and disconnects from
  34     the currently active controller.
  35
  36     Note: Unlike CleanModel, this will not create a new controller for you,
  37     and an active controller must already be available.
  38     """
  39     def __init__(self):
  40         self._controller = None
  41
  42     async def __aenter__(self):
  43         self._controller = Controller()
  44         await self._controller.connect()
  45         return self._controller
  46
  47     async def __aexit__(self, exc_type, exc, tb):
  48         await self._controller.disconnect()
  49
  50
  51 class CleanModel():
  52     """
  53     Context manager that automatically connects to the currently active
  54     controller, adds a fresh model, returns the connection to that model,
  55     and automatically disconnects and cleans up the model.
  56
  57     The new model is also set as the current default for the controller
  58     connection.
  59     """
  60     def __init__(self, bakery_client=None):
  61         self._controller = None
  62         self._model = None
  63         self._model_uuid = None
  64         self._bakery_client = bakery_client
  65
  66     async def __aenter__(self):
  67         model_nonce = uuid.uuid4().hex[-4:]
  68         frame = inspect.stack()[1]
  69         test_name = frame.function.replace('_', '-')
  70         jujudata = TestJujuData()
  71         self._controller = Controller(
  72             jujudata=jujudata,
  73             bakery_client=self._bakery_client,
  74         )
  75         controller_name = jujudata.current_controller()
  76         user_name = jujudata.accounts()[controller_name]['user']
  77         await self._controller.connect(controller_name)
  78
  79         model_name = 'test-{}-{}-{}'.format(
  80             test_run_nonce,
  81             test_name,
  82             model_nonce,
  83         )
  84         self._model = await self._controller.add_model(model_name)
  85
  86         # Change the JujuData instance so that it will return the new
  87         # model as the current model name, so that we'll connect
  88         # to it by default.
  89         jujudata.set_model(
  90             controller_name,
  91             user_name + "/" + model_name,
  92             self._model.info.uuid,
  93         )
  94
  95         # save the model UUID in case test closes model
  96         self._model_uuid = self._model.info.uuid
  97
  98         return self._model
  99
 100     async def __aexit__(self, exc_type, exc, tb):
 101         await self._model.disconnect()
 102         await self._controller.destroy_model(self._model_uuid)
 103         await self._controller.disconnect()
 104
 105
 106 class TestJujuData(FileJujuData):
 107     def __init__(self):
 108         self.__controller_name = None
 109         self.__model_name = None
 110         self.__model_uuid = None
 111         super().__init__()
 112
 113     def set_model(self, controller_name, model_name, model_uuid):
 114         self.__controller_name = controller_name
 115         self.__model_name = model_name
 116         self.__model_uuid = model_uuid
 117
 118     def current_model(self, *args, **kwargs):
 119         return self.__model_name or super().current_model(*args, **kwargs)
 120
 121     def models(self):
 122         all_models = super().models()
 123         if self.__model_name is None:
 124             return all_models
 125         all_models.setdefault(self.__controller_name, {})
 126         all_models[self.__controller_name].setdefault('models', {})
 127         cmodels = all_models[self.__controller_name]['models']
 128         cmodels[self.__model_name] = {'uuid': self.__model_uuid}
 129         return all_models
 130
 131
 132 class AsyncMock(mock.MagicMock):
 133     async def __call__(self, *args, **kwargs):
 134         return super().__call__(*args, **kwargs)
 135
 136
 137 @contextmanager
 138 def patch_file(filename):
 139     """
 140     "Patch" a file so that its current contents are automatically restored
 141     when the context is exited.
 142     """
 143     filepath = Path(filename).expanduser()
 144     data = filepath.read_bytes()
 145     try:
 146         yield
 147     finally:
 148         filepath.write_bytes(data)