[tcgc] map discriminator string value type to enum value type (#834)

fix: Azure/autorest.go#1206
Azure · May 15, 2024 · c9398d0 · c9398d0
1 parent 242c8b7
commit c9398d0
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 3 deletions.
diff --git a/.chronus/changes/fix_value_discriminator-2024-4-15-10-46-20.md b/.chronus/changes/fix_value_discriminator-2024-4-15-10-46-20.md
@@ -0,0 +1,7 @@
+---
+changeKind: fix
+packages:
+  - "@azure-tools/typespec-client-generator-core"
+---
+
+map discriminator string value type to enum value type
diff --git a/packages/typespec-client-generator-core/src/types.ts b/packages/typespec-client-generator-core/src/types.ts
@@ -424,6 +424,14 @@ function addDiscriminatorToModelType(
   const discriminator = getDiscriminator(context.program, type);
   const diagnostics = createDiagnosticCollector();
   if (discriminator) {
+    let discriminatorType: SdkType | undefined = undefined;
+    for (let i = 0; i < model.properties.length; i++) {
+      const property = model.properties[i];
+      if (property.kind === "property" && property.__raw?.name === discriminator.propertyName) {
+        discriminatorType = property.type;
+      }
+    }
+
     let discriminatorProperty;
     for (const childModel of type.derivedModels) {
       const childModelSdkType = diagnostics.pipe(getSdkModelWithDiagnostics(context, childModel));
@@ -450,12 +458,20 @@ function addDiscriminatorToModelType(
                 })
               );
             } else {
-              childModelSdkType.discriminatorValue = property.type.value;
+              // map string value type to enum value type
+              if (property.type.kind === "constant" && discriminatorType?.kind === "enum") {
+                for (const value of discriminatorType.values) {
+                  if (value.value === property.type.value) {
+                    property.type = value;
+                  }
+                }
+              }
+              childModelSdkType.discriminatorValue = property.type.value as string;
               property.discriminator = true;
               if (model.discriminatedSubtypes === undefined) {
                 model.discriminatedSubtypes = {};
               }
-              model.discriminatedSubtypes[property.type.value] = childModelSdkType;
+              model.discriminatedSubtypes[property.type.value as string] = childModelSdkType;
               discriminatorProperty = property;
             }
           }
@@ -470,7 +486,7 @@ function addDiscriminatorToModelType(
         return diagnostics.wrap(undefined);
       }
     }
-    let discriminatorType: SdkType;
+
     if (discriminatorProperty) {
       if (discriminatorProperty.type.kind === "constant") {
         discriminatorType = { ...discriminatorProperty.type.valueType };

diff --git a/packages/typespec-client-generator-core/test/types.test.ts b/packages/typespec-client-generator-core/test/types.test.ts
@@ -2599,6 +2599,63 @@ describe("typespec-client-generator-core: types", () => {
       strictEqual(kindType.isFixed, false);
     });
 
+    it("string discriminator map to enum value", async () => {
+      await runner.compileWithBuiltInService(`
+      union KindType {
+        string,
+        shark: "shark",
+        salmon: "salmon"
+      };
+
+      @discriminator("kind")
+      model Fish {
+        kind: KindType;
+        age: int32;
+      }
+
+      model Shark extends Fish {
+        kind: "shark";
+        hasFin: boolean;
+      }
+
+      model Salmon extends Fish {
+        kind: "salmon";
+        norweigan: boolean;
+      }
+
+      @get
+      op getModel(): Fish;
+      `);
+      const models = runner.context.experimental_sdkPackage.models;
+      strictEqual(models.length, 3);
+      const fish = models.find((x) => x.name === "Fish");
+      ok(fish);
+      let kindTypeProperty = fish.properties.find((x) => x.name === "kind");
+      ok(kindTypeProperty);
+      strictEqual(kindTypeProperty.type.kind, "enum");
+      strictEqual(kindTypeProperty.type.isUnionAsEnum, true);
+      strictEqual(fish.discriminatorProperty, kindTypeProperty);
+      const shark = models.find((x) => x.name === "Shark");
+      ok(shark);
+      strictEqual(shark.discriminatorValue, "shark");
+      kindTypeProperty = shark.properties.find((x) => x.name === "kind");
+      ok(kindTypeProperty);
+      strictEqual(kindTypeProperty.type.kind, "enumvalue");
+      const salmon = models.find((x) => x.name === "Salmon");
+      ok(salmon);
+      kindTypeProperty = salmon.properties.find((x) => x.name === "kind");
+      ok(kindTypeProperty);
+      strictEqual(kindTypeProperty.type.kind, "enumvalue");
+      strictEqual(salmon.discriminatorValue, "salmon");
+
+      strictEqual(runner.context.experimental_sdkPackage.enums.length, 1);
+      const kindType = runner.context.experimental_sdkPackage.enums.find(
+        (x) => x.name === "KindType"
+      );
+      ok(kindType);
+      strictEqual(kindType.isFixed, false);
+    });
+
     it("discriminator rename", async () => {
       await runner.compileWithBuiltInService(`
       @discriminator("kind")