From 3f3d96c4fc1883bdd8131a8fe39aeeb9fb5ca45f Mon Sep 17 00:00:00 2001
From: Rodge Fu <fwj_ac@msn.com>
Date: Tue, 8 Aug 2023 15:09:54 +0800
Subject: [PATCH] Fix the support of generating Example Model of
 DataFactoryKeyVaultReference type of DataFactoryElement (#6561)

* support keyvaultreference dfe type

* update changelog and version

* fix lint error
---
 .../autorest.testmodeler/CHANGELOG.json       | 12 +++++++++
 .../autorest.testmodeler/CHANGELOG.md         |  9 ++++++-
 .../autorest.testmodeler/package.json         |  2 +-
 .../autorest.testmodeler/src/core/model.ts    | 25 ++++++++++++++-----
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.json b/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.json
index b5b6cf95175..2ef04b3f916 100644
--- a/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.json
+++ b/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.json
@@ -1,6 +1,18 @@
 {
   "name": "@autorest/testmodeler",
   "entries": [
+    {
+      "version": "2.6.1",
+      "tag": "@autorest/testmodeler_v2.6.1",
+      "date": "Fri, 21 Jul 2023 08:36:29 GMT",
+      "comments": {
+        "patch": [
+          {
+            "comment": "Support DataFactoryKeyVaultSecretReference of DataFactoryElement"
+          }
+        ]
+      }
+    },
     {
       "version": "2.6.0",
       "tag": "@autorest/testmodeler_v2.6.0",
diff --git a/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.md b/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.md
index ef2635f44ff..c188b218489 100644
--- a/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.md
+++ b/tools/sdk-testgen/packages/autorest.testmodeler/CHANGELOG.md
@@ -1,6 +1,13 @@
 # Change Log - @autorest/testmodeler
 
-This log was last generated on Mon, 19 Jun 2023 08:50:04 GMT and should not be manually modified.
+This log was last generated on Fri, 21 Jul 2023 08:36:29 GMT and should not be manually modified.
+
+## 2.6.1
+Fri, 21 Jul 2023 08:36:29 GMT
+
+### Patches
+
+- Support DataFactoryKeyVaultSecretReference of DataFactoryElement
 
 ## 2.6.0
 Mon, 19 Jun 2023 08:50:04 GMT
diff --git a/tools/sdk-testgen/packages/autorest.testmodeler/package.json b/tools/sdk-testgen/packages/autorest.testmodeler/package.json
index 1e052a69044..c719c66ba0a 100644
--- a/tools/sdk-testgen/packages/autorest.testmodeler/package.json
+++ b/tools/sdk-testgen/packages/autorest.testmodeler/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@autorest/testmodeler",
-    "version": "2.6.0",
+    "version": "2.6.1",
     "description": "Autorest extension for testmodeler",
     "main": "dist/index.js",
     "scripts": {
diff --git a/tools/sdk-testgen/packages/autorest.testmodeler/src/core/model.ts b/tools/sdk-testgen/packages/autorest.testmodeler/src/core/model.ts
index cf9b37de312..eb1373dd99c 100644
--- a/tools/sdk-testgen/packages/autorest.testmodeler/src/core/model.ts
+++ b/tools/sdk-testgen/packages/autorest.testmodeler/src/core/model.ts
@@ -218,7 +218,7 @@ export class ExampleValue {
             const format = extensions[xMsFormat];
             const elementFormat = extensions[xMsFormatElementType];
 
-            const dfeObjSchema = ExampleValue.createSchemaForDfeObject(rawValue, format);
+            const dfeObjSchema = ExampleValue.createSchemaForDfeObject(session, rawValue, format);
             if (dfeObjSchema) {
                 return this.createInstance(session, rawValue, usedProperties, dfeObjSchema, language, undefined, searchDescents);
             } else {
@@ -231,17 +231,30 @@ export class ExampleValue {
         return instance;
     }
 
-    private static createSchemaForDfeObject(raw: any, dfeFormat: string): ObjectSchema | undefined {
+    private static createSchemaForDfeObject(session: Session<TestCodeModel>, raw: any, dfeFormat: string): ObjectSchema | undefined {
         const dfeObjectType = 'type';
         const dfeObjectValue = 'value';
-        const dfeObjectTypeValues = ['Expression', 'SecureString', 'AzureKeyVaultSecretReference'];
         const dfeObjectSchemaPrefix = 'DataFactoryElement-';
 
-        if (Object(raw) && raw[dfeObjectType] && raw[dfeObjectValue] && dfeObjectTypeValues.includes(raw[dfeObjectType])) {
+        if (Object(raw) && raw[dfeObjectType] && raw[dfeObjectValue]) {
             const r = new ObjectSchema(dfeObjectSchemaPrefix + raw[dfeObjectType], '');
             r.addProperty(new Property(dfeObjectType, '', new StringSchema(`${dfeFormat}-${dfeObjectType}`, '')));
-            r.addProperty(new Property(dfeObjectValue, '', new StringSchema(`${dfeFormat}-${dfeObjectValue}`, '')));
-            return r;
+            switch (raw[dfeObjectType]) {
+                case 'Expression':
+                case 'SecureString':
+                    r.addProperty(new Property(dfeObjectValue, '', new StringSchema(`${dfeFormat}-${dfeObjectValue}`, '')));
+                    return r;
+                case 'AzureKeyVaultSecretReference': {
+                    const valueSchema = session.model.schemas.objects.find((s) => s.language.default.name === `AzureKeyVaultSecretReference`);
+                    if (!valueSchema) {
+                        throw new Error('Cant find schema for the value of DataFactoryElement KeyVaultSecret Reference');
+                    }
+                    r.addProperty(new Property(dfeObjectValue, '', valueSchema));
+                    return r;
+                }
+                default:
+                    return undefined;
+            }
         }
         return undefined;
     }