frontend-plugin-api: extension data container runtime + tests

Co-authored-by: Fredrik Adelöw <freben@gmail.com>
Co-authored-by: blam <ben@blam.sh>
Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>

packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.test.tsx

@@ -21,10 +21,21 @@ import { createExtensionTester } from '@backstage/frontend-test-utils';
 import { createExtensionDataRef } from './createExtensionDataRef';
 import { createExtensionInput } from './createExtensionInput';
 import { RouteRef } from '../routing';
-import { toInternalExtensionDefinition } from './createExtension';
+import {
+  ExtensionDefinition,
+  toInternalExtensionDefinition,
+} from './createExtension';
 
 function unused(..._any: any[]) {}
 
+function factoryOutput(ext: ExtensionDefinition<any, any>) {
+  const int = toInternalExtensionDefinition(ext);
+  if (int.version !== 'v2') {
+    throw new Error('Expected v2 extension');
+  }
+  return Array.from(int.factory({} as any));
+}
+
 describe('createExtensionBlueprint', () => {
   it('should allow creation of extension blueprints', () => {
     const TestExtensionBlueprint = createExtensionBlueprint({
@@ -476,6 +487,8 @@ describe('createExtensionBlueprint', () => {
     );
 
     expect(ext.output).toEqual([testDataRef2]);
+
+    expect(factoryOutput(ext)).toEqual([testDataRef2('foobar')]);
   });
 
   it('should allow returning of the parent data container', () => {
@@ -491,23 +504,28 @@ describe('createExtensionBlueprint', () => {
       },
     });
 
-    blueprint.make({
-      output: [testDataRef1, testDataRef2],
-      *factory(origFactory) {
-        yield* origFactory({});
-        yield testDataRef2('bar');
-      },
-    });
+    expect(
+      factoryOutput(
+        blueprint.make({
+          output: [testDataRef1, testDataRef2],
+          *factory(origFactory) {
+            yield* origFactory({});
+            yield testDataRef2('bar');
+          },
+        }),
+      ),
+    ).toEqual([testDataRef1('foo'), testDataRef2('bar')]);
 
-    blueprint.make({
-      output: [testDataRef1, testDataRef2],
-      factory(origFactory) {
-        return [...origFactory({}), testDataRef2('bar')];
-      },
-    });
-
-    expect(true).toBe(true);
-    // todo: test that the data is actually available
+    expect(
+      factoryOutput(
+        blueprint.make({
+          output: [testDataRef1, testDataRef2],
+          factory(origFactory) {
+            return [...origFactory({}), testDataRef2('bar')];
+          },
+        }),
+      ),
+    ).toEqual([testDataRef1('foo'), testDataRef2('bar')]);
   });
 
   it('should not allow returning parent output if outputs are overridden', () => {
@@ -523,13 +541,58 @@ describe('createExtensionBlueprint', () => {
       },
     });
 
+    // @ts-expect-error
     blueprint.make({
       output: [testDataRef2.optional()],
       *factory() {
+        yield testDataRef1('foo');
         yield testDataRef2('bar');
       },
     });
 
-    expect(true).toBe(true);
+    expect(
+      factoryOutput(
+        blueprint.make({
+          output: [testDataRef2.optional()],
+          *factory() {
+            yield testDataRef2('bar');
+          },
+        }),
+      ),
+    ).toEqual([testDataRef2('bar')]);
+  });
+
+  it('should not rely on optional outputs when forwarding from parent', () => {
+    const testDataRef1 = createExtensionDataRef<string>().with({ id: 'test1' });
+    const testDataRef2 = createExtensionDataRef<string>().with({ id: 'test2' });
+
+    const blueprint = createExtensionBlueprint({
+      kind: 'test-extension',
+      attachTo: { id: 'test', input: 'default' },
+      output: [testDataRef1, testDataRef2.optional()],
+      factory() {
+        return [testDataRef1('foo')];
+      },
+    });
+
+    // @ts-expect-error
+    blueprint.make({
+      output: [testDataRef1, testDataRef2],
+      *factory(origFactory) {
+        yield* origFactory({});
+      },
+    });
+
+    expect(
+      factoryOutput(
+        blueprint.make({
+          output: [testDataRef1, testDataRef2],
+          *factory(origFactory) {
+            yield* origFactory({});
+            yield testDataRef2('bar');
+          },
+        }),
+      ),
+    ).toEqual([testDataRef1('foo'), testDataRef2('bar')]);
   });
 });

packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.ts

@@ -28,6 +28,7 @@ import { z } from 'zod';
 import { ExtensionInput } from './createExtensionInput';
 import {
   AnyExtensionDataRef,
+  ExtensionDataRef,
   ExtensionDataValue,
 } from './createExtensionDataRef';
 
@@ -169,6 +170,30 @@ export interface ExtensionBlueprint<
   >;
 }
 
+/** @internal */
+function createDataContainer<UData extends AnyExtensionDataRef>(
+  values: Iterable<
+    UData extends ExtensionDataRef<infer IData, infer IId>
+      ? ExtensionDataValue<IData, IId>
+      : never
+  >,
+): ExtensionDataContainer<UData> {
+  const container = new Map<string, ExtensionDataValue<any, any>>();
+
+  for (const output of values) {
+    container.set(output.id, output);
+  }
+
+  return {
+    get(ref) {
+      return container.get(ref.id)?.value;
+    },
+    [Symbol.iterator]() {
+      return container.values();
+    },
+  } as ExtensionDataContainer<UData>;
+}
+
 /**
  * @internal
  */
@@ -312,11 +337,13 @@ class ExtensionBlueprintImpl<
                 inputs?: Expand<ResolvedExtensionInputs<TInputs>>;
               },
             ): ExtensionDataContainer<UOutput> => {
-              return this.options.factory(innerParams, {
-                node,
-                config: innerContext?.config ?? config,
-                inputs: (innerContext?.inputs ?? inputs) as any, // TODO: Fix the way input values are overridden
-              });
+              return createDataContainer<UOutput>(
+                this.options.factory(innerParams, {
+                  node,
+                  config: innerContext?.config ?? config,
+                  inputs: (innerContext?.inputs ?? inputs) as any, // TODO: Fix the way input values are overridden
+                }),
+              );
             },
             {
               node,