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