From 92138541bf0b1bf12819c9f94d68658657b656f1 Mon Sep 17 00:00:00 2001
From: Andrew Bulat <andrii.bulat@gmail.com>
Date: Fri, 14 Feb 2025 15:01:37 +0000
Subject: [PATCH] Allow a string value to be passed to `Capability` constructor

This implements RSA9f [1] and TK2b [2], where a capability JSON text can
be provided for a `Auth#createTokenRequest` function and when passing
`TokenParams` object

Resolves #579

[1] https://sdk.ably.com/builds/ably/specification/main/features/#RSA9f
[2] https://sdk.ably.com/builds/ably/specification/main/features/#TK2b
---
 ably/rest/auth.py                     |  2 +-
 ably/types/capability.py              | 19 ++++++++++++++-----
 ably/types/tokendetails.py            |  8 +-------
 test/ably/rest/restcapability_test.py | 14 ++++++++++++++
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/ably/rest/auth.py b/ably/rest/auth.py
index 06af2438..ab255a3e 100644
--- a/ably/rest/auth.py
+++ b/ably/rest/auth.py
@@ -229,7 +229,7 @@ async def request_token(self, token_params: Optional[dict] = None,
         log.debug("Token: %s" % str(response_dict.get("token")))
         return TokenDetails.from_dict(response_dict)
 
-    async def create_token_request(self, token_params: Optional[dict] = None, key_name: Optional[str] = None,
+    async def create_token_request(self, token_params: Optional[dict | str] = None, key_name: Optional[str] = None,
                                    key_secret: Optional[str] = None, query_time=None):
         token_params = token_params or {}
         token_request = {}
diff --git a/ably/types/capability.py b/ably/types/capability.py
index 5d209d7c..0c35940e 100644
--- a/ably/types/capability.py
+++ b/ably/types/capability.py
@@ -1,4 +1,5 @@
 from collections.abc import MutableMapping
+from typing import Optional, Union
 import json
 import logging
 
@@ -7,11 +8,19 @@
 
 
 class Capability(MutableMapping):
-    def __init__(self, obj=None):
-        if obj is None:
-            obj = {}
-        self.__dict = dict(obj)
-        for k, v in obj.items():
+    def __init__(self, capability: Optional[Union[dict, str]] = None):
+        # RSA9f: provided capability can be a JSON string
+        if capability and isinstance(capability, str):
+            try:
+                capability = json.loads(capability)
+            except json.JSONDecodeError:
+                capability = json.loads(capability.replace("'", '"'))
+
+        if capability is None:
+            capability = {}
+
+        self.__dict = dict(capability)
+        for k, v in capability.items():
             self[k] = v
 
     def __eq__(self, other):
diff --git a/ably/types/tokendetails.py b/ably/types/tokendetails.py
index f3b79e47..771b29ec 100644
--- a/ably/types/tokendetails.py
+++ b/ably/types/tokendetails.py
@@ -20,13 +20,7 @@ def __init__(self, token=None, expires=None, issued=0,
             self.__expires = expires
         self.__token = token
         self.__issued = issued
-        if capability and isinstance(capability, str):
-            try:
-                self.__capability = Capability(json.loads(capability))
-            except json.JSONDecodeError:
-                self.__capability = Capability(json.loads(capability.replace("'", '"')))
-        else:
-            self.__capability = Capability(capability or {})
+        self.__capability = Capability(capability or {})
         self.__client_id = client_id
 
     @property
diff --git a/test/ably/rest/restcapability_test.py b/test/ably/rest/restcapability_test.py
index f7c761ab..cb74ae8e 100644
--- a/test/ably/rest/restcapability_test.py
+++ b/test/ably/rest/restcapability_test.py
@@ -240,3 +240,17 @@ async def test_invalid_capabilities_3(self):
         the_exception = excinfo.value
         assert 400 == the_exception.status_code
         assert 40000 == the_exception.code
+
+    @dont_vary_protocol
+    def test_capability_from_string(self):
+        capability_from_str = Capability('{"cansubscribe":["subscribe"]}')
+        capability_from_str_single_quote = Capability('{\'cansubscribe\':[\'subscribe\']}')
+
+        capability_from_dict = Capability({
+            "cansubscribe": ["subscribe"]
+        })
+
+        assert capability_from_str == capability_from_dict, "Unexpected Capability constructed from string"
+        assert (
+            capability_from_str_single_quote == capability_from_dict
+        ), "Unexpected Capability constructed from string"