From 83df9deb6a71eb4708f3bc9e7c9de4986c55469c Mon Sep 17 00:00:00 2001
From: Half-Shot <will@half-shot.uk>
Date: Tue, 18 Jan 2022 15:43:49 +0000
Subject: [PATCH] Ensure tests support E2EE

---
 test/integration/clientHelper.ts      | 11 ++++++++---
 test/integration/mjolnirSetupUtils.ts | 13 ++++++++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/test/integration/clientHelper.ts b/test/integration/clientHelper.ts
index 2afb902..c005914 100644
--- a/test/integration/clientHelper.ts
+++ b/test/integration/clientHelper.ts
@@ -1,7 +1,7 @@
 import axios from "axios";
 import { HmacSHA1 } from "crypto-js";
 import e from "express";
-import { LogService, MatrixClient, MemoryStorageProvider, PantalaimonClient } from "matrix-bot-sdk";
+import { LogService, MatrixClient, MemoryStorageProvider, PantalaimonClient, RustSdkCryptoStorageProvider } from "matrix-bot-sdk";
 import config from "../../src/config";
 
 const REGISTRATION_ATTEMPTS = 10;
@@ -82,8 +82,13 @@ export async function registerNewTestUser(isAdmin: boolean, label: string = "")
  */
 export async function newTestUser(isAdmin: boolean = false, label: string = ""): Promise<MatrixClient> {
     const username = await registerNewTestUser(isAdmin, label);
-    const pantalaimon = new PantalaimonClient(config.homeserverUrl, new MemoryStorageProvider());
-    return await pantalaimon.createClientWithCredentials(username, username);
+    if (config.pantalaimon) {
+        const pantalaimon = new PantalaimonClient(config.homeserverUrl, new MemoryStorageProvider());
+        return await pantalaimon.createClientWithCredentials(username, username);
+    } else {
+        const client = new MatrixClient(config.homeserverUrl, config.accessToken, new MemoryStorageProvider(), new RustSdkCryptoStorageProvider(config.dataPath));
+        client.crypto.prepare(await client.getJoinedRooms());
+    }
 }
 
 /**
diff --git a/test/integration/mjolnirSetupUtils.ts b/test/integration/mjolnirSetupUtils.ts
index b30b3ae..de47315 100644
--- a/test/integration/mjolnirSetupUtils.ts
+++ b/test/integration/mjolnirSetupUtils.ts
@@ -19,7 +19,8 @@ import {
     MemoryStorageProvider,
     LogService,
     LogLevel,
-    RichConsoleLogger
+    RichConsoleLogger,
+    RustSdkCryptoStorageProvider
 } from "matrix-bot-sdk";
 import { Mjolnir}  from '../../src/Mjolnir';
 import config from "../../src/config";
@@ -80,8 +81,14 @@ export async function makeMjolnir(): Promise<Mjolnir> {
     LogService.setLogger(new RichConsoleLogger());
     LogService.setLevel(LogLevel.fromString(config.logLevel, LogLevel.DEBUG));
     LogService.info("test/mjolnirSetupUtils", "Starting bot...");
-    const pantalaimon = new PantalaimonClient(config.homeserverUrl, new MemoryStorageProvider());
-    const client = await pantalaimon.createClientWithCredentials(config.pantalaimon.username, config.pantalaimon.password);
+    let client: MatrixClient;
+    if (config.pantalaimon) {
+        const pantalaimon = new PantalaimonClient(config.homeserverUrl, new MemoryStorageProvider());
+        client = await pantalaimon.createClientWithCredentials(config.pantalaimon.username, config.pantalaimon.password);
+    } else {
+        client = new MatrixClient(config.homeserverUrl, config.accessToken, new MemoryStorageProvider(), new RustSdkCryptoStorageProvider(config.dataPath));
+        client.crypto.prepare(await client.getJoinedRooms());
+    }
     patchMatrixClient();
     await ensureAliasedRoomExists(client, config.managementRoom);
     let mjolnir = await Mjolnir.setupMjolnirFromConfig(client);