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
:
25 bootstrapped
= pytest
.mark
.skipif(
26 not is_bootstrapped(),
27 reason
='bootstrapped Juju environment required')
29 test_run_nonce
= uuid
.uuid4().hex[-4:]
32 class CleanController():
34 Context manager that automatically connects and disconnects from
35 the currently active controller.
37 Note: Unlike CleanModel, this will not create a new controller for you,
38 and an active controller must already be available.
41 self
._controller
= None
43 async def __aenter__(self
):
44 self
._controller
= Controller()
45 await self
._controller
.connect()
46 return self
._controller
48 async def __aexit__(self
, exc_type
, exc
, tb
):
49 await self
._controller
.disconnect()
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.
58 The new model is also set as the current default for the controller
61 def __init__(self
, bakery_client
=None):
62 self
._controller
= None
64 self
._model
_uuid
= None
65 self
._bakery
_client
= bakery_client
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(
74 bakery_client
=self
._bakery
_client
,
76 controller_name
= jujudata
.current_controller()
77 user_name
= jujudata
.accounts()[controller_name
]['user']
78 await self
._controller
.connect(controller_name
)
80 model_name
= 'test-{}-{}-{}'.format(
85 self
._model
= await self
._controller
.add_model(model_name
)
87 # Change the JujuData instance so that it will return the new
88 # model as the current model name, so that we'll connect
92 user_name
+ "/" + model_name
,
93 self
._model
.info
.uuid
,
96 # save the model UUID in case test closes model
97 self
._model
_uuid
= self
._model
.info
.uuid
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()
107 class TestJujuData(FileJujuData
):
109 self
.__controller
_name
= None
110 self
.__model
_name
= None
111 self
.__model
_uuid
= None
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
119 def current_model(self
, *args
, **kwargs
):
120 return self
.__model
_name
or super().current_model(*args
, **kwargs
)
123 all_models
= super().models()
124 if self
.__model
_name
is None:
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
}
133 class AsyncMock(mock
.MagicMock
):
134 async def __call__(self
, *args
, **kwargs
):
135 return super().__call
__(*args
, **kwargs
)
139 def patch_file(filename
):
141 "Patch" a file so that its current contents are automatically restored
142 when the context is exited.
144 filepath
= Path(filename
).expanduser()
145 data
= filepath
.read_bytes()
149 filepath
.write_bytes(data
)