Split out fetchVariableDefinitionByPredicate function (#DH-21947)

Author: bmingles

package-lock.json

@@ -1,6 +1,9 @@
 import dh from '@deephaven/jsapi-shim';
 import { TimeoutError } from '@deephaven/utils';
-import { fetchVariableDefinition } from './ConnectionUtils';
+import {
+  fetchVariableDefinition,
+  fetchVariableDefinitionByPredicate,
+} from './ConnectionUtils';
 
 const mockRemoveListener = jest.fn();
 const mockSubscribeToFieldUpdates = jest.fn(
@@ -94,3 +97,81 @@ it('throws an Error if variable not found', async () => {
   await expect(fetchPromise).rejects.not.toThrow(TimeoutError);
   expect(mockRemoveListener).toHaveBeenCalled();
 });
+
+describe('fetchVariableDefinitionByPredicate', () => {
+  it('finds definition matching predicate', async () => {
+    const predicate = (def: dh.ide.VariableDefinition) =>
+      def.type === dh.VariableType.FIGURE;
+
+    const fetchPromise = fetchVariableDefinitionByPredicate(
+      connection,
+      predicate
+    );
+
+    expect(mockSubscribeToFieldUpdates).toHaveBeenCalled();
+    expect(mockRemoveListener).not.toHaveBeenCalled();
+
+    const listener = mockSubscribeToFieldUpdates.mock.calls[0][0];
+
+    listener({
+      created: [testDefinition1, testDefinition2],
+      updated: [],
+      removed: [],
+    });
+
+    const result = await fetchPromise;
+
+    expect(result).toBe(testDefinition2);
+    expect(mockRemoveListener).toHaveBeenCalled();
+  });
+
+  it('throws a TimeoutError if finding the variable timed out', async () => {
+    const timeout = 1000;
+    const predicate = (def: dh.ide.VariableDefinition) =>
+      def.title.startsWith('NON_EXISTENT');
+
+    const fetchPromise = fetchVariableDefinitionByPredicate(
+      connection,
+      predicate,
+      timeout,
+      'Custom error message'
+    );
+
+    expect(mockSubscribeToFieldUpdates).toHaveBeenCalled();
+    expect(mockRemoveListener).not.toHaveBeenCalled();
+
+    jest.advanceTimersByTime(timeout + 2000);
+
+    await expect(fetchPromise).rejects.toThrow(TimeoutError);
+    await expect(fetchPromise).rejects.toThrow('Timeout: Custom error message');
+    expect(mockRemoveListener).toHaveBeenCalled();
+  });
+
+  it('throws an Error if no definition matches predicate', async () => {
+    const predicate = (def: dh.ide.VariableDefinition) =>
+      def.title === 'NON_EXISTENT';
+
+    const fetchPromise = fetchVariableDefinitionByPredicate(
+      connection,
+      predicate,
+      undefined,
+      'Custom not found message'
+    );
+
+    expect(mockSubscribeToFieldUpdates).toHaveBeenCalled();
+    expect(mockRemoveListener).not.toHaveBeenCalled();
+
+    const listener = mockSubscribeToFieldUpdates.mock.calls[0][0];
+
+    listener({
+      created: [testDefinition1],
+      updated: [],
+      removed: [],
+    });
+
+    await expect(fetchPromise).rejects.toThrow(Error);
+    await expect(fetchPromise).rejects.toThrow('Custom not found message');
+    await expect(fetchPromise).rejects.not.toThrow(TimeoutError);
+    expect(mockRemoveListener).toHaveBeenCalled();
+  });
+});

packages/jsapi-utils/src/ConnectionUtils.test.ts

@@ -9,33 +9,55 @@ export const FETCH_TIMEOUT = 10_000;
  * @param connection Connection to get the variable from
  * @param name Name of the definition to fetch
  * @param timeout Timeout for the fetch
- * @returns Promise the resolves to the variable definition if found, or rejects if there's an error or the timeout has exceeded
+ * @returns Promise that resolves to the variable definition if found, or rejects if there's an error or the timeout has exceeded
  */
 export function fetchVariableDefinition(
   connection: dh.IdeConnection,
   name: string,
   timeout = FETCH_TIMEOUT
+): Promise<dh.ide.VariableDefinition> {
+  return fetchVariableDefinitionByPredicate(
+    connection,
+    def => def.title === name,
+    timeout,
+    `Variable ${name} not found`
+  );
+}
+
+/**
+ * Fetch the definition for a variable given a connection. Subscribes to field updates and triggers when a variable matching the predicate is found.
+ * @param connection Connection to get the variable from
+ * @param predicate Predicate function to test each variable definition
+ * @param timeout Timeout for the fetch
+ * @param errorMessage Optional error message for timeout and not found errors
+ * @returns Promise that resolves to the variable definition if found, or rejects if there's an error or the timeout has exceeded
+ */
+export function fetchVariableDefinitionByPredicate(
+  connection: dh.IdeConnection,
+  predicate: (definition: dh.ide.VariableDefinition) => boolean,
+  timeout = FETCH_TIMEOUT,
+  errorMessage = 'Variable not found'
 ): Promise<dh.ide.VariableDefinition> {
   return new Promise<dh.ide.VariableDefinition>((resolve, reject) => {
     let removeListener: () => void;
 
     const timeoutId = setTimeout(() => {
       removeListener?.();
-      reject(new TimeoutError(`Timeout looking for variable ${name}`));
+      reject(new TimeoutError(`Timeout: ${errorMessage}`));
     }, timeout);
 
     /**
-     * Checks if the variable we're looking for is in the changes, and resolves the promise if it does
+     * Checks if a variable matching the predicate is in the changes, and resolves the promise if it does
      * @param changes Variables changes that have occurred
      */
     function handleFieldUpdates(changes: dh.ide.VariableChanges): void {
-      const definition = changes.created.find(def => def.title === name);
+      const definition = changes.created.find(predicate);
       clearTimeout(timeoutId);
       removeListener?.();
       if (definition != null) {
         resolve(definition);
       } else {
-        reject(new Error(`Variable ${name} not found`));
+        reject(new Error(errorMessage));
       }
     }