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.controller = None
  38         self.model = None
  39
  40     async def __aenter__(self):
  41         self.controller = Controller()
  42         await self.controller.connect_current()
  43
  44         model_name = 'model-{}'.format(uuid.uuid4())
  45         self.model = await self.controller.add_model(model_name)
  46
  47         # save the model UUID in case test closes model
  48         self.model_uuid = self.model.info.uuid
  49
  50         # Ensure that we connect to the new model by default.  This also
  51         # prevents failures if test was started with no current model.
  52         self._patch_cm = mock.patch.object(JujuData, 'current_model',
  53                                            return_value=model_name)
  54         self._patch_cm.start()
  55
  56         return self.model
  57
  58     async def __aexit__(self, exc_type, exc, tb):
  59         self._patch_cm.stop()
  60         await self.model.disconnect()
  61         await self.controller.destroy_model(self.model_uuid)
  62         await self.controller.disconnect()
  63
  64
  65 class AsyncMock(mock.MagicMock):
  66     async def __call__(self, *args, **kwargs):
  67         return super().__call__(*args, **kwargs)