4 from contextlib
import contextmanager
5 from pathlib
import Path
8 from juju
.client
.jujudata
import FileJujuData
9 from juju
.controller
import Controller
14 def is_bootstrapped():
16 result
= subprocess
.run(['juju', 'switch'], stdout
=subprocess
.PIPE
)
18 result
.returncode
== 0 and
19 len(result
.stdout
.decode().strip()) > 0)
20 except FileNotFoundError
:
24 bootstrapped
= pytest
.mark
.skipif(
25 not is_bootstrapped(),
26 reason
='bootstrapped Juju environment required')
28 test_run_nonce
= uuid
.uuid4().hex[-4:]
31 class CleanController():
33 Context manager that automatically connects and disconnects from
34 the currently active controller.
36 Note: Unlike CleanModel, this will not create a new controller for you,
37 and an active controller must already be available.
40 self
._controller
= None
42 async def __aenter__(self
):
43 self
._controller
= Controller()
44 await self
._controller
.connect()
45 return self
._controller
47 async def __aexit__(self
, exc_type
, exc
, tb
):
48 await self
._controller
.disconnect()
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.
57 The new model is also set as the current default for the controller
60 def __init__(self
, bakery_client
=None):
61 self
._controller
= None
63 self
._model
_uuid
= None
64 self
._bakery
_client
= bakery_client
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(
73 bakery_client
=self
._bakery
_client
,
75 controller_name
= jujudata
.current_controller()
76 user_name
= jujudata
.accounts()[controller_name
]['user']
77 await self
._controller
.connect(controller_name
)
79 model_name
= 'test-{}-{}-{}'.format(
84 self
._model
= await self
._controller
.add_model(model_name
)
86 # Change the JujuData instance so that it will return the new
87 # model as the current model name, so that we'll connect
91 user_name
+ "/" + model_name
,
92 self
._model
.info
.uuid
,
95 # save the model UUID in case test closes model
96 self
._model
_uuid
= self
._model
.info
.uuid
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()
106 class TestJujuData(FileJujuData
):
108 self
.__controller
_name
= None
109 self
.__model
_name
= None
110 self
.__model
_uuid
= None
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
118 def current_model(self
, *args
, **kwargs
):
119 return self
.__model
_name
or super().current_model(*args
, **kwargs
)
122 all_models
= super().models()
123 if self
.__model
_name
is None:
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
}
132 class AsyncMock(mock
.MagicMock
):
133 async def __call__(self
, *args
, **kwargs
):
134 return super().__call
__(*args
, **kwargs
)
138 def patch_file(filename
):
140 "Patch" a file so that its current contents are automatically restored
141 when the context is exited.
143 filepath
= Path(filename
).expanduser()
144 data
= filepath
.read_bytes()
148 filepath
.write_bytes(data
)