modules/libjuju/tests/base.py

   1 import mock
   2 import subprocess
   3 import uuid
   4
   5 import pytest
   6
   7 from juju.controller import Controller
   8 from juju.client.connection import JujuData
   9
  10
  11 def is_bootstrapped():
  12     result = subprocess.run(['juju', 'switch'], stdout=subprocess.PIPE)
  13     return (
  14         result.returncode == 0 and
  15         len(result.stdout.decode().strip()) > 0)
  16
  17 bootstrapped = pytest.mark.skipif(
  18     not is_bootstrapped(),
  19     reason='bootstrapped Juju environment required')
  20
  21
  22 class CleanController():
  23     def __init__(self):
  24         self.controller = None
  25
  26     async def __aenter__(self):
  27         self.controller = Controller()
  28         await self.controller.connect_current()
  29         return self.controller
  30
  31     async def __aexit__(self, exc_type, exc, tb):
  32         await self.controller.disconnect()
  33
  34
  35 class CleanModel():
  36     def __init__(self):
  37         self.user_name = None
  38         self.controller = None
  39         self.controller_name = None
  40         self.model = None
  41         self.model_name = None
  42         self.model_uuid = None
  43
  44     async def __aenter__(self):
  45         self.controller = Controller()
  46         juju_data = JujuData()
  47         self.controller_name = juju_data.current_controller()
  48         self.user_name = juju_data.accounts()[self.controller_name]['user']
  49         await self.controller.connect_controller(self.controller_name)
  50
  51         self.model_name = 'test-{}'.format(uuid.uuid4())
  52         self.model = await self.controller.add_model(self.model_name)
  53
  54         # save the model UUID in case test closes model
  55         self.model_uuid = self.model.info.uuid
  56
  57         # Ensure that we connect to the new model by default.  This also
  58         # prevents failures if test was started with no current model.
  59         self._patch_cm = mock.patch.object(JujuData, 'current_model',
  60                                            return_value=self.model_name)
  61         self._patch_cm.start()
  62
  63         # Ensure that the models data includes this model, since it doesn't
  64         # get added to the client store by Controller.add_model().
  65         self._orig_models = JujuData().models
  66         self._patch_models = mock.patch.object(JujuData, 'models',
  67                                                side_effect=self._models)
  68         self._patch_models.start()
  69
  70         return self.model
  71
  72     def _models(self):
  73         result = self._orig_models()
  74         models = result[self.controller_name]['models']
  75         full_model_name = '{}/{}'.format(self.user_name, self.model_name)
  76         if full_model_name not in models:
  77             models[full_model_name] = {'uuid': self.model_uuid}
  78         return result
  79
  80     async def __aexit__(self, exc_type, exc, tb):
  81         self._patch_models.stop()
  82         self._patch_cm.stop()
  83         await self.model.disconnect()
  84         await self.controller.destroy_model(self.model_uuid)
  85         await self.controller.disconnect()
  86
  87
  88 class AsyncMock(mock.MagicMock):
  89     async def __call__(self, *args, **kwargs):
  90         return super().__call__(*args, **kwargs)