From 2b5394e4a795dab81e5f404ec44cc0acf42d0db0 Mon Sep 17 00:00:00 2001
From: wuwentao <wuwentao2000@126.com>
Date: Thu, 4 Jul 2024 15:50:13 +0000
Subject: [PATCH] fix: rollback get_keys with 3 keys

---
 midealocal/cloud.py                           |  5 +--
 tests/cloud_test.py                           | 38 +++++++++++++++++--
 ...et_keys.json => meijucloud_get_keys1.json} |  4 +-
 tests/responses/meijucloud_get_keys2.json     | 12 ++++++
 4 files changed, 50 insertions(+), 9 deletions(-)
 rename tests/responses/{meijucloud_get_keys.json => meijucloud_get_keys1.json} (61%)
 create mode 100644 tests/responses/meijucloud_get_keys2.json

diff --git a/midealocal/cloud.py b/midealocal/cloud.py
index 0ac9b098..f18dac4d 100644
--- a/midealocal/cloud.py
+++ b/midealocal/cloud.py
@@ -206,9 +206,8 @@ async def get_keys(self, appliance_id: int) -> dict[int, dict[str, Any]]:
                             "token": token["token"].lower(),
                             "key": token["key"].lower(),
                         }
-                break
-        if not result:
-            result.update(default_keys)
+        # add default key with method 1 and method 2 key
+        result.update(default_keys)
         return result
 
     @staticmethod
diff --git a/tests/cloud_test.py b/tests/cloud_test.py
index 565dba41..7e446ac0 100644
--- a/tests/cloud_test.py
+++ b/tests/cloud_test.py
@@ -121,7 +121,12 @@ async def test_meijucloud_get_keys(self) -> None:
         response = Mock()
         response.read = AsyncMock(
             side_effect=[
-                self.responses["meijucloud_get_keys.json"],
+                self.responses["meijucloud_get_keys1.json"],
+                self.responses["meijucloud_get_keys2.json"],
+                self.responses["meijucloud_get_keys1.json"],
+                self.responses["cloud_invalid_response.json"],
+                self.responses["cloud_invalid_response.json"],
+                self.responses["meijucloud_get_keys2.json"],
                 self.responses["cloud_invalid_response.json"],
                 self.responses["cloud_invalid_response.json"],
             ],
@@ -134,11 +139,36 @@ async def test_meijucloud_get_keys(self) -> None:
             password="password",
         )
         assert cloud is not None
-        keys: dict = await cloud.get_keys(100)
-        assert keys[1]["token"] == "returnedappliancetoken"
-        assert keys[1]["key"] == "returnedappliancekey"
 
+        # test method1 + method2 + default key
+        keys3: dict = await cloud.get_keys(100)
+        # test response token/key
+        assert keys3[1]["token"] == "method1_return_token1"
+        assert keys3[1]["key"] == "method1_return_key1"
+        assert keys3[2]["token"] == "method2_return_token2"
+        assert keys3[2]["key"] == "method2_return_key2"
+        # simple test default key with length
+        assert len(keys3) == 3
+
+        # test method1 + default key
+        keys1: dict = await cloud.get_keys(100)
+        # test response token/key
+        assert keys1[1]["token"] == "method1_return_token1"
+        assert keys1[1]["key"] == "method1_return_key1"
+        # simple test default key with length
+        assert len(keys1) == 2
+
+        # test method2 + default key
+        keys2: dict = await cloud.get_keys(100)
+        # test response token/key
+        assert keys2[2]["token"] == "method2_return_token2"
+        assert keys2[2]["key"] == "method2_return_key2"
+        # simple test default key with length
+        assert len(keys2) == 2
+
+        # test only default key
         keys = await cloud.get_keys(100)
+        assert len(keys) == 1
         assert keys == default_keys
 
     async def test_meijucloud_list_home(self) -> None:
diff --git a/tests/responses/meijucloud_get_keys.json b/tests/responses/meijucloud_get_keys1.json
similarity index 61%
rename from tests/responses/meijucloud_get_keys.json
rename to tests/responses/meijucloud_get_keys1.json
index d62cfbdf..550c1b50 100644
--- a/tests/responses/meijucloud_get_keys.json
+++ b/tests/responses/meijucloud_get_keys1.json
@@ -4,8 +4,8 @@
     "tokenlist": [
       {
         "udpId": "dec4da86e0aeefadde14a4e553680b9b",
-        "token": "returnedappliancetoken",
-        "key": "returnedappliancekey"
+        "token": "method1_return_token1",
+        "key": "method1_return_key1"
       }
     ]
   }
diff --git a/tests/responses/meijucloud_get_keys2.json b/tests/responses/meijucloud_get_keys2.json
new file mode 100644
index 00000000..5587b855
--- /dev/null
+++ b/tests/responses/meijucloud_get_keys2.json
@@ -0,0 +1,12 @@
+{
+  "code": 0,
+  "data": {
+    "tokenlist": [
+      {
+        "udpId": "b2dd199071d527a33c8239833a5bf5fb",
+        "token": "method2_return_token2",
+        "key": "method2_return_key2"
+      }
+    ]
+  }
+}