Fix upgrade middleware unit test

Signed-off-by: jamshale <[email protected]>
jamshale · Apr 11, 2024 · a29ccde · a29ccde
1 parent dd4ee19
commit a29ccde
Showing 1 changed file with 17 additions and 6 deletions.
diff --git a/aries_cloudagent/admin/tests/test_admin_server.py b/aries_cloudagent/admin/tests/test_admin_server.py
@@ -14,6 +14,9 @@
 from ...core.goal_code_registry import GoalCodeRegistry
 from ...core.in_memory import InMemoryProfile
 from ...core.protocol_registry import ProtocolRegistry
+from ...storage.base import BaseStorage
+from ...storage.record import StorageRecord
+from ...storage.type import RECORD_TYPE_ACAPY_UPGRADING
 from ...utils.stats import Collector
 from ...utils.task_queue import TaskQueue
 from .. import server as test_module
@@ -479,9 +482,8 @@ async def test_server_health_state(self):
         await server.stop()
 
     async def test_upgrade_middleware(self):
-        self.context = AdminRequestContext.test_context(
-            {}, InMemoryProfile.test_profile()
-        )
+        profile = InMemoryProfile.test_profile()
+        self.context = AdminRequestContext.test_context({}, profile)
         self.request_dict = {
             "context": self.context,
         }
@@ -495,10 +497,19 @@ async def test_upgrade_middleware(self):
 
         await test_module.upgrade_middleware(request, handler)
 
-        with self.assertRaises(test_module.web.HTTPServiceUnavailable):
-            await test_module.upgrade_middleware(request, handler)
+        async with profile.session() as session:
+            storage = session.inject(BaseStorage)
+            upgrading_record = StorageRecord(
+                RECORD_TYPE_ACAPY_UPGRADING,
+                "true",
+            )
+            await storage.add_record(upgrading_record)
 
-        await test_module.upgrade_middleware(request, handler)
+            with self.assertRaises(test_module.web.HTTPServiceUnavailable):
+                await test_module.upgrade_middleware(request, handler)
+
+            await storage.delete_record(upgrading_record)
+            await test_module.upgrade_middleware(request, handler)
 
 
 @pytest.fixture