From 22922d429e6e6e92765e2331db57dd17c5a00eda Mon Sep 17 00:00:00 2001
From: Dima Tisnek <dima.tisnek@canonical.com>
Date: Thu, 3 Oct 2024 17:01:50 +0900
Subject: [PATCH] chore: more hookup

---
 tests/unit/test_wait_for_idle.py | 64 ++++++++++++++++++++++++--------
 1 file changed, 48 insertions(+), 16 deletions(-)

diff --git a/tests/unit/test_wait_for_idle.py b/tests/unit/test_wait_for_idle.py
index e0139e61c..4ec750ec5 100644
--- a/tests/unit/test_wait_for_idle.py
+++ b/tests/unit/test_wait_for_idle.py
@@ -4,8 +4,9 @@
 
 import json
 import pytest
-from typing import Dict, List
+from typing import Any, Dict, List
 from typing import reveal_type as reveal_type
+from unittest.mock import Mock
 
 from juju.application import Application
 from juju.client.facade import _convert_response
@@ -45,15 +46,25 @@ def test_foo(full_status_response: dict):
 
 class ModelFake(Model):
     _applications: Dict[str, Application]
+    _response: Dict
 
     @property
     def applications(self) -> Dict[str, Application]:
         return self._applications
-    
+
     def __init__(self):
         super().__init__()
         self._applications = {}
 
+    def connection(self):
+        rv = Mock()
+        rv.facades = {"Client": 6}  # Must match juju.client.connection.client_facades
+        rv.rpc = self.rpc
+        return rv
+
+    async def rpc(self, msg, encoder=None):
+        return self._response
+
 
 class ApplicationFake(Application):
     _status: str = ""
@@ -70,19 +81,40 @@ def status_message(self) -> str:
 
     @property
     def units(self) -> List[Unit]:
+        1/0
         return self._units
 
-async def test_something(full_status_response):
-    m = Model()
-    rv = await m._legacy_check_idle(
-            apps=[],
-            raise_on_error=False,
-            raise_on_blocked=False,
-            status=None,
-            wait_for_at_least_units=None,
-            wait_for_exact_units=None,
-            timeout=100,
-            idle_period=100,
-            _wait_for_units=1,
-            )
-    assert rv == 42
+
+@pytest.fixture
+def kwargs() -> Dict[str, Any]:
+    return dict(
+        apps=["hexanator", "grafana-agent-k8s"],
+        raise_on_error=False,
+        raise_on_blocked=False,
+        status=None,
+        wait_for_at_least_units=None,
+        wait_for_exact_units=None,
+        timeout=100,
+        idle_period=100,
+        _wait_for_units=1,
+    )
+
+
+async def test_something(full_status_response, kwargs: Dict[str, Any]):
+    m = ModelFake()
+    m._response = full_status_response
+    m._applications["hexanator"] = ApplicationFake("hexanator", m)
+    m._applications["grafana-agent-k8s"] = ApplicationFake("grafana-agent-k8s", m)
+
+    try:
+        idle = await m._check_idle(**kwargs)
+    except Exception as e:
+        idle = e
+
+    try:
+        legacy = await m._legacy_check_idle(**kwargs)
+    except Exception as e:
+        raise
+        legacy = e
+
+    assert idle == legacy