Examine all m.direct rooms to find a DM

Author: weeman1337

src/utils/DMRoomMap.ts

@@ -192,6 +192,16 @@ export default class DMRoomMap {
             .reduce((obj, r) => (obj[r.userId] = r.room) && obj, {});
     }
 
+    /**
+     * @returns all room Ids from m.direct
+     */
+    public getRoomIds(): Set<string> {
+        return Object.values(this.mDirectEvent).reduce((prevRoomIds: Set<string>, roomIds: string[]): Set<string> => {
+            roomIds.forEach((roomId) => prevRoomIds.add(roomId));
+            return prevRoomIds;
+        }, new Set<string>());
+    }
+
     private getUserToRooms(): { [key: string]: string[] } {
         if (!this.userToRooms) {
             const userToRooms = this.mDirectEvent;

src/utils/dm/findDMForUser.ts

@@ -29,8 +29,8 @@ import { getFunctionalMembers } from "../room/getFunctionalMembers";
  * @returns {Room} Room if found
  */
 export function findDMForUser(client: MatrixClient, userId: string): Room {
-    const roomIds = DMRoomMap.shared().getDMRoomsForUserId(userId);
-    const rooms = roomIds.map((id) => client.getRoom(id));
+    const roomIds = DMRoomMap.shared().getRoomIds();
+    const rooms = Array.from(roomIds).map((id) => client.getRoom(id));
     const suitableDMRooms = rooms
         .filter((r) => {
             // Validate that we are joined and the other person is also joined. We'll also make sure

test/LegacyCallHandler-test.ts

@@ -225,20 +225,10 @@ describe("LegacyCallHandler", () => {
                     return null;
                 }
             },
-            getDMRoomsForUserId: (userId: string) => {
-                if (userId === NATIVE_ALICE) {
-                    return [NATIVE_ROOM_ALICE];
-                } else if (userId === NATIVE_BOB) {
-                    return [NATIVE_ROOM_BOB];
-                } else if (userId === NATIVE_CHARLIE) {
-                    return [NATIVE_ROOM_CHARLIE];
-                } else if (userId === VIRTUAL_BOB) {
-                    return [VIRTUAL_ROOM_BOB];
-                } else {
-                    return [];
-                }
+            getRoomIds: () => {
+                return [NATIVE_ROOM_ALICE, NATIVE_ROOM_BOB, NATIVE_ROOM_CHARLIE, VIRTUAL_ROOM_BOB];
             },
-        } as DMRoomMap;
+        } as unknown as DMRoomMap;
         DMRoomMap.setShared(dmRoomMap);
 
         pstnLookup = null;

test/utils/DMRoomMap-test.ts

@@ -0,0 +1,54 @@
+/*
+Copyright 2023 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { mocked, Mocked } from "jest-mock";
+import { EventType, IContent, MatrixClient } from "matrix-js-sdk/src/matrix";
+
+import DMRoomMap from "../../src/utils/DMRoomMap";
+import { mkEvent, stubClient } from "../test-utils";
+
+describe("DMRoomMap", () => {
+    const roomId1 = "!room1:example.com";
+    const roomId2 = "!room2:example.com";
+    const roomId3 = "!room3:example.com";
+    const roomId4 = "!room4:example.com";
+
+    const mDirectContent = {
+        "user@example.com": [roomId1, roomId2],
+        "@user:example.com": [roomId1, roomId3, roomId4],
+        "@user2:example.com": [] as string[],
+    } satisfies IContent;
+
+    let client: Mocked<MatrixClient>;
+    let dmRoomMap: DMRoomMap;
+
+    beforeEach(() => {
+        client = mocked(stubClient());
+
+        const mDirectEvent = mkEvent({
+            event: true,
+            type: EventType.Direct,
+            user: client.getSafeUserId(),
+            content: mDirectContent,
+        });
+        client.getAccountData.mockReturnValue(mDirectEvent);
+        dmRoomMap = new DMRoomMap(client);
+    });
+
+    it("getRoomIds should return the room Ids", () => {
+        expect(dmRoomMap.getRoomIds()).toEqual(new Set([roomId1, roomId2, roomId3, roomId4]));
+    });
+});

test/utils/dm/findDMForUser-test.ts

@@ -15,10 +15,10 @@ limitations under the License.
 */
 
 import { mocked } from "jest-mock";
-import { MatrixClient, Room } from "matrix-js-sdk/src/matrix";
+import { EventType, MatrixClient, Room } from "matrix-js-sdk/src/matrix";
 
 import DMRoomMap from "../../../src/utils/DMRoomMap";
-import { createTestClient, makeMembershipEvent } from "../../test-utils";
+import { createTestClient, makeMembershipEvent, mkEvent } from "../../test-utils";
 import { LocalRoom } from "../../../src/models/LocalRoom";
 import { findDMForUser } from "../../../src/utils/dm/findDMForUser";
 import { getFunctionalMembers } from "../../../src/utils/room/getFunctionalMembers";
@@ -39,6 +39,19 @@ describe("findDMForUser", () => {
     let dmRoomMap: DMRoomMap;
     let mockClient: MatrixClient;
 
+    const setUpMDirect = (mDirect: { [key: string]: string[] }) => {
+        const mDirectEvent = mkEvent({
+            event: true,
+            type: EventType.Direct,
+            user: mockClient.getSafeUserId(),
+            content: mDirect,
+        });
+        mocked(mockClient).getAccountData.mockReturnValue(mDirectEvent);
+
+        dmRoomMap = new DMRoomMap(mockClient);
+        jest.spyOn(DMRoomMap, "shared").mockReturnValue(dmRoomMap);
+    };
+
     beforeEach(() => {
         mockClient = createTestClient();
 
@@ -87,42 +100,48 @@ describe("findDMForUser", () => {
                 [room5.roomId]: room5,
             }[roomId];
         });
+    });
 
-        dmRoomMap = {
-            getDMRoomForIdentifiers: jest.fn(),
-            getDMRoomsForUserId: jest.fn(),
-        } as unknown as DMRoomMap;
-        jest.spyOn(DMRoomMap, "shared").mockReturnValue(dmRoomMap);
-        mocked(dmRoomMap.getDMRoomsForUserId).mockReturnValue([
-            room1.roomId,
-            room2.roomId,
-            room3.roomId,
-            room4.roomId,
-            room5.roomId,
-        ]);
+    afterAll(() => {
+        jest.restoreAllMocks();
     });
 
     describe("for an empty DM room list", () => {
         beforeEach(() => {
-            mocked(dmRoomMap.getDMRoomsForUserId).mockReturnValue([]);
+            setUpMDirect({});
         });
 
         it("should return undefined", () => {
             expect(findDMForUser(mockClient, userId1)).toBeUndefined();
         });
     });
 
-    it("should find a room ordered by last activity 1", () => {
-        room1.getLastActiveTimestamp = () => 2;
-        room3.getLastActiveTimestamp = () => 1;
+    describe("when there are soom rooms", () => {
+        beforeEach(() => {
+            setUpMDirect({
+                [userId1]: [room1.roomId, room2.roomId, room3.roomId, room4.roomId, room5.roomId],
+            });
+        });
+
+        it("should find a room ordered by last activity 1", () => {
+            room1.getLastActiveTimestamp = () => 2;
+            room3.getLastActiveTimestamp = () => 1;
 
-        expect(findDMForUser(mockClient, userId1)).toBe(room1);
-    });
+            expect(findDMForUser(mockClient, userId1)).toBe(room1);
+        });
 
-    it("should find a room ordered by last activity 2", () => {
-        room1.getLastActiveTimestamp = () => 1;
-        room3.getLastActiveTimestamp = () => 2;
+        it("should find a room ordered by last activity 2", () => {
+            room1.getLastActiveTimestamp = () => 1;
+            room3.getLastActiveTimestamp = () => 2;
 
-        expect(findDMForUser(mockClient, userId1)).toBe(room3);
+            expect(findDMForUser(mockClient, userId1)).toBe(room3);
+        });
+
+        it("should find a room for a user without an m.direct entry but a DM-like room exists", () => {
+            room1.getLastActiveTimestamp = () => 1;
+            room3.getLastActiveTimestamp = () => 2;
+
+            expect(findDMForUser(mockClient, userId2)).toBe(room3);
+        });
     });
 });