From cac9c3f851221bebef7ae9a3a0b6cc3b816eb9a3 Mon Sep 17 00:00:00 2001
From: blam <ben@blam.sh>
Date: Wed, 24 Jul 2024 11:48:07 +0200
Subject: [PATCH] chore: added some tests

Signed-off-by: blam <ben@blam.sh>
---
 .../src/schema/createSchemaFromZod.ts         |  5 +-
 .../src/wiring/createExtension.test.ts        | 62 ++++++++++++++++++
 .../src/wiring/createExtension.ts             | 17 +++--
 .../wiring/createExtensionBlueprint.test.tsx  | 65 +++++++++++++++++++
 .../src/wiring/createExtensionBlueprint.ts    | 12 ++--
 5 files changed, 149 insertions(+), 12 deletions(-)

diff --git a/packages/frontend-plugin-api/src/schema/createSchemaFromZod.ts b/packages/frontend-plugin-api/src/schema/createSchemaFromZod.ts
index bfaf57f522..c69ef5dc60 100644
--- a/packages/frontend-plugin-api/src/schema/createSchemaFromZod.ts
+++ b/packages/frontend-plugin-api/src/schema/createSchemaFromZod.ts
@@ -19,7 +19,10 @@ import { z, ZodSchema, ZodTypeDef } from 'zod';
 import zodToJsonSchema from 'zod-to-json-schema';
 import { PortableSchema } from './types';
 
-/** @public */
+/**
+ * @public
+ * @deprecated Use the `config.schema` option of `createExtension` instead, or use `createExtensionBlueprint`.
+ */
 export function createSchemaFromZod<TOutput, TInput>(
   schemaCreator: (zImpl: typeof z) => ZodSchema<TOutput, ZodTypeDef, TInput>,
 ): PortableSchema<TOutput> {
diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
index 4323737540..468da21eb7 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
@@ -291,4 +291,66 @@ describe('createExtension', () => {
       'ExtensionDefinition{namespace=test,attachTo=root@default}',
     );
   });
+
+  it('should create an extension with config', () => {
+    const extension = createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      config: {
+        schema: {
+          foo: z => z.string(),
+          bar: z => z.string().default('bar'),
+          baz: z => z.string().optional(),
+        },
+      },
+      output: {
+        foo: stringData,
+      },
+      factory({ config }) {
+        const a1: string = config.foo;
+        const a2: string = config.bar;
+        // @ts-expect-error
+        const a3: string = config.baz;
+        // @ts-expect-error
+        const c1: number = config.foo;
+        // @ts-expect-error
+        const c2: number = config.bar;
+        // @ts-expect-error
+        const c3: number = config.baz;
+        unused(a1, a2, a3, c1, c2, c3);
+
+        return {
+          foo: 'bar',
+        };
+      },
+    });
+    expect(extension.namespace).toBe('test');
+    expect(String(extension)).toBe(
+      'ExtensionDefinition{namespace=test,attachTo=root@default}',
+    );
+
+    expect(
+      extension.configSchema?.parse({
+        foo: 'x',
+        bar: 'y',
+        baz: 'z',
+        qux: 'w',
+      }),
+    ).toEqual({
+      foo: 'x',
+      bar: 'y',
+      baz: 'z',
+    });
+    expect(
+      extension.configSchema?.parse({
+        foo: 'x',
+      }),
+    ).toEqual({
+      foo: 'x',
+      bar: 'bar',
+    });
+    expect(() => extension.configSchema?.parse({})).toThrow(
+      "Missing required value at 'foo'",
+    );
+  });
 });
diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.ts b/packages/frontend-plugin-api/src/wiring/createExtension.ts
index 6fa3a50634..7783f8e14d 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.ts
@@ -82,7 +82,7 @@ export interface CreateExtensionOptions<
   TOutput extends AnyExtensionDataMap,
   TInputs extends AnyExtensionInputMap,
   TConfig,
-  TConfigSchema extends { [key: string]: z.ZodTypeAny },
+  TConfigSchema extends { [key: string]: z.ZodType },
 > {
   kind?: string;
   namespace?: string;
@@ -100,7 +100,9 @@ export interface CreateExtensionOptions<
   };
   factory(context: {
     node: AppNode;
-    config: TConfig;
+    config: TConfig & {
+      [key in keyof TConfigSchema]: z.infer<TConfigSchema[key]>;
+    };
     inputs: Expand<ResolvedExtensionInputs<TInputs>>;
   }): Expand<ExtensionDataValues<TOutput>>;
 }
@@ -152,12 +154,14 @@ export function createExtension<
   TOutput extends AnyExtensionDataMap,
   TInputs extends AnyExtensionInputMap,
   TConfig,
-  TConfigSchema extends { [key: string]: z.ZodTypeAny },
+  TConfigSchema extends { [key: string]: z.ZodType },
 >(
   options: CreateExtensionOptions<TOutput, TInputs, TConfig, TConfigSchema>,
 ): ExtensionDefinition<TConfig> {
-  // TODO: resovle configSchema from either configSchema or config.schema
   const newConfigSchema = options.config?.schema;
+  if (newConfigSchema && options.configSchema) {
+    throw new Error(`Cannot provide both configSchema and config.schema`);
+  }
   const configSchema = newConfigSchema
     ? createSchemaFromZod(innerZ =>
         innerZ.object(
@@ -179,10 +183,13 @@ export function createExtension<
     inputs: options.inputs ?? {},
     output: options.output,
     configSchema,
-    factory({ inputs, ...rest }) {
+    factory({ inputs, config, ...rest }) {
       // TODO: Simplify this, but TS wouldn't infer the input type for some reason
       return options.factory({
         inputs: inputs as Expand<ResolvedExtensionInputs<TInputs>>,
+        config: config as TConfig & {
+          [key in keyof TConfigSchema]: z.infer<TConfigSchema[key]>;
+        },
         ...rest,
       });
     },
diff --git a/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.test.tsx b/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.test.tsx
index d2373518a3..cf7de0fed3 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.test.tsx
+++ b/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.test.tsx
@@ -20,6 +20,8 @@ import { createExtensionBlueprint } from './createExtensionBlueprint';
 import { createExtensionTester } from '@backstage/frontend-test-utils';
 import { createExtensionDataRef } from './createExtensionDataRef';
 
+function unused(..._any: any[]) {}
+
 describe('createExtensionBlueprint', () => {
   it('should allow creation of extension blueprints', () => {
     const TestExtensionBlueprint = createExtensionBlueprint({
@@ -124,4 +126,67 @@ describe('createExtensionBlueprint', () => {
       data: dataRef,
     });
   });
+
+  it('should allow defining a config schema with additional properties in the instance', () => {
+    const TestExtensionBlueprint = createExtensionBlueprint({
+      kind: 'test-extension',
+      attachTo: { id: 'test', input: 'default' },
+      output: {
+        element: coreExtensionData.reactElement,
+      },
+      config: {
+        schema: {
+          text: z => z.string(),
+        },
+      },
+      factory(_, { config }) {
+        // @ts-expect-error
+        const b = config.something;
+
+        const a: string = config.text;
+        unused(a);
+
+        expect(config.text).toBe('Hello, world!');
+
+        return {
+          element: <h1>{config.text}</h1>,
+        };
+      },
+    });
+
+    const extension = TestExtensionBlueprint.make({
+      name: 'my-extension',
+      params: {
+        text: 'Hello, world!',
+      },
+      config: {
+        schema: {
+          something: z => z.string(),
+          defaulted: z => z.string().optional().default('default'),
+        },
+      },
+      factory(origFactory, { config }) {
+        const b: string = config.something;
+        const c: string = config.text;
+        const d: string = config.defaulted;
+
+        unused(b, c, d);
+
+        expect(config.text).toBe('Hello, world!');
+        expect(config.something).toBe('something new!');
+        expect(config.defaulted).toBe('lolz');
+        return origFactory({});
+      },
+    });
+
+    expect.assertions(4);
+
+    createExtensionTester(extension, {
+      config: {
+        something: 'something new!',
+        text: 'Hello, world!',
+        defaulted: 'lolz',
+      },
+    }).render();
+  });
 });
diff --git a/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.ts b/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.ts
index c54131622f..56edcd8141 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtensionBlueprint.ts
@@ -33,7 +33,7 @@ export interface CreateExtensionBlueprintOptions<
   TParams,
   TInputs extends AnyExtensionInputMap,
   TOutput extends AnyExtensionDataMap,
-  TConfigSchema extends { [key in string]: z.ZodTypeAny },
+  TConfigSchema extends { [key in string]: z.ZodType },
   TDataRefs extends AnyExtensionDataMap,
 > {
   kind: string;
@@ -79,7 +79,7 @@ export interface ExtensionBlueprint<
    * You must either pass `params` directly, or define a `factory` that can
    * optionally call the original factory with the same params.
    */
-  make<TExtensionConfigSchema extends { [key in string]: z.ZodTypeAny }>(
+  make<TExtensionConfigSchema extends { [key in string]: z.ZodType }>(
     args: {
       namespace?: string;
       name?: string;
@@ -135,7 +135,7 @@ class ExtensionBlueprintImpl<
   TParams,
   TInputs extends AnyExtensionInputMap,
   TOutput extends AnyExtensionDataMap,
-  TConfigSchema extends { [key in string]: z.ZodTypeAny },
+  TConfigSchema extends { [key in string]: z.ZodType },
   TDataRefs extends AnyExtensionDataMap,
 > {
   constructor(
@@ -153,7 +153,7 @@ class ExtensionBlueprintImpl<
   dataRefs: TDataRefs;
 
   public make<
-    TExtensionConfigSchema extends { [key in string]: z.ZodTypeAny },
+    TExtensionConfigSchema extends { [key in string]: z.ZodType },
   >(args: {
     namespace?: string;
     name?: string;
@@ -202,7 +202,7 @@ class ExtensionBlueprintImpl<
   > {
     const schema = {
       ...this.options.config?.schema,
-      ...args.config,
+      ...args.config?.schema,
     } as {
       [key in keyof TConfigSchema]: (zImpl: typeof z) => TConfigSchema[key];
     } & {
@@ -266,7 +266,7 @@ export function createExtensionBlueprint<
   TParams,
   TInputs extends AnyExtensionInputMap,
   TOutput extends AnyExtensionDataMap,
-  TConfigSchema extends { [key in string]: z.ZodTypeAny },
+  TConfigSchema extends { [key in string]: z.ZodType },
   TDataRefs extends AnyExtensionDataMap = never,
 >(
   options: CreateExtensionBlueprintOptions<