From e578d84464403d4a15ee8a7cf3ac643f4fb86d69 Mon Sep 17 00:00:00 2001
From: Kerry Archibald <kerrya@element.io>
Date: Wed, 4 Jan 2023 17:48:55 +1300
Subject: [PATCH] allow alt event types in relations model

---
 spec/unit/relations.spec.ts | 94 +++++++++++++++++++++++++++++++++++++
 src/models/relations.ts     |  8 +++-
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/spec/unit/relations.spec.ts b/spec/unit/relations.spec.ts
index 91b77dd1212..9da81f0a146 100644
--- a/spec/unit/relations.spec.ts
+++ b/spec/unit/relations.spec.ts
@@ -14,13 +14,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+import { M_POLL_START } from "matrix-events-sdk";
+
 import { EventTimelineSet } from "../../src/models/event-timeline-set";
 import { MatrixEvent, MatrixEventEvent } from "../../src/models/event";
 import { Room } from "../../src/models/room";
 import { Relations } from "../../src/models/relations";
 import { TestClient } from "../TestClient";
+import { RelationType } from "../../src";
+import { logger } from "../../src/logger";
 
 describe("Relations", function () {
+    afterEach(() => {
+        jest.spyOn(logger, "error").mockRestore();
+    });
+
     it("should deduplicate annotations", function () {
         const room = new Room("room123", null!, null!);
         const relations = new Relations("m.annotation", "m.reaction", room);
@@ -75,6 +83,92 @@ describe("Relations", function () {
         }
     });
 
+    describe("addEvent()", () => {
+        const relationType = RelationType.Reference;
+        const eventType = M_POLL_START.stable!;
+        const altEventTypes = [M_POLL_START.unstable!];
+        const room = new Room("room123", null!, null!);
+
+        it("should not add events without a relation", async () => {
+            // dont pollute console
+            const logSpy = jest.spyOn(logger, "error").mockImplementation(() => {});
+            const relations = new Relations(relationType, eventType, room);
+            const emitSpy = jest.spyOn(relations, "emit");
+            const event = new MatrixEvent({ type: eventType });
+
+            await relations.addEvent(event);
+            expect(logSpy).toHaveBeenCalledWith("Event must have relation info");
+            // event not added
+            expect(relations.getRelations().length).toBe(0);
+            expect(emitSpy).not.toHaveBeenCalled();
+        });
+
+        it("should not add events of incorrect event type", async () => {
+            // dont pollute console
+            const logSpy = jest.spyOn(logger, "error").mockImplementation(() => {});
+            const relations = new Relations(relationType, eventType, room);
+            const emitSpy = jest.spyOn(relations, "emit");
+            const event = new MatrixEvent({
+                type: "different-event-type",
+                content: {
+                    "m.relates_to": {
+                        event_id: "$2s4yYpEkVQrPglSCSqB_m6E8vDhWsg0yFNyOJdVIb_o",
+                        rel_type: relationType,
+                    },
+                },
+            });
+
+            await relations.addEvent(event);
+
+            expect(logSpy).toHaveBeenCalledWith(`Event relation info doesn't match this container`);
+            // event not added
+            expect(relations.getRelations().length).toBe(0);
+            expect(emitSpy).not.toHaveBeenCalled();
+        });
+
+        it("adds events that match alt event types", async () => {
+            const relations = new Relations(relationType, eventType, room, altEventTypes);
+            const emitSpy = jest.spyOn(relations, "emit");
+            const event = new MatrixEvent({
+                type: M_POLL_START.unstable!,
+                content: {
+                    "m.relates_to": {
+                        event_id: "$2s4yYpEkVQrPglSCSqB_m6E8vDhWsg0yFNyOJdVIb_o",
+                        rel_type: relationType,
+                    },
+                },
+            });
+
+            await relations.addEvent(event);
+
+            // event added
+            expect(relations.getRelations()).toEqual([event]);
+            expect(emitSpy).toHaveBeenCalled();
+        });
+
+        it("should not add events of incorrect relation type", async () => {
+            const logSpy = jest.spyOn(logger, "error").mockImplementation(() => {});
+            const relations = new Relations(relationType, eventType, room);
+            const event = new MatrixEvent({
+                type: eventType,
+                content: {
+                    "m.relates_to": {
+                        event_id: "$2s4yYpEkVQrPglSCSqB_m6E8vDhWsg0yFNyOJdVIb_o",
+                        rel_type: "m.annotation",
+                    },
+                },
+            });
+
+            await relations.addEvent(event);
+            const emitSpy = jest.spyOn(relations, "emit");
+
+            expect(logSpy).toHaveBeenCalledWith(`Event relation info doesn't match this container`);
+            // event not added
+            expect(relations.getRelations().length).toBe(0);
+            expect(emitSpy).not.toHaveBeenCalled();
+        });
+    });
+
     it("should emit created regardless of ordering", async function () {
         const targetEvent = new MatrixEvent({
             sender: "@bob:example.com",
diff --git a/src/models/relations.ts b/src/models/relations.ts
index bd6f0093814..714d5577a85 100644
--- a/src/models/relations.ts
+++ b/src/models/relations.ts
@@ -33,6 +33,9 @@ export type EventHandlerMap = {
     [RelationsEvent.Redaction]: (event: MatrixEvent) => void;
 };
 
+const matchesEventType = (eventType: string, targetEventType: string, altTargetEventTypes: string[] = []): boolean =>
+    [targetEventType, ...altTargetEventTypes].includes(eventType);
+
 /**
  * A container for relation events that supports easy access to common ways of
  * aggregating such events. Each instance holds events that of a single relation
@@ -60,6 +63,7 @@ export class Relations extends TypedEventEmitter<RelationsEvent, EventHandlerMap
         public readonly relationType: RelationType | string,
         public readonly eventType: string,
         client: MatrixClient | Room,
+        public readonly altEventTypes?: string[],
     ) {
         super();
         this.client = client instanceof Room ? client.client : client;
@@ -84,7 +88,7 @@ export class Relations extends TypedEventEmitter<RelationsEvent, EventHandlerMap
         const relationType = relation.rel_type;
         const eventType = event.getType();
 
-        if (this.relationType !== relationType || this.eventType !== eventType) {
+        if (this.relationType !== relationType || !matchesEventType(eventType, this.eventType, this.altEventTypes)) {
             logger.error("Event relation info doesn't match this container");
             return;
         }
@@ -131,7 +135,7 @@ export class Relations extends TypedEventEmitter<RelationsEvent, EventHandlerMap
         const relationType = relation.rel_type;
         const eventType = event.getType();
 
-        if (this.relationType !== relationType || this.eventType !== eventType) {
+        if (this.relationType !== relationType || !matchesEventType(eventType, this.eventType, this.altEventTypes)) {
             logger.error("Event relation info doesn't match this container");
             return;
         }