Merge ~martin-kalcok/juju-verify:target-machine into ~canonical-solutions-engineering/juju-verify:master

[Martin Kalcok]

Allow targeting juju-verify on machines with -M/--machine option

Added option -M/--machines to juju-verify that allows targeting juju machines instead of Units.

This does not have impact on overall juju-verify functionality, it serves as a shortcut that removes the need to list every principal charm of the machine manually.

[🤖 Canonical IS Merge Bot (canonical-is-mergebot)]

This merge proposal is being monitored by mergebot. Change the status to Approved to merge.

[Alvaro Uria (aluria)]

All good. I've added a minor suggestion to use list comprehension instead of .append

[Alvaro Uria (aluria)]

+1, ty

review: Approve

[🤖 Canonical IS Merge Bot (canonical-is-mergebot)]

Change successfully merged at revision 12f01da7d8512b654b725c489c131757db214d3f

diff --git a/juju_verify/juju_verify.py b/juju_verify/juju_verify.py
index bb105c0..6469229 100644
--- a/juju_verify/juju_verify.py
+++ b/juju_verify/juju_verify.py
@@ -59,6 +59,18 @@ async def find_units(model: Model, units: List[str]) -> List[Unit]:
     return selected_units
 
 
+async def find_units_on_machine(model: Model, machines: List[str]) -> List[Unit]:
+    """Find all principal units that run on selected machines.
+
+    :param model: Juju model to search units in
+    :param machines: names of Juju machines on which to search units
+    :return: List of juju.Unit objects that match units running on the machines
+    """
+    return [unit for _, unit in model.units.items()
+            if unit.machine.entity_id in machines
+            and not unit.data.get("subordinate")]
+
+
 async def connect_model(model_name: Union[str, None]) -> Model:
     """Connect to a custom or default Juju model.
 
@@ -88,11 +100,15 @@ def parse_args() -> argparse.Namespace:
                         help='Connect to specific model.')
     parser.add_argument('check', choices=BaseVerifier.supported_checks(),
                         type=str.lower, help='Check to verify.')
-    parser.add_argument('units', nargs='+', help='Units to check.')
     parser.add_argument('-l', '--log-level', type=str.lower,
                         help='Set amount of displayed information',
                         default='info', choices=['trace', 'debug', 'info'])
 
+    target = parser.add_mutually_exclusive_group(required=True)
+    target.add_argument('--units', '-u', nargs='+', type=str,
+                        help='Units to check.')
+    target.add_argument('--machines', '-M', nargs='+', type=str,
+                        help='Check all units on the machine.')
     return parser.parse_args()
 
 
@@ -119,7 +135,15 @@ def main() -> None:
     args = parse_args()
     config_logger(args.log_level)
     model = loop.run(connect_model(args.model))
-    units = loop.run(find_units(model, args.units))
+    units: List[Unit] = []
+
+    if args.units:
+        units = loop.run(find_units(model, args.units))
+    elif args.machines:
+        units = loop.run(find_units_on_machine(model, args.machines))
+    else:
+        fail('juju-verify must target either juju units or juju machines')
+
     try:
         verifier = get_verifier(units)
         result = verifier.verify(args.check)
diff --git a/tests/unit/test_juju_verify.py b/tests/unit/test_juju_verify.py
index 5b938aa..6bdbbb2 100644
--- a/tests/unit/test_juju_verify.py
+++ b/tests/unit/test_juju_verify.py
@@ -99,6 +99,32 @@ async def test_find_units(model, all_units, fail):
     fail.assert_called_once_with(expected_message)
 
 
[e-mail]
+async def test_find_units_on_machine(model, all_units):
+    """Test that find_units_on_machine function returns correct units."""
+    machine_1_name = '0'
+    machine_2_name = '1'
+
+    machine_1 = MagicMock()
+    machine_1.entity_id = machine_1_name
+    machine_2 = MagicMock()
+    machine_2.entity_id = machine_2_name
+
+    machine_1_units = all_units[:3]
+    machine_2_units = all_units[3:]
+
+    for unit_name, unit in model.units.items():
+        if unit_name in machine_1_units:
+            unit.machine = machine_1
+        elif unit_name in machine_2_units:
+            unit.machine = machine_2
+
+    found_units = await juju_verify.find_units_on_machine(model,
+                                                          [machine_1_name])
+
+    assert machine_1_units == [unit.entity_id for unit in found_units]
+
+
 @pytest.mark.parametrize('log_level, log_constant',
                          [('debug', logging.DEBUG),
                           ('DEBUG', logging.DEBUG),
@@ -148,13 +174,14 @@ def test_parse_args(mocker):
     parser.assert_called_once()
 
 
-def test_main_entrypoint(mocker):
-    """Verify workflow if the main entrypoint."""
+def test_main_entrypoint_target_units(mocker):
+    """Verify workflow of the main entrypoint when script targets units."""
     args = MagicMock()
     args.log_level = 'info'
     args.model = None
     args.check = 'shutdown'
     args.units = ['nova-compute/0']
+    args.machines = None
 
     result = Result(True, 'Passed')
     verifier = MagicMock()
@@ -175,6 +202,56 @@ def test_main_entrypoint(mocker):
     logger.info.assert_called_with(result.format())
 
 
+def test_main_entrypoint_target_machine(mocker):
+    """Verify workflow of the main entrypoint when script targets machines."""
+    args = MagicMock()
+    args.log_level = 'info'
+    args.model = None
+    args.check = 'shutdown'
+    args.machines = ['0']
+    args.units = None
+    expected_units = ['nova-compute/0']
+
+    result = Result(True, 'Passed')
+    verifier = MagicMock()
+    verifier.verify.return_value = result
+
+    mocker.patch.object(juju_verify, 'parse_args').return_value = args
+    mocker.patch.object(juju_verify, 'get_verifier').return_value = verifier
+    mocker.patch.object(juju_verify, 'loop')
+    mocker.patch.object(juju_verify, 'connect_model', new_callable=MagicMock())
+    mocker.patch.object(juju_verify, 'find_units_on_machine',
+                        new_callable=MagicMock()).return_value = expected_units
+    logger = mocker.patch.object(juju_verify, 'logger')
+
+    juju_verify.main()
+
+    juju_verify.connect_model.assert_called_with(args.model)
+    juju_verify.find_units_on_machine.assert_called_with(ANY, args.machines)
+    verifier.verify.asssert_called_with(args.check)
+    logger.info.assert_called_with(result.format())
+
+
+def test_main_entrypoint_no_target_fail(mocker, fail):
+    """Test that main fails if not target (units/machines) is specified."""
+    args = MagicMock()
+    args.log_level = 'info'
+    args.model = None
+    args.check = 'shutdown'
+    args.machines = None
+    args.units = None
+
+    expected_msg = 'juju-verify must target either juju units or juju machines'
+
+    mocker.patch.object(juju_verify, 'parse_args').return_value = args
+    mocker.patch.object(juju_verify, 'connect_model', new_callable=MagicMock())
+    mocker.patch.object(juju_verify, 'loop')
+
+    juju_verify.main()
+
+    fail.assert_any_call(expected_msg)
+
+
 @pytest.mark.parametrize('error, error_msg',
                          [(CharmException, 'Charm failure'),
                           (VerificationError, 'Verification Failure'),