fix: actually return mediated conn ids

Signed-off-by: Daniel Bluhm <dbluhm@pm.me>

Author: dbluhm

aries_cloudagent/protocols/coordinate_mediation/v1_0/manager.py

@@ -6,12 +6,14 @@
 
 from ....core.error import BaseError
 from ....core.profile import Profile, ProfileSession
+from ....connections.models.conn_record import ConnRecord
 from ....storage.base import BaseStorage
 from ....storage.error import StorageNotFoundError
 from ....storage.record import StorageRecord
 from ....wallet.base import BaseWallet
 from ....wallet.did_info import DIDInfo
 from ....wallet.did_method import DIDMethod
+from ....wallet.error import WalletNotFoundError
 from ....wallet.key_type import KeyType
 from ...routing.v1_0.manager import RoutingManager
 from ...routing.v1_0.models.route_record import RouteRecord
@@ -600,21 +602,34 @@ async def store_update_results(
             for record_for_removal in to_remove:
                 await record_for_removal.delete_record(session)
 
+    async def _conn_id_from_recipient_key(
+        self, session: ProfileSession, wallet: BaseWallet, recipient_key: str
+    ) -> str:
+        try:
+            conn = await ConnRecord.retrieve_by_invitation_key(
+                session, invitation_key=normalize_from_did_key(recipient_key)
+            )
+        except StorageNotFoundError:
+            did_info = await wallet.get_local_did_for_verkey(
+                normalize_from_did_key(recipient_key)
+            )
+            conn = await ConnRecord.retrieve_by_did(session, my_did=did_info.did)
+        return conn.connection_id
+
     async def notify_keylist_updated(
         self, connection_id: str, response: KeylistUpdateResponse
     ):
         """Notify of keylist update response received."""
-        # Retrieve connection IDs associated with recipient keys
-        # recipient key -> connection id
         async with self._profile.session() as session:
+            wallet = session.inject(BaseWallet)
             try:
-                routes = [
-                    await RouteRecord.retrieve_by_recipient_key(
-                        session, updated.recipient_key
+                routes = {
+                    updated.recipient_key: await self._conn_id_from_recipient_key(
+                        session, wallet, updated.recipient_key
                     )
                     for updated in response.updated
-                ]
-            except StorageNotFoundError as err:
+                }
+            except (StorageNotFoundError, WalletNotFoundError) as err:
                 raise MediationManagerError(
                     "Unknown recipient key received in keylist update response"
                 ) from err
@@ -625,9 +640,7 @@ async def notify_keylist_updated(
                 "connection_id": connection_id,
                 "thread_id": response._thread_id,
                 "updated": [update.serialize() for update in response.updated],
-                "mediated_connections": {
-                    route.recipient_key: route.connection_id for route in routes
-                },
+                "mediated_connections": routes,
             },
         )
 

aries_cloudagent/protocols/coordinate_mediation/v1_0/tests/test_mediation_manager.py

@@ -1,7 +1,8 @@
 """Test MediationManager."""
 import logging
-from typing import AsyncIterable, Iterable
+from typing import AsyncGenerator, AsyncIterable, Iterable
 
+from functools import partial
 from asynctest import mock as async_mock
 import pytest
 
@@ -482,37 +483,40 @@ async def test_notify_keylist_updated(
         self,
         manager: MediationManager,
         mock_event_bus: MockEventBus,
-        session: ProfileSession,
     ):
         """test notify_keylist_updated."""
-        await RouteRecord(
-            role=RouteRecord.ROLE_CLIENT,
-            connection_id="conn_id_1",
-            recipient_key=TEST_ROUTE_VERKEY,
-        ).save(session)
-        await RouteRecord(
-            role=RouteRecord.ROLE_CLIENT,
-            connection_id="conn_id_2",
-            recipient_key=TEST_VERKEY,
-        ).save(session)
 
-        response = KeylistUpdateResponse(
-            updated=[
-                KeylistUpdated(
-                    recipient_key=TEST_ROUTE_VERKEY,
-                    action=KeylistUpdateRule.RULE_ADD,
-                    result=KeylistUpdated.RESULT_SUCCESS,
-                ),
-                KeylistUpdated(
-                    recipient_key=TEST_VERKEY,
-                    action=KeylistUpdateRule.RULE_REMOVE,
-                    result=KeylistUpdated.RESULT_SUCCESS,
-                ),
-            ],
-        )
+        async def _result_generator():
+            yield "conn_id_1"
+            yield "conn_id_2"
 
-        response.assign_thread_id(TEST_THREAD_ID)
-        await manager.notify_keylist_updated(TEST_CONN_ID, response)
+        async def _retrieve_by_invitation_key(
+            generator: AsyncGenerator, *args, **kwargs
+        ):
+            return await generator.__anext__()
+
+        with async_mock.patch.object(
+            manager,
+            "_conn_id_from_recipient_key",
+            partial(_retrieve_by_invitation_key, _result_generator()),
+        ):
+            response = KeylistUpdateResponse(
+                updated=[
+                    KeylistUpdated(
+                        recipient_key=TEST_ROUTE_VERKEY,
+                        action=KeylistUpdateRule.RULE_ADD,
+                        result=KeylistUpdated.RESULT_SUCCESS,
+                    ),
+                    KeylistUpdated(
+                        recipient_key=TEST_VERKEY,
+                        action=KeylistUpdateRule.RULE_REMOVE,
+                        result=KeylistUpdated.RESULT_SUCCESS,
+                    ),
+                ],
+            )
+
+            response.assign_thread_id(TEST_THREAD_ID)
+            await manager.notify_keylist_updated(TEST_CONN_ID, response)
         assert mock_event_bus.events
         assert mock_event_bus.events[0][1].topic == manager.KEYLIST_UPDATED_EVENT
         assert mock_event_bus.events[0][1].payload == {