Add unit tests / small refactor

Signed-off-by: jamshale <[email protected]>

aries_cloudagent/wallet/anoncreds_upgrade.py

@@ -529,14 +529,16 @@ async def convert_records_to_anoncreds(profile) -> None:
 
 
 async def retry_converting_records(
-    profile: Profile, storage: BaseStorage, upgrading_record: StorageRecord, retry: int
+    profile: Profile, upgrading_record: StorageRecord, retry: int
 ) -> None:
     """Retry converting records to anoncreds."""
 
     async def clear_upgrade():
-        upgrade_singleton = UpgradeSingleton()
-        upgrade_singleton.remove_wallet(profile.name)
-        await storage.delete_record(upgrading_record)
+        async with profile.session() as session:
+            storage = session.inject(BaseStorage)
+            upgrade_singleton = UpgradeSingleton()
+            upgrade_singleton.remove_wallet(profile.name)
+            await storage.delete_record(upgrading_record)
 
     try:
         await convert_records_to_anoncreds(profile)
@@ -547,9 +549,7 @@ async def clear_upgrade():
         if retry < max_retries:
             LOGGER.info(f"Retry attempt {retry + 1} to upgrade wallet {profile.name}")
             await asyncio.sleep(1)
-            await retry_converting_records(
-                profile, storage, upgrading_record, retry + 1
-            )
+            await retry_converting_records(profile, upgrading_record, retry + 1)
         else:
             LOGGER.error(
                 f"""Failed to upgrade wallet: {profile.name} after 5 retries. 
@@ -560,7 +560,7 @@ async def clear_upgrade():
 
 
 async def upgrade_wallet_to_anoncreds(profile: Profile, is_subwallet=False) -> None:
-    """Get upgrading record and set singleton."""
+    """Get upgrading record and attempt to upgrade wallet to anoncreds."""
     async with profile.session() as session:
         storage = session.inject(BaseStorage)
         try:
@@ -581,10 +581,9 @@ async def upgrade_wallet_to_anoncreds(profile: Profile, is_subwallet=False) -> N
             )
             upgrade_singleton.remove_wallet(profile.name)
             await storage.delete_record(upgrading_record)
-            return
         except Exception as e:
             LOGGER.error(f"Error when upgrading wallet {profile.name} : {e} ")
-            await retry_converting_records(profile, storage, upgrading_record, 0)
+            await retry_converting_records(profile, upgrading_record, 0)
 
 
 async def set_storage_type_to_anoncreds(profile: Profile):

aries_cloudagent/wallet/tests/test_anoncreds_upgrade.py

@@ -7,8 +7,9 @@
 from ...multitenant.manager import MultitenantManager
 from ...storage.base import BaseStorage
 from ...storage.record import StorageRecord
-from ...storage.type import RECORD_TYPE_ACAPY_STORAGE_TYPE
+from ...storage.type import RECORD_TYPE_ACAPY_STORAGE_TYPE, RECORD_TYPE_ACAPY_UPGRADING
 from .. import anoncreds_upgrade
+from ..upgrade_singleton import UpgradeSingleton
 
 
 class TestAnoncredsUpgrade(IsolatedAsyncioTestCase):
@@ -21,6 +22,72 @@ def setUp(self) -> None:
             BaseMultitenantManager, mock.MagicMock(MultitenantManager, autospec=True)
         )
 
+    async def test_retry_converting_records(self):
+        with mock.patch.object(
+            anoncreds_upgrade, "convert_records_to_anoncreds", mock.CoroutineMock()
+        ) as mock_convert_records_to_anoncreds:
+            mock_convert_records_to_anoncreds.side_effect = [
+                Exception("Error"),
+                Exception("Error"),
+                None,
+            ]
+            async with self.profile.session() as session:
+                storage = session.inject(BaseStorage)
+                upgrading_record = StorageRecord(
+                    RECORD_TYPE_ACAPY_UPGRADING,
+                    "true",
+                )
+                await storage.add_record(upgrading_record)
+            await anoncreds_upgrade.retry_converting_records(
+                self.profile, upgrading_record, 0
+            )
+            upgrade_singleton = UpgradeSingleton()
+            assert mock_convert_records_to_anoncreds.called
+            assert mock_convert_records_to_anoncreds.call_count == 3
+            assert not session.profile.records
+            assert not upgrade_singleton
+
+    async def test_upgrade_wallet_to_anoncreds(self):
+        # upgrading record not present
+        await anoncreds_upgrade.upgrade_wallet_to_anoncreds(self.profile)
+
+        # upgrading record present
+        async with self.profile.session() as session:
+            storage = session.inject(BaseStorage)
+            await storage.add_record(
+                StorageRecord(
+                    RECORD_TYPE_ACAPY_UPGRADING,
+                    "true",
+                )
+            )
+            await anoncreds_upgrade.upgrade_wallet_to_anoncreds(self.profile)
+            # Upgrading record should be deleted
+            assert len(storage.profile.records) == 1
+            _, storage_type_record = next(iter(self.profile.records.items()))
+            assert storage_type_record.value == "askar-anoncreds"
+
+            upgrade_singleton = UpgradeSingleton()
+            assert not upgrade_singleton
+
+        # retry called on exception
+        with mock.patch.object(
+            anoncreds_upgrade,
+            "convert_records_to_anoncreds",
+            mock.CoroutineMock(side_effect=[Exception("Error")]),
+        ) as mock_convert_records_to_anoncreds, mock.patch.object(
+            anoncreds_upgrade, "retry_converting_records", mock.CoroutineMock()
+        ) as mock_retry_converting_records:
+            async with self.profile.session() as session:
+                storage = session.inject(BaseStorage)
+                await storage.add_record(
+                    StorageRecord(
+                        RECORD_TYPE_ACAPY_UPGRADING,
+                        "true",
+                    )
+                )
+            await anoncreds_upgrade.upgrade_wallet_to_anoncreds(self.profile)
+            assert mock_retry_converting_records.called
+
     async def test_set_storage_type_to_anoncreds_no_existing_record(self):
         await anoncreds_upgrade.set_storage_type_to_anoncreds(self.profile)
         _, storage_type_record = next(iter(self.profile.records.items()))

aries_cloudagent/wallet/tests/test_routes.py

@@ -1006,6 +1006,25 @@ async def test_rotate_did_keypair_x(self):
         with self.assertRaises(test_module.web.HTTPBadRequest):
             await test_module.wallet_rotate_did_keypair(self.request)
 
+    async def test_upgrade_anoncreds(self):
+        self.profile.settings["wallet.name"] = "test_wallet"
+        self.request.query = {"wallet_name": "not_test_wallet"}
+        with self.assertRaises(test_module.web.HTTPBadRequest):
+            await test_module.upgrade_anoncreds(self.request)
+
+        self.request.query = {"wallet_name": "not_test_wallet"}
+        self.profile.settings["wallet.type"] = "askar-anoncreds"
+        with self.assertRaises(test_module.web.HTTPBadRequest):
+            await test_module.upgrade_anoncreds(self.request)
+
+        self.request.query = {"wallet_name": "test_wallet"}
+        self.profile.settings["wallet.type"] = "askar"
+        result = await test_module.upgrade_anoncreds(self.request)
+        print(result)
+        _, upgrade_record = next(iter(self.profile.records.items()))
+        assert upgrade_record.type == "acapy_upgrading"
+        assert upgrade_record.value == "true"
+
     async def test_register(self):
         mock_app = mock.MagicMock()
         mock_app.add_routes = mock.MagicMock()