Add helper utils for interacting with juju environment
This commit is contained in:
Liam Young
@@ -1,3 +1,4 @@
|
||||
juju
|
||||
juju_wait
|
||||
PyYAML
|
||||
flake8>=2.2.4,<=3.5.0
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
import functools
|
||||
import mock
|
||||
import zaza.model as model
|
||||
import unit_tests.utils as ut_utils
|
||||
from juju import loop
|
||||
|
||||
|
||||
class TestModel(ut_utils.BaseTestCase):
|
||||
|
||||
def setUp(self):
|
||||
super(TestModel, self).setUp()
|
||||
self.unit1 = mock.MagicMock()
|
||||
self.unit1.public_address = 'ip1'
|
||||
self.unit1.name = 'app/2'
|
||||
self.unit1.entity_id = 'app/2'
|
||||
self.unit1.machine = 'machine3'
|
||||
self.unit2 = mock.MagicMock()
|
||||
self.unit2.public_address = 'ip2'
|
||||
self.unit2.name = 'app/4'
|
||||
self.unit2.entity_id = 'app/4'
|
||||
self.unit2.machine = 'machine7'
|
||||
self.units = [self.unit1, self.unit2]
|
||||
_units = mock.MagicMock()
|
||||
_units.units = self.units
|
||||
self.mymodel = mock.MagicMock()
|
||||
self.mymodel.applications = {
|
||||
'app': _units
|
||||
}
|
||||
self.Model_mock = mock.MagicMock()
|
||||
|
||||
async def _connect_model(model_name):
|
||||
return model_name
|
||||
|
||||
async def _disconnect():
|
||||
return
|
||||
self.Model_mock.connect_model.side_effect = _connect_model
|
||||
self.Model_mock.disconnect.side_effect = _disconnect
|
||||
self.Model_mock.applications = self.mymodel.applications
|
||||
|
||||
def test_run_in_model(self):
|
||||
self.patch_object(model, 'Model')
|
||||
|
||||
async def _test_func(arg):
|
||||
return arg * 2
|
||||
self.Model.return_value = self.Model_mock
|
||||
func = functools.partial(_test_func, 'hello')
|
||||
out = loop.run(
|
||||
model.run_in_model(
|
||||
'mymodel',
|
||||
func,
|
||||
awaitable=True))
|
||||
self.assertEqual(out, 'hellohello')
|
||||
|
||||
def test_run_in_model_not_awaitable(self):
|
||||
self.patch_object(model, 'Model')
|
||||
|
||||
def _test_func(arg):
|
||||
return arg * 3
|
||||
self.Model.return_value = self.Model_mock
|
||||
func = functools.partial(_test_func, 'hello')
|
||||
out = loop.run(
|
||||
model.run_in_model(
|
||||
'mymodel',
|
||||
func,
|
||||
awaitable=False))
|
||||
self.assertEqual(out, 'hellohellohello')
|
||||
|
||||
def test_run_in_model_add_model_arg(self):
|
||||
self.patch_object(model, 'Model')
|
||||
|
||||
def _test_func(arg, model):
|
||||
return model
|
||||
self.Model.return_value = self.Model_mock
|
||||
func = functools.partial(_test_func, 'hello')
|
||||
out = loop.run(
|
||||
model.run_in_model(
|
||||
'mymodel',
|
||||
func,
|
||||
add_model_arg=True,
|
||||
awaitable=False))
|
||||
self.assertEqual(out, self.Model_mock)
|
||||
|
||||
def test_scp_to_unit(self):
|
||||
self.patch_object(model, 'Model')
|
||||
self.patch_object(model, 'get_unit_from_name')
|
||||
unit_mock = mock.MagicMock()
|
||||
self.get_unit_from_name.return_value = unit_mock
|
||||
self.Model.return_value = self.Model_mock
|
||||
model.scp_to_unit('app/1', 'modelname', '/tmp/src', '/tmp/dest')
|
||||
unit_mock.scp_to.assert_called_once_with(
|
||||
'/tmp/src', '/tmp/dest', proxy=False, scp_opts='', user='ubuntu')
|
||||
|
||||
def test_scp_from_unit(self):
|
||||
self.patch_object(model, 'Model')
|
||||
self.patch_object(model, 'get_unit_from_name')
|
||||
unit_mock = mock.MagicMock()
|
||||
self.get_unit_from_name.return_value = unit_mock
|
||||
self.Model.return_value = self.Model_mock
|
||||
model.scp_from_unit('app/1', 'modelname', '/tmp/src', '/tmp/dest')
|
||||
unit_mock.scp_from.assert_called_once_with(
|
||||
'/tmp/src', '/tmp/dest', proxy=False, scp_opts='', user='ubuntu')
|
||||
|
||||
def test_get_units(self):
|
||||
self.patch_object(model, 'Model')
|
||||
self.Model.return_value = self.Model_mock
|
||||
self.assertEqual(
|
||||
model.get_units('modelname', 'app'),
|
||||
self.units)
|
||||
|
||||
def test_get_machines(self):
|
||||
self.patch_object(model, 'Model')
|
||||
self.Model.return_value = self.Model_mock
|
||||
self.assertEqual(
|
||||
model.get_machines('modelname', 'app'),
|
||||
['machine3', 'machine7'])
|
||||
|
||||
def test_get_first_unit_name(self):
|
||||
self.patch_object(model, 'get_units')
|
||||
self.get_units.return_value = self.units
|
||||
self.assertEqual(
|
||||
model.get_first_unit_name('model', 'app'),
|
||||
'app/2')
|
||||
|
||||
def test_get_unit_from_name(self):
|
||||
self.assertEqual(
|
||||
model.get_unit_from_name('app/4', self.mymodel),
|
||||
self.unit2)
|
||||
|
||||
def test_get_app_ips(self):
|
||||
self.patch_object(model, 'get_units')
|
||||
self.get_units.return_value = self.units
|
||||
self.assertEqual(model.get_app_ips('model', 'app'), ['ip1', 'ip2'])
|
||||
+191
@@ -1,3 +1,5 @@
|
||||
import functools
|
||||
|
||||
from juju import loop
|
||||
from juju.model import Model
|
||||
|
||||
@@ -16,6 +18,195 @@ async def deployed(filter=None):
|
||||
await model.disconnect()
|
||||
|
||||
|
||||
def get_unit_from_name(unit_name, model):
|
||||
"""Return the units that corresponds to the name in the given model
|
||||
|
||||
:param unit_name: Name of unit to match
|
||||
:type unit_name: str
|
||||
:param model: Model to perform lookup in
|
||||
:type model: juju.model.Model
|
||||
:returns: Unit matching given name
|
||||
:rtype: juju.unit.Unit or None
|
||||
"""
|
||||
app = unit_name.split('/')[0]
|
||||
unit = None
|
||||
for u in model.applications[app].units:
|
||||
if u.entity_id == unit_name:
|
||||
unit = u
|
||||
break
|
||||
else:
|
||||
raise Exception
|
||||
return unit
|
||||
|
||||
|
||||
async def run_in_model(model_name, f, add_model_arg=False, awaitable=True):
|
||||
"""Run the given function in the model matching the model_name
|
||||
|
||||
:param model_name: Name of model to run function in
|
||||
:type model_name: str
|
||||
:param f: Function to run with given moel in focus
|
||||
:type f: functools.partial
|
||||
:param add_model_arg: Whether to add kwarg pointing at model to the given
|
||||
function before running it
|
||||
:type add_model_arg: boolean
|
||||
:param awaitable: Whether f is awaitable
|
||||
:type awaitable: boolean
|
||||
:returns: Output of f
|
||||
:rtype: Unknown, depends on the passed in function
|
||||
"""
|
||||
model = Model()
|
||||
await model.connect_model(model_name)
|
||||
output = None
|
||||
try:
|
||||
if add_model_arg:
|
||||
f.keywords.update(model=model)
|
||||
if awaitable:
|
||||
output = await f()
|
||||
else:
|
||||
output = f()
|
||||
finally:
|
||||
# Disconnect from the api server and cleanup.
|
||||
await model.disconnect()
|
||||
return output
|
||||
|
||||
|
||||
def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu',
|
||||
proxy=False, scp_opts=''):
|
||||
"""Transfer files from to unit_name in model_name.
|
||||
|
||||
:param unit_name: Name of unit to scp to
|
||||
:type unit_name: str
|
||||
:param model_name: Name of model unit is in
|
||||
:type model_name: str
|
||||
:param source: Local path of file(s) to transfer
|
||||
:type source: str
|
||||
:param destination: Remote destination of transferred files
|
||||
:type source: str
|
||||
:param user: Remote username
|
||||
:type source: str
|
||||
:param proxy: Proxy through the Juju API server
|
||||
:type proxy: bool
|
||||
:param scp_opts: Additional options to the scp command
|
||||
:type scp_opts: str
|
||||
"""
|
||||
async def _scp_to_unit(unit_name, source, destination, user, proxy,
|
||||
scp_opts, model):
|
||||
unit = get_unit_from_name(unit_name, model)
|
||||
await unit.scp_to(source, destination, user=user, proxy=proxy,
|
||||
scp_opts=scp_opts)
|
||||
scp_func = functools.partial(
|
||||
_scp_to_unit,
|
||||
unit_name,
|
||||
source,
|
||||
destination,
|
||||
user=user,
|
||||
proxy=proxy,
|
||||
scp_opts=scp_opts)
|
||||
loop.run(
|
||||
run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True))
|
||||
|
||||
|
||||
def scp_from_unit(unit_name, model_name, source, destination, user='ubuntu',
|
||||
proxy=False, scp_opts=''):
|
||||
"""Transfer files from to unit_name in model_name.
|
||||
|
||||
:param unit_name: Name of unit to scp from
|
||||
:type unit_name: str
|
||||
:param model_name: Name of model unit is in
|
||||
:type model_name: str
|
||||
:param source: Remote path of file(s) to transfer
|
||||
:type source: str
|
||||
:param destination: Local destination of transferred files
|
||||
:type source: str
|
||||
:param user: Remote username
|
||||
:type source: str
|
||||
:param proxy: Proxy through the Juju API server
|
||||
:type proxy: bool
|
||||
:param scp_opts: Additional options to the scp command
|
||||
:type scp_opts: str
|
||||
"""
|
||||
async def _scp_from_unit(unit_name, source, destination, user, proxy,
|
||||
scp_opts, model):
|
||||
unit = get_unit_from_name(unit_name, model)
|
||||
await unit.scp_from(source, destination, user=user, proxy=proxy,
|
||||
scp_opts=scp_opts)
|
||||
scp_func = functools.partial(
|
||||
_scp_from_unit,
|
||||
unit_name,
|
||||
source,
|
||||
destination,
|
||||
user=user,
|
||||
proxy=proxy,
|
||||
scp_opts=scp_opts)
|
||||
loop.run(
|
||||
run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True))
|
||||
|
||||
|
||||
def get_units(model_name, application_name):
|
||||
"""Return all the units of a given application
|
||||
|
||||
:param model_name: Name of model to query.
|
||||
:type model_name: str
|
||||
:param application_name: Name of application to retrieve units for
|
||||
:type application_name: str
|
||||
|
||||
:returns: List of juju units
|
||||
:rtype: [juju.unit.Unit, juju.unit.Unit,...]
|
||||
"""
|
||||
async def _get_units(application_name, model):
|
||||
return model.applications[application_name].units
|
||||
f = functools.partial(_get_units, application_name)
|
||||
return loop.run(run_in_model(model_name, f, add_model_arg=True))
|
||||
|
||||
|
||||
def get_machines(model_name, application_name):
|
||||
"""Return all the machines of a given application
|
||||
|
||||
:param model_name: Name of model to query.
|
||||
:type model_name: str
|
||||
:param application_name: Name of application to retrieve units for
|
||||
:type application_name: str
|
||||
|
||||
:returns: List of juju machines
|
||||
:rtype: [juju.machine.Machine, juju.machine.Machine,...]
|
||||
"""
|
||||
async def _get_machines(application_name, model):
|
||||
machines = []
|
||||
for unit in model.applications[application_name].units:
|
||||
machines.append(unit.machine)
|
||||
return machines
|
||||
f = functools.partial(_get_machines, application_name)
|
||||
return loop.run(run_in_model(model_name, f, add_model_arg=True))
|
||||
|
||||
|
||||
def get_first_unit_name(model_name, application_name):
|
||||
"""Return name of lowest numbered unit of given application
|
||||
|
||||
:param model_name: Name of model to query.
|
||||
:type model_name: str
|
||||
:param application_name: Name of application
|
||||
:type application_name: str
|
||||
|
||||
:returns: Name of lowest numbered unit
|
||||
:rtype: str
|
||||
"""
|
||||
return get_units(model_name, application_name)[0].name
|
||||
|
||||
|
||||
def get_app_ips(model_name, application_name):
|
||||
"""Return public address of all units of an application
|
||||
|
||||
:param model_name: Name of model to query.
|
||||
:type model_name: str
|
||||
:param application_name: Name of application
|
||||
:type application_name: str
|
||||
|
||||
:returns: List of ip addresses
|
||||
:rtype: [str, str,...]
|
||||
"""
|
||||
return [u.public_address for u in get_units(model_name, application_name)]
|
||||
|
||||
|
||||
def main():
|
||||
# Run the deploy coroutine in an asyncio event loop, using a helper
|
||||
# that abstracts loop creation and teardown.
|
||||
|
||||
Reference in New Issue
Block a user