From 46a9dccc8039763623519c05fee2fe638c8f07d6 Mon Sep 17 00:00:00 2001
From: jneo8 <james.lin@canonical.com>
Date: Tue, 11 Jun 2024 17:58:38 +0800
Subject: [PATCH] feat: Restart ceilometer-agent-compute after nova-compute
 upgrade

Fix: #427
---
 cou/utils/juju_utils.py             | 38 ++++++++++++++++++++++++++---
 tests/unit/utils/test_juju_utils.py | 29 +++++++++++++++++++---
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/cou/utils/juju_utils.py b/cou/utils/juju_utils.py
index 124198bf..047c728f 100644
--- a/cou/utils/juju_utils.py
+++ b/cou/utils/juju_utils.py
@@ -19,9 +19,9 @@
 import asyncio
 import logging
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Callable, Optional
+from typing import Any, Callable, List, Optional
 
 from juju.action import Action
 from juju.application import Application as JujuApplication
@@ -142,6 +142,22 @@ class Machine:
     az: Optional[str] = None  # simple deployments may not have azs
 
 
+@dataclass(frozen=True)
+class SubordinationUnit:
+    """Representation of a single unit of subordinate unit."""
+
+    name: str
+    workload_version: str
+
+    def __repr__(self) -> str:
+        """App representation.
+
+        :return: Name of the application
+        :rtype: str
+        """
+        return self.name
+
+
 @dataclass(frozen=True)
 class Unit:
     """Representation of a single unit of application."""
@@ -149,6 +165,7 @@ class Unit:
     name: str
     machine: Machine
     workload_version: str
+    subordinates: List[SubordinationUnit] = field(default_factory=lambda: [], compare=False)
 
     def __repr__(self) -> str:
         """App representation.
@@ -367,6 +384,13 @@ async def get_applications(self) -> dict[str, Application]:
         full_status = await self.get_status()
         machines = await self._get_machines()
 
+        for app, status in full_status.applications.items():
+            logger.debug(app)
+            for name, unit in status.units.items():
+                logger.debug(name)
+                logger.debug(unit)
+                logger.debug(unit.subordinates)
+
         return {
             app: Application(
                 name=app,
@@ -383,7 +407,15 @@ async def get_applications(self) -> dict[str, Application]:
                 series=status.series,
                 subordinate_to=status.subordinate_to,
                 units={
-                    name: Unit(name, machines[unit.machine], unit.workload_version)
+                    name: Unit(
+                        name,
+                        machines[unit.machine],
+                        unit.workload_version,
+                        [
+                            SubordinationUnit(subordinate_name, subordinate.workload_version)
+                            for subordinate_name, subordinate in unit.subordinates.items()
+                        ],
+                    )
                     for name, unit in status.units.items()
                 },
                 workload_version=status.workload_version,
diff --git a/tests/unit/utils/test_juju_utils.py b/tests/unit/utils/test_juju_utils.py
index 54aafad6..839f6429 100644
--- a/tests/unit/utils/test_juju_utils.py
+++ b/tests/unit/utils/test_juju_utils.py
@@ -625,10 +625,16 @@ def _generate_juju_machine(machine_id: str) -> MagicMock:
     return machine
 
 
-def _generate_unit_status(app: str, unit_id: int, machine_id: str) -> tuple[str, MagicMock]:
+def _generate_unit_status(
+    app: str,
+    unit_id: int,
+    machine_id: str,
+    subordinates: dict[str, MagicMock] = {},
+) -> tuple[str, MagicMock]:
     """Generate unit name and status."""
     status = MagicMock(spec_set=UnitStatus)()
     status.machine = machine_id
+    status.subordinates = subordinates
     return f"{app}/{unit_id}", status
 
 
@@ -678,7 +684,9 @@ async def test_get_applications(mock_get_machines, mock_get_status, mocked_model
     }
     exp_units_from_status = {
         "app1": dict([_generate_unit_status("app1", i, f"{i}") for i in range(3)]),
-        "app2": dict([_generate_unit_status("app2", 0, "0")]),
+        "app2": dict(
+            [_generate_unit_status("app2", 0, "0", dict([_generate_unit_status("app4", 0, "")]))]
+        ),
         "app3": dict([_generate_unit_status("app3", 0, "1")]),
         "app4": {},  # subordinate application has no units defined in juju status
     }
@@ -712,7 +720,18 @@ async def test_get_applications(mock_get_machines, mock_get_status, mocked_model
             series=status.series,
             subordinate_to=status.subordinate_to,
             units={
-                name: juju_utils.Unit(name, exp_machines[unit.machine], unit.workload_version)
+                name: juju_utils.Unit(
+                    name,
+                    exp_machines[unit.machine],
+                    unit.workload_version,
+                    [
+                        juju_utils.SubordinationUnit(
+                            subordinate_name,
+                            subordinate.workload_version,
+                        )
+                        for subordinate_name, subordinate in unit.subordinates.items()
+                    ],
+                )
                 for name, unit in exp_units_from_status[app].items()
             },
             workload_version=status.workload_version,
@@ -746,3 +765,7 @@ async def test_get_applications(mock_get_machines, mock_get_status, mocked_model
 def test_unit_repr():
     unit = juju_utils.Unit(name="foo/0", machine=MagicMock(), workload_version="1")
     assert repr(unit) == "foo/0"
+
+def test_suborinate_unit_repr():
+    unit = juju_utils.SubordinationUnit(name="foo/0", workload_version="1")
+    assert repr(unit) == "foo/0"