From 226993488c1da529631f70c9bd6e366592182921 Mon Sep 17 00:00:00 2001
From: Dima Tisnek <dima.tisnek@canonical.com>
Date: Wed, 2 Oct 2024 14:36:09 +0900
Subject: [PATCH] chore: refactor jrpc response type coercion for testing

---
 juju/client/facade.py | 59 +++++++++++++++++++++++++++----------------
 juju/model.py         |  9 +++++++
 2 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/juju/client/facade.py b/juju/client/facade.py
index 57a5970e..b2bf159c 100644
--- a/juju/client/facade.py
+++ b/juju/client/facade.py
@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import argparse
 import builtins
@@ -13,7 +14,7 @@
 from collections import defaultdict
 from glob import glob
 from pathlib import Path
-from typing import Any, Mapping, Sequence, TypeVar
+from typing import overload, Any, Dict, Mapping, Optional, Sequence, TypeVar, Type as TypingType
 
 import typing_inspect
 
@@ -482,31 +483,42 @@ def ReturnMapping(cls):
     def decorator(f):
         @functools.wraps(f)
         async def wrapper(*args, **kwargs):
-            nonlocal cls
             reply = await f(*args, **kwargs)
-            if cls is None:
-                return reply
-            if 'error' in reply:
+            return _convert_response(reply, cls=cls)
+        return wrapper
+    return decorator
+
+
+@overload
+def _convert_response(response: Dict, *, cls: TypingType[SomeType]) -> SomeType: ...
+
+
+@overload
+def _convert_response(response: Dict, *, cls: None) -> Dict: ...
+
+
+def _convert_response(response: dict, *, cls: Optional[TypingType[Type]]):
+    if cls is None:
+        return response
+    if 'error' in response:
+        cls = CLASSES['Error']
+    if typing_inspect.is_generic_type(cls) and issubclass(typing_inspect.get_origin(cls), Sequence):
+        parameters = typing_inspect.get_parameters(cls)
+        result = []
+        item_cls = parameters[0]
+        for item in response:
+            result.append(item_cls.from_json(item))
+            """
+            if 'error' in item:
                 cls = CLASSES['Error']
-            if typing_inspect.is_generic_type(cls) and issubclass(typing_inspect.get_origin(cls), Sequence):
-                parameters = typing_inspect.get_parameters(cls)
-                result = []
-                item_cls = parameters[0]
-                for item in reply:
-                    result.append(item_cls.from_json(item))
-                    """
-                    if 'error' in item:
-                        cls = CLASSES['Error']
-                    else:
-                        cls = item_cls
-                    result.append(cls.from_json(item))
-                    """
             else:
-                result = cls.from_json(reply['response'])
+                cls = item_cls
+            result.append(cls.from_json(item))
+            """
+    else:
+        result = cls.from_json(response['response'])
 
-            return result
-        return wrapper
-    return decorator
+    return result
 
 
 def makeFunc(cls, name, description, params, result, _async=True):
@@ -737,6 +749,9 @@ def get(self, key, default=None):
         return getattr(self, attr, default)
 
 
+SomeType = TypeVar("SomeType", bound=Type)
+
+
 class Schema(dict):
     def __init__(self, schema):
         self.name = schema['Name']
diff --git a/juju/model.py b/juju/model.py
index 83abf5b1..6535243a 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -2849,6 +2849,15 @@ async def _get_source_api(self, url):
         await controller.connect(controller_name=controller_name)
         return controller
 
+    # FIXME scraped from charm collection
+    # apps
+    # - most common: [explicit, list, of, apps]
+    # - rare: implicit None
+    # status
+    # - common implicit None
+    # - common "active"
+    # - rare "blocked"
+    # - very rare "unknown"
     async def wait_for_idle(self,
         apps: Optional[List[str]] = None,
         raise_on_error=True,