From 59a9b7f0cb1f7506466d2b705a43a976017da701 Mon Sep 17 00:00:00 2001 From: Liam Young Date: Wed, 28 Mar 2018 14:07:48 +0000 Subject: [PATCH] Add helper utils for interacting with juju environment --- test-requirements.txt | 1 + unit_tests/test_zaza_model.py | 132 +++++++++++++++++++++++ zaza/model.py | 191 ++++++++++++++++++++++++++++++++++ 3 files changed, 324 insertions(+) create mode 100644 unit_tests/test_zaza_model.py diff --git a/test-requirements.txt b/test-requirements.txt index 7d9f11e..9a731b5 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,3 +1,4 @@ +juju juju_wait PyYAML flake8>=2.2.4,<=3.5.0 diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py new file mode 100644 index 0000000..010c390 --- /dev/null +++ b/unit_tests/test_zaza_model.py @@ -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']) diff --git a/zaza/model.py b/zaza/model.py index a03eed5..c5fd41d 100644 --- a/zaza/model.py +++ b/zaza/model.py @@ -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.