From 311b755b3b57cf4e4d655f7926eb9d580cd65dc9 Mon Sep 17 00:00:00 2001
From: James Campbell <james.campbell@tanti.org.uk>
Date: Wed, 2 Aug 2023 15:38:12 +0200
Subject: [PATCH] Add tests for timeout

---
 src/dkg.ts                          |  12 ++--
 test/integration/dkg-client.test.ts | 100 +++++++++++++++++++++++++++-
 2 files changed, 104 insertions(+), 8 deletions(-)

diff --git a/src/dkg.ts b/src/dkg.ts
index ca57868fd..89eb44f1e 100644
--- a/src/dkg.ts
+++ b/src/dkg.ts
@@ -156,18 +156,16 @@ export class DkgClient {
   private static performRitual = async (
     web3Provider: ethers.providers.Web3Provider,
     ritualId: number
-    ): Promise<void> => {
+  ): Promise<void> => {
     const isSuccessful = await DkgClient.waitUntilRitualEnd(
       web3Provider,
       ritualId
     );
-  
+
     if (!isSuccessful) {
-      throw new Error(
-        `Ritual initialization failed. Ritual id ${ritualId}`
-      );
+      throw new Error(`Ritual initialization failed. Ritual id ${ritualId}`);
     }
-  }
+  };
 
   private static waitForBlockTime = async (
     web3Provider: ethers.providers.Web3Provider,
@@ -181,7 +179,7 @@ export class DkgClient {
         await new Promise((resolve) => setTimeout(resolve, 1000)); // Wait for 1 second before checking again
       }
     } while (currentBlockTime < endTimestamp);
-  }
+  };
 
   private static waitUntilRitualEnd = async (
     web3Provider: ethers.providers.Web3Provider,
diff --git a/test/integration/dkg-client.test.ts b/test/integration/dkg-client.test.ts
index ff2150465..64cc59a75 100644
--- a/test/integration/dkg-client.test.ts
+++ b/test/integration/dkg-client.test.ts
@@ -1,6 +1,9 @@
 import { SecretKey } from '@nucypher/nucypher-core';
 
-import { DkgCoordinatorAgent } from '../../src/agents/coordinator';
+import {
+  DkgCoordinatorAgent,
+  DkgRitualState,
+} from '../../src/agents/coordinator';
 import { DkgClient } from '../../src/dkg';
 import {
   fakeCoordinatorRitual,
@@ -81,4 +84,99 @@ describe('DkgClient', () => {
     );
     expect(getParticipantPublicKeysSpy).toHaveBeenCalled();
   });
+
+  it('waits until the ritual end time during initialization', async () => {
+    jest.useFakeTimers();
+    const fakeProvider = fakeWeb3Provider(SecretKey.random().toBEBytes());
+    const fakeUrsulas = ['ursula1', 'ursula2', 'ursula3'];
+    const fakeRitualId = 123;
+    const initTimestamp = Math.floor(Date.now() / 1000);
+    const timeout = 10;
+
+    jest
+      .spyOn(DkgCoordinatorAgent, 'initializeRitual')
+      .mockResolvedValue(fakeRitualId);
+
+    jest
+      .spyOn(DkgCoordinatorAgent, 'getRitualInitTime')
+      .mockResolvedValue(initTimestamp);
+
+    jest.spyOn(DkgCoordinatorAgent, 'getTimeout').mockResolvedValue(timeout);
+
+    jest.spyOn(DkgClient as any, 'performRitual').mockResolvedValue(undefined);
+
+    const promise = DkgClient.initializeRitual(fakeProvider, fakeUrsulas, true);
+
+    jest.advanceTimersByTime(timeout * 1000);
+
+    await expect(promise).resolves.toBe(fakeRitualId);
+
+    expect(DkgCoordinatorAgent.initializeRitual).toHaveBeenCalledWith(
+      fakeProvider,
+      fakeUrsulas
+    );
+
+    expect(DkgCoordinatorAgent.getRitualInitTime).toHaveBeenCalledWith(
+      fakeProvider,
+      fakeRitualId
+    );
+
+    expect(DkgCoordinatorAgent.getTimeout).toHaveBeenCalledWith(fakeProvider);
+
+    expect((DkgClient as any).performRitual).toHaveBeenCalledWith(
+      fakeProvider,
+      fakeRitualId
+    );
+  });
+
+  it('throws an error when initialization times out', async () => {
+    jest.useFakeTimers();
+    const fakeProvider = fakeWeb3Provider(SecretKey.random().toBEBytes());
+    const fakeUrsulas = ['ursula1', 'ursula2', 'ursula3'];
+    const fakeRitualId = 123;
+    const initTimestamp = Math.floor(Date.now() / 1000);
+    const timeout = 10;
+
+    jest
+      .spyOn(DkgCoordinatorAgent, 'initializeRitual')
+      .mockResolvedValue(fakeRitualId);
+
+    jest
+      .spyOn(DkgCoordinatorAgent, 'getRitualInitTime')
+      .mockResolvedValue(initTimestamp);
+
+    jest.spyOn(DkgCoordinatorAgent, 'getTimeout').mockResolvedValue(timeout);
+
+    const performRitualSpy = jest
+      .spyOn(DkgClient as any, 'performRitual')
+      .mockRejectedValue(
+        new Error(`Ritual initialization failed. Ritual id ${fakeRitualId}`)
+      );
+
+    jest
+      .spyOn(DkgCoordinatorAgent, 'getRitualState')
+      .mockResolvedValue(DkgRitualState.TIMEOUT);
+
+    const promise = DkgClient.initializeRitual(fakeProvider, fakeUrsulas, true);
+
+    jest.advanceTimersByTime(timeout * 1000);
+
+    await expect(promise).rejects.toThrow(
+      `Ritual initialization failed. Ritual id ${fakeRitualId} is in state TIMEOUT`
+    );
+
+    expect(DkgCoordinatorAgent.initializeRitual).toHaveBeenCalledWith(
+      fakeProvider,
+      fakeUrsulas
+    );
+
+    expect(DkgCoordinatorAgent.getRitualInitTime).toHaveBeenCalledWith(
+      fakeProvider,
+      fakeRitualId
+    );
+
+    expect(DkgCoordinatorAgent.getTimeout).toHaveBeenCalledWith(fakeProvider);
+
+    expect(performRitualSpy).toHaveBeenCalledWith(fakeProvider, fakeRitualId);
+  });
 });