feat(lit-client, networks): handle edge cases

Author: Ansonhkg

packages/e2e/src/helper/shiva-client.ts renamed to packages/e2e/src/helper/ShivaClient/createShivaClient.ts

@@ -1,3 +1,9 @@
+import { createLitClient } from '@lit-protocol/lit-client';
+import {
+  createEpochSnapshot,
+  EpochSnapshot,
+} from './helpers/createEpochSnapshot';
+
 /**
  * Options used when Shiva spins up a brand-new testnet instance.
  * Values mirror the Rust manager contract; all fields are optional for our wrapper.
@@ -40,23 +46,17 @@ type FetchOptions = {
   body?: unknown;
 };
 
-/**
- * Snapshot returned from {@link ShivaClient.inspectEpoch} and {@link ShivaClient.waitForEpochChange}.
- */
-// type EpochSnapshot = {
-//   epoch: number | undefined;
-//   nodeEpochs: Array<{ url: string; epoch: number | undefined }>;
-//   threshold: number | undefined;
-//   connectedCount: number | undefined;
-//   latestBlockhash: string | undefined;
-//   rawContext: any;
-// };
-
 /**
  * Options for {@link ShivaClient.waitForEpochChange}.
  */
 type WaitForEpochOptions = {
-  baselineEpoch: number | undefined;
+  expectedEpoch: number | undefined;
+  timeoutMs?: number;
+  intervalMs?: number;
+};
+
+type PollTestnetStateOptions = {
+  waitFor?: TestNetState | TestNetState[];
   timeoutMs?: number;
   intervalMs?: number;
 };
@@ -68,26 +68,42 @@ export type ShivaClient = {
   baseUrl: string;
   testnetId: string;
   /** Fetch a one-off snapshot of the Lit context and per-node epochs. */
-  // inspectEpoch: () => Promise<EpochSnapshot>;
+  inspectEpoch: () => Promise<EpochSnapshot>;
   /**
    * Poll the Lit client until it reports an epoch different from {@link WaitForEpochOptions.baselineEpoch}.
    * Useful immediately after triggering an epoch change via Shiva.
    */
-  // waitForEpochChange: (options: WaitForEpochOptions) => Promise<EpochSnapshot>;
+  waitForEpochChange: (options: WaitForEpochOptions) => Promise<EpochSnapshot>;
   /** Invoke Shiva's `/test/action/transition/epoch/wait/<id>` and wait for completion. */
   transitionEpochAndWait: () => Promise<boolean>;
   /** Stop a random node and wait for the subsequent epoch change. */
   stopRandomNodeAndWait: () => Promise<boolean>;
   /** Query the current state of the managed testnet (Busy, Active, etc.). */
-  pollTestnetState: () => Promise<TestNetState>;
+  /**
+   * @example
+   * ```ts
+   * // Wait up to two minutes for the testnet to become active.
+   * await client.pollTestnetState({ waitFor: 'Active', timeoutMs: 120_000 });
+   * ```
+   */
+  pollTestnetState: (
+    options?: PollTestnetStateOptions
+  ) => Promise<TestNetState>;
   /** Retrieve the full testnet configuration (contract ABIs, RPC URL, etc.). */
   getTestnetInfo: () => Promise<unknown>;
   /** Shut down the underlying testnet through the Shiva manager. */
   deleteTestnet: () => Promise<boolean>;
+
+  // Setters
+  setLitClient: (
+    litClient: Awaited<ReturnType<typeof createLitClient>>
+  ) => void;
 };
 
-// const DEFAULT_POLL_INTERVAL = 2000;
-// const DEFAULT_TIMEOUT = 60_000;
+const DEFAULT_POLL_INTERVAL = 2000;
+const DEFAULT_TIMEOUT = 60_000;
+const DEFAULT_STATE_POLL_INTERVAL = 2000;
+const DEFAULT_STATE_POLL_TIMEOUT = 60_000;
 
 const normaliseBaseUrl = (baseUrl: string) => {
   return baseUrl.endsWith('/') ? baseUrl.slice(0, -1) : baseUrl;
@@ -186,33 +202,6 @@ const getOrCreateTestnetId = async (
   return response.testnetId;
 };
 
-// const buildEpochSnapshot = (ctx: any): EpochSnapshot => {
-//   const nodeEpochEntries = Object.entries(
-//     ctx?.handshakeResult?.serverKeys ?? {}
-//   );
-//   const nodeEpochs = nodeEpochEntries.map(([url, data]: [string, any]) => ({
-//     url,
-//     epoch: data?.epoch,
-//   }));
-
-//   const connected = ctx?.handshakeResult?.connectedNodes;
-//   const connectedCount =
-//     typeof connected?.size === 'number'
-//       ? connected.size
-//       : Array.isArray(connected)
-//       ? connected.length
-//       : undefined;
-
-//   return {
-//     epoch: ctx?.latestConnectionInfo?.epochInfo?.number,
-//     nodeEpochs,
-//     threshold: ctx?.handshakeResult?.threshold,
-//     connectedCount,
-//     latestBlockhash: ctx?.latestBlockhash,
-//     rawContext: ctx,
-//   };
-// };
-
 /**
  * Creates a Shiva client wrapper for the provided Lit client instance.
  * The wrapper talks to the Shiva manager REST endpoints, auto-discovers (or optionally creates) a testnet,
@@ -228,30 +217,47 @@ export const createShivaClient = async (
     options.createRequest
   );
 
-  // const inspectEpoch = async () => {
-  //   const ctx = await litClient.getContext();
-  //   return buildEpochSnapshot(ctx);
-  // };
-
-  // const waitForEpochChange = async ({
-  //   baselineEpoch,
-  //   timeoutMs = DEFAULT_TIMEOUT,
-  //   intervalMs = DEFAULT_POLL_INTERVAL,
-  // }: WaitForEpochOptions) => {
-  //   const deadline = Date.now() + timeoutMs;
-
-  //   while (Date.now() < deadline) {
-  //     await new Promise((resolve) => setTimeout(resolve, intervalMs));
-  //     const snapshot = await inspectEpoch();
-  //     if (snapshot.epoch !== baselineEpoch) {
-  //       return snapshot;
-  //     }
-  //   }
-
-  //   throw new Error(
-  //     `Epoch did not change from ${baselineEpoch} within ${timeoutMs}ms`
-  //   );
-  // };
+  let litClientInstance:
+    | Awaited<ReturnType<typeof createLitClient>>
+    | undefined;
+
+  const setLitClient = (
+    client: Awaited<ReturnType<typeof createLitClient>>
+  ) => {
+    litClientInstance = client;
+  };
+
+  const inspectEpoch = async () => {
+    if (!litClientInstance) {
+      throw new Error(
+        `Lit client not set. Please call setLitClient() before using inspectEpoch().`
+      );
+    }
+
+    return createEpochSnapshot(litClientInstance);
+  };
+
+  const waitForEpochChange = async ({
+    expectedEpoch,
+    timeoutMs = DEFAULT_TIMEOUT,
+    intervalMs = DEFAULT_POLL_INTERVAL,
+  }: WaitForEpochOptions) => {
+    const deadline = Date.now() + timeoutMs;
+
+    while (Date.now() < deadline) {
+      await new Promise((resolve) => setTimeout(resolve, intervalMs));
+      const snapshot = await inspectEpoch();
+      if (
+        snapshot.latestConnectionInfo.epochState.currentNumber !== expectedEpoch
+      ) {
+        return snapshot;
+      }
+    }
+
+    throw new Error(
+      `Epoch did not change from ${expectedEpoch} within ${timeoutMs}ms`
+    );
+  };
 
   const transitionEpochAndWait = async () => {
     const response = await fetchShiva<boolean>(
@@ -266,15 +272,52 @@ export const createShivaClient = async (
       baseUrl,
       `/test/action/stop/random/wait/${testnetId}`
     );
+
+    // wait briefly to allow the node to drop from the network
+    await new Promise((resolve) => setTimeout(resolve, 5000));
+
     return Boolean(response.body);
   };
 
-  const pollTestnetState = async () => {
-    const response = await fetchShiva<string>(
-      baseUrl,
-      `/test/poll/testnet/${testnetId}`
-    );
-    return (response.body ?? 'UNKNOWN') as TestNetState;
+  const pollTestnetState = async (
+    options: PollTestnetStateOptions = {}
+  ): Promise<TestNetState> => {
+    const {
+      waitFor,
+      timeoutMs = DEFAULT_STATE_POLL_TIMEOUT,
+      intervalMs = DEFAULT_STATE_POLL_INTERVAL,
+    } = options;
+
+    const desiredStates = Array.isArray(waitFor)
+      ? waitFor
+      : waitFor
+      ? [waitFor]
+      : undefined;
+    const deadline = Date.now() + timeoutMs;
+
+    // Continue polling until we hit a desired state or timeout.
+    // If no desired state is provided, return the first observation .
+    for (;;) {
+      const response = await fetchShiva<string>(
+        baseUrl,
+        `/test/poll/testnet/${testnetId}`
+      );
+      const state = (response.body ?? 'UNKNOWN') as TestNetState;
+
+      if (!desiredStates || desiredStates.includes(state)) {
+        return state;
+      }
+
+      if (Date.now() >= deadline) {
+        throw new Error(
+          `Timed out after ${timeoutMs}ms waiting for testnet ${testnetId} to reach state ${desiredStates.join(
+            ', '
+          )}. Last observed state: ${state}.`
+        );
+      }
+
+      await new Promise((resolve) => setTimeout(resolve, intervalMs));
+    }
   };
 
   const getTestnetInfo = async () => {
@@ -296,12 +339,15 @@ export const createShivaClient = async (
   return {
     baseUrl,
     testnetId,
-    // inspectEpoch,
-    // waitForEpochChange,
+    setLitClient,
     transitionEpochAndWait,
     stopRandomNodeAndWait,
     pollTestnetState,
     getTestnetInfo,
     deleteTestnet,
+
+    // utils
+    inspectEpoch,
+    waitForEpochChange,
   };
 };

packages/e2e/src/helper/ShivaClient/helpers/createEpochSnapshot.ts

@@ -0,0 +1,82 @@
+type EpochInfo = {
+  epochLength: number;
+  number: number;
+  endTime: number;
+  retries: number;
+  timeout: number;
+};
+
+type EpochState = {
+  currentNumber: number;
+  startTime: number;
+};
+
+type NetworkPrice = {
+  url: string;
+  prices: Array<number | bigint>;
+};
+
+type PriceFeedInfo = {
+  epochId: number;
+  minNodeCount: number;
+  networkPrices: NetworkPrice[];
+};
+
+type LatestConnectionInfo = {
+  epochInfo: EpochInfo;
+  epochState: EpochState;
+  minNodeCount: number;
+  bootstrapUrls: string[];
+  priceFeedInfo: PriceFeedInfo;
+};
+
+type ServerKeyDetails = {
+  serverPublicKey: string;
+  subnetPublicKey: string;
+  networkPublicKey: string;
+  networkPublicKeySet: string;
+  clientSdkVersion: string;
+  hdRootPubkeys: string[];
+  attestation?: string | null;
+  latestBlockhash: string;
+  nodeIdentityKey: string;
+  nodeVersion: string;
+  epoch: number;
+};
+
+type CoreNodeConfig = {
+  subnetPubKey: string;
+  networkPubKey: string;
+  networkPubKeySet: string;
+  hdRootPubkeys: string[];
+  latestBlockhash: string;
+};
+
+type HandshakeResult = {
+  serverKeys: Record<string, ServerKeyDetails>;
+  connectedNodes: Record<string, unknown> | Set<string>;
+  coreNodeConfig: CoreNodeConfig | null;
+  threshold: number;
+};
+
+type EpochSnapshotSource = {
+  latestConnectionInfo?: LatestConnectionInfo | null;
+  handshakeResult?: HandshakeResult | null;
+};
+
+export type EpochSnapshot = EpochSnapshotSource;
+
+export const createEpochSnapshot = async (
+  litClient: Awaited<
+    ReturnType<typeof import('@lit-protocol/lit-client').createLitClient>
+  >
+): Promise<EpochSnapshot> => {
+  const ctx = await litClient.getContext();
+
+  const snapshot = {
+    latestConnectionInfo: ctx?.latestConnectionInfo,
+    handshakeResult: ctx?.handshakeResult,
+  };
+
+  return snapshot;
+};

packages/e2e/src/index.ts

@@ -5,7 +5,6 @@ export * from './helper/tests';
 export { init } from './init';
 
 export { getOrCreatePkp } from './helper/pkp-utils';
-export { createShivaClient } from './helper/shiva-client';
 export { printAligned } from './helper/utils';
 export type { AuthContext } from './types';
 
@@ -16,3 +15,6 @@ export { createTestAccount } from './helper/createTestAccount';
 export { createTestEnv } from './helper/createTestEnv';
 export type { CreateTestAccountResult } from './helper/createTestAccount';
 export { registerPaymentDelegationTicketSuite } from './tickets/delegation.suite';
+
+// -- Shiva
+export { createShivaClient } from './helper/ShivaClient/createShivaClient';