chore: update the types a little bit to help with parsing

Signed-off-by: blam <ben@blam.sh>

Signed-off-by: blam <ben@blam.sh>

packages/frontend-plugin-api/src/blueprints/SignInPageBlueprint.test.tsx

@@ -62,7 +62,7 @@ describe('SignInPageBlueprint', () => {
 
     const tester = createExtensionTester(extension);
 
-    const Element = tester.get(SignInPageBlueprint.dataRefs.component)!;
+    const Element = tester.get(SignInPageBlueprint.dataRefs.component);
 
     expect(Element).toBeDefined();
 

packages/frontend-test-utils/src/app/createExtensionTester.test.tsx

@@ -240,8 +240,8 @@ describe('createExtensionTester', () => {
     const extension = createExtension({
       namespace: 'test',
       attachTo: { id: 'ignored', input: 'ignored' },
-      output: { text: stringDataRef },
-      factory: () => ({ text: 'test-text' }),
+      output: [stringDataRef],
+      factory: () => [stringDataRef('test-text')],
     });
 
     const tester = createExtensionTester(extension);
@@ -297,6 +297,76 @@ describe('createExtensionTester', () => {
     );
   });
 
+  it('should not allow getting extension data for an output that was not defined in the extension', () => {
+    const internalRef = createExtensionDataRef<number>().with({
+      id: 'test.internal',
+    });
+
+    const internalRef2 = createExtensionDataRef<number>().with({
+      id: 'test.internal2',
+    });
+
+    const extension = createExtension({
+      namespace: 'test',
+      name: 'e1',
+      attachTo: { id: 'ignored', input: 'ignored' },
+      output: [stringDataRef, internalRef.optional()],
+      factory: () => [stringDataRef('test-text')],
+    });
+
+    const tester = createExtensionTester(extension);
+
+    const test: string = tester.get(stringDataRef);
+
+    // @ts-expect-error - internalRef is optional
+    const test2: number = tester.get(internalRef);
+
+    // @ts-expect-error - internalRef2 is not defined in the extension
+    const test3: number = tester.get(internalRef2);
+
+    expect([test, test2, test3]).toBeDefined();
+  });
+
+  it('should support getting outputs from a query response', () => {
+    const internalRef = createExtensionDataRef<number>().with({
+      id: 'test.internal',
+    });
+
+    const internalRef2 = createExtensionDataRef<number>().with({
+      id: 'test.internal2',
+    });
+
+    const extension = createExtension({
+      namespace: 'test',
+      name: 'e1',
+      attachTo: { id: 'ignored', input: 'ignored' },
+      output: [coreExtensionData.reactElement],
+      factory: () => [coreExtensionData.reactElement(<div>bob</div>)],
+    });
+
+    const extraExtension = createExtension({
+      namespace: 'test',
+      name: 'e1',
+      attachTo: { id: 'ignored', input: 'ignored' },
+      output: [stringDataRef, internalRef.optional()],
+      factory: () => [stringDataRef('test-text')],
+    });
+
+    const tester = createExtensionTester(extension)
+      .add(extraExtension)
+      .query(extraExtension);
+
+    const test: string = tester.get(stringDataRef);
+
+    // @ts-expect-error - internalRef is optional
+    const test2: number = tester.get(internalRef);
+
+    // @ts-expect-error - internalRef2 is not defined in the extension
+    const test3: number = tester.get(internalRef2);
+
+    expect([test, test2, test3]).toBeDefined();
+  });
+
   // TODO: this should be implemented
   // eslint-disable-next-line jest/no-disabled-tests
   it.skip('should allow querying an extension and getting outputs', () => {
@@ -328,7 +398,9 @@ describe('createExtensionTester', () => {
 
     const tester = createExtensionTester(extension).add(extension2);
 
+    // @ts-expect-error
     expect(tester.query(extension).get(stringDataRef)).toBe('nest-test-text');
+    // @ts-expect-error
     expect(tester.query(extension2).get(stringDataRef)).toBe('test-text');
     // @ts-expect-error
     expect(tester.query(extension).input('input').data(stringDataRef)).toBe(
@@ -362,6 +434,7 @@ describe('createExtensionTester', () => {
       inputs: { input: 'test-text' },
     });
 
+    // @ts-expect-error
     expect(tester.query(extension).get(stringDataRef)).toBe('nest-test-text');
   });
 });

packages/frontend-test-utils/src/app/createExtensionTester.tsx

@@ -19,6 +19,7 @@ import { MemoryRouter, Link } from 'react-router-dom';
 import { RenderResult, render } from '@testing-library/react';
 import { createSpecializedApp } from '@backstage/frontend-app-api';
 import {
+  AnyExtensionDataRef,
   AppNode,
   AppTree,
   Extension,
@@ -101,7 +102,7 @@ const TestAppNavExtension = createExtension({
 });
 
 /** @public */
-export class ExtensionQuery {
+export class ExtensionQuery<UOutput extends AnyExtensionDataRef> {
   #node: AppNode;
 
   constructor(node: AppNode) {
@@ -124,18 +125,24 @@ export class ExtensionQuery {
     return instance;
   }
 
-  get<T>(ref: ExtensionDataRef<T>): T | undefined {
+  get<TId extends UOutput['id']>(
+    ref: ExtensionDataRef<any, TId, any>,
+  ): UOutput extends ExtensionDataRef<infer IData, TId, infer IConfig>
+    ? IConfig['optional'] extends true
+      ? IData | undefined
+      : IData
+    : never {
     return this.instance.getData(ref);
   }
 }
 
 /** @public */
-export class ExtensionTester {
+export class ExtensionTester<UOutput extends AnyExtensionDataRef> {
   /** @internal */
-  static forSubject<TConfig, TConfigInput>(
+  static forSubject<TConfig, TConfigInput, UOutput extends AnyExtensionDataRef>(
     subject: ExtensionDefinition<TConfig, TConfigInput>,
     options?: { config?: TConfigInput },
-  ): ExtensionTester {
+  ): ExtensionTester<UOutput> {
     const tester = new ExtensionTester();
     const internal = toInternalExtensionDefinition(subject);
 
@@ -192,7 +199,7 @@ export class ExtensionTester {
   add<TConfig, TConfigInput>(
     extension: ExtensionDefinition<TConfig, TConfigInput>,
     options?: { config?: TConfigInput },
-  ): ExtensionTester {
+  ): ExtensionTester<UOutput> {
     if (this.#tree) {
       throw new Error(
         'Cannot add more extensions accessing the extension tree',
@@ -219,17 +226,24 @@ export class ExtensionTester {
     return this;
   }
 
-  get<T>(ref: ExtensionDataRef<T>): T | undefined {
+  get<TId extends UOutput['id']>(
+    ref: ExtensionDataRef<any, TId, any>,
+  ): UOutput extends ExtensionDataRef<infer IData, TId, infer IConfig>
+    ? IConfig['optional'] extends true
+      ? IData | undefined
+      : IData
+    : never {
     const tree = this.#resolveTree();
 
     return new ExtensionQuery(tree.root).get(ref);
   }
 
-  query(id: string | ExtensionDefinition<any, any>): ExtensionQuery {
+  query<UQueryExtensionOutput extends AnyExtensionDataRef>(
+    extension: ExtensionDefinition<any, any, UQueryExtensionOutput>,
+  ): ExtensionQuery<UQueryExtensionOutput> {
     const tree = this.#resolveTree();
 
-    const actualId =
-      typeof id === 'string' ? id : resolveExtensionDefinition(id).id;
+    const actualId = resolveExtensionDefinition(extension).id;
 
     const node = tree.nodes.get(actualId);
 
@@ -355,9 +369,13 @@ export class ExtensionTester {
 }
 
 /** @public */
-export function createExtensionTester<TConfig>(
-  subject: ExtensionDefinition<TConfig>,
-  options?: { config?: TConfig },
-): ExtensionTester {
+export function createExtensionTester<
+  TConfig,
+  TConfigInput,
+  UOutput extends AnyExtensionDataRef,
+>(
+  subject: ExtensionDefinition<TConfig, TConfigInput, UOutput>,
+  options?: { config?: TConfigInput },
+): ExtensionTester<UOutput> {
   return ExtensionTester.forSubject(subject, options);
 }