Feat/naga shiva

packages/e2e/src/e2e.spec.ts

@@ -14,9 +14,9 @@ import {
   createViewPKPsByAddressTest,
   createViewPKPsByAuthDataTest,
   init,
+  registerPaymentDelegationTicketSuite,
 } from '@lit-protocol/e2e';
 import type { AuthContext } from '@lit-protocol/e2e';
-import { registerPaymentDelegationTicketSuite } from './tickets/delegation.suite';
 
 const RPC_OVERRIDE = process.env['LIT_YELLOWSTONE_PRIVATE_RPC_URL'];
 if (RPC_OVERRIDE) {

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

@@ -1,4 +1,8 @@
-import type { LitClientInstance } from '../types';
+import { createLitClient } from '@lit-protocol/lit-client';
+import {
+  createEpochSnapshot,
+  EpochSnapshot,
+} from './helpers/createEpochSnapshot';
 
 /**
  * Options used when Shiva spins up a brand-new testnet instance.
@@ -17,11 +21,11 @@ type TestNetCreateRequest = {
 };
 
 type TestNetResponse<T> = {
-  testnet_id: string;
+  testnetId: string;
   command: string;
-  was_canceled: boolean;
+  wasCanceled: boolean;
   body: T | null;
-  last_state_observed: string | null;
+  lastStateObserved: string | null;
   messages: string[] | null;
   errors: string[] | null;
 };
@@ -42,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;
 };
@@ -81,15 +79,31 @@ export type ShivaClient = {
   /** 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_STATE_POLL_INTERVAL = 2000;
+const DEFAULT_STATE_POLL_TIMEOUT = 60_000;
 
 const normaliseBaseUrl = (baseUrl: string) => {
   return baseUrl.endsWith('/') ? baseUrl.slice(0, -1) : baseUrl;
@@ -153,7 +167,7 @@ const getTestnetIds = async (baseUrl: string): Promise<string[]> => {
   return (await response.json()) as string[];
 };
 
-const ensureTestnetId = async (
+const getOrCreateTestnetId = async (
   baseUrl: string,
   providedId?: string,
   createRequest?: TestNetCreateRequest
@@ -178,38 +192,14 @@ const ensureTestnetId = async (
     body: createRequest,
   });
 
-  if (!response.testnet_id) {
-    throw new Error('Shiva create testnet response did not include testnet_id');
+  if (!response.testnetId) {
+    throw new Error(
+      'Shiva create testnet response did not include testnetId. Received: ' +
+        JSON.stringify(response)
+    );
   }
 
-  return response.testnet_id;
-};
-
-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,
-  };
+  return response.testnetId;
 };
 
 /**
@@ -218,23 +208,37 @@ const buildEpochSnapshot = (ctx: any): EpochSnapshot => {
  * and exposes helpers for triggering and validating epoch transitions.
  */
 export const createShivaClient = async (
-  litClient: LitClientInstance,
   options: CreateShivaClientOptions
 ): Promise<ShivaClient> => {
   const baseUrl = normaliseBaseUrl(options.baseUrl);
-  const testnetId = await ensureTestnetId(
+  const testnetId = await getOrCreateTestnetId(
     baseUrl,
     options.testnetId,
     options.createRequest
   );
 
+  let litClientInstance:
+    | Awaited<ReturnType<typeof createLitClient>>
+    | undefined;
+
+  const setLitClient = (
+    client: Awaited<ReturnType<typeof createLitClient>>
+  ) => {
+    litClientInstance = client;
+  };
+
   const inspectEpoch = async () => {
-    const ctx = await litClient.getContext();
-    return buildEpochSnapshot(ctx);
+    if (!litClientInstance) {
+      throw new Error(
+        `Lit client not set. Please call setLitClient() before using inspectEpoch().`
+      );
+    }
+
+    return createEpochSnapshot(litClientInstance);
   };
 
   const waitForEpochChange = async ({
-    baselineEpoch,
+    expectedEpoch,
     timeoutMs = DEFAULT_TIMEOUT,
     intervalMs = DEFAULT_POLL_INTERVAL,
   }: WaitForEpochOptions) => {
@@ -243,13 +247,15 @@ export const createShivaClient = async (
     while (Date.now() < deadline) {
       await new Promise((resolve) => setTimeout(resolve, intervalMs));
       const snapshot = await inspectEpoch();
-      if (snapshot.epoch !== baselineEpoch) {
+      if (
+        snapshot.latestConnectionInfo.epochState.currentNumber !== expectedEpoch
+      ) {
         return snapshot;
       }
     }
 
     throw new Error(
-      `Epoch did not change from ${baselineEpoch} within ${timeoutMs}ms`
+      `Epoch did not change from ${expectedEpoch} within ${timeoutMs}ms`
     );
   };
 
@@ -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

@@ -1,10 +1,20 @@
 // re-export
-export { init } from './init';
 export * from './helper/auth-contexts';
-export * from './helper/tests';
 export * from './helper/NetworkManager';
+export * from './helper/tests';
+export { init } from './init';
 
-export { printAligned } from './helper/utils';
 export { getOrCreatePkp } from './helper/pkp-utils';
-export { createShivaClient } from './helper/shiva-client';
+export { printAligned } from './helper/utils';
 export type { AuthContext } from './types';
+
+// re-export new helpers that should be used to refactor the `init.ts` proces
+// see packages/e2e/src/tickets/delegation.suite.ts for usage examples
+export { createEnvVars } from './helper/createEnvVars';
+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';