[tcgc] allow union variant discriminators (#260)

Azure · Feb 15, 2024 · 3f8486e · 3f8486e
1 parent ee7d779
commit 3f8486e
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 4 deletions.
diff --git a/.chronus/changes/union_discriminator-2024-1-13-16-23-22.md b/.chronus/changes/union_discriminator-2024-1-13-16-23-22.md
@@ -0,0 +1,7 @@
+---
+changeKind: fix
+packages:
+  - "@azure-tools/typespec-client-generator-core"
+---
+
+allow models to have a union variant as a discriminator
diff --git a/packages/typespec-client-generator-core/src/types.ts b/packages/typespec-client-generator-core/src/types.ts
@@ -689,10 +689,7 @@ export function getClientType(context: SdkContext, type: Type, operation?: Opera
       addFormatInfo(context, type, innerType);
       return getKnownValuesEnum(context, type, operation) ?? innerType;
     case "UnionVariant":
-      return {
-        ...getSdkTypeBaseHelper(context, type, "any"),
-        encode: getEncodeHelper(context, type, "any"),
-      };
+      return getClientType(context, type.type, operation);
     case "EnumMember":
       const enumType = getSdkEnum(context, type.enum, operation);
       return getSdkEnumValue(context, enumType, type);

diff --git a/packages/typespec-client-generator-core/test/types.test.ts b/packages/typespec-client-generator-core/test/types.test.ts
@@ -1488,6 +1488,40 @@ describe("typespec-client-generator-core: types", () => {
       strictEqual(sharktypeProperty.type.kind, "string");
     });
 
+    it("union discriminator", async () => {
+      await runner.compileWithBuiltInService(`
+      union KindType {
+        string,
+        shark: "shark",
+        salmon: "salmon"
+      };
+
+      @discriminator("kind")
+      model Fish {
+        age: int32;
+      }
+
+      model Shark extends Fish {
+        kind: KindType.shark;
+        hasFin: boolean;
+      }
+
+      model Salmon extends Fish {
+        kind: KindType.salmon;
+        norweigan: boolean;
+      }
+
+      @get
+      op getModel(): Fish;
+      `);
+      const models = Array.from(getAllModels(runner.context));
+      strictEqual(models.length, 3);
+      const shark = models.find((x) => x.name === "Shark")! as SdkModelType;
+      strictEqual(shark.discriminatorValue, "shark");
+      const salmon = models.find((x) => x.name === "Salmon")! as SdkModelType;
+      strictEqual(salmon.discriminatorValue, "salmon");
+    });
+
     it("filterOutCoreModels true", async () => {
       const runnerWithCore = await createSdkTestRunner({
         librariesToAdd: [AzureCoreTestLibrary],