From 043809128232bc9aaac84bbae3eb35386dc2e66c Mon Sep 17 00:00:00 2001
From: Travis Prescott <trpresco@microsoft.com>
Date: Tue, 27 Feb 2024 15:02:38 -0800
Subject: [PATCH] Refactor logic.

---
 .../src/rules/arm-no-record.ts                | 43 +++++++------------
 .../test/rules/arm-no-record.test.ts          |  2 +-
 2 files changed, 16 insertions(+), 29 deletions(-)

diff --git a/packages/typespec-azure-resource-manager/src/rules/arm-no-record.ts b/packages/typespec-azure-resource-manager/src/rules/arm-no-record.ts
index 9ec08835d8..d4cbf1ed1e 100644
--- a/packages/typespec-azure-resource-manager/src/rules/arm-no-record.ts
+++ b/packages/typespec-azure-resource-manager/src/rules/arm-no-record.ts
@@ -1,4 +1,4 @@
-import { Model, SemanticNodeListener, createRule } from "@typespec/compiler";
+import { DiagnosticTarget, Model, SemanticNodeListener, createRule } from "@typespec/compiler";
 import { getArmResources } from "../resource.js";
 
 export const armNoRecordRule = createRule({
@@ -15,36 +15,23 @@ export const armNoRecordRule = createRule({
   },
   create(context): SemanticNodeListener {
     return {
-      root: (program) => {
-        function isRecordType(model: Model): boolean {
-          return model.name === "Record";
-        }
-
-        function checkModel(model: Model) {
-          if (model.baseModel !== undefined && isRecordType(model.baseModel)) {
-            context.reportDiagnostic({
-              code: "arm-no-record",
-              target: model,
-              messageId: "extends",
-            });
-          }
-          if (model.sourceModel !== undefined && isRecordType(model.sourceModel)) {
+      root: (_) => {
+        function checkModel(model: Model, target: DiagnosticTarget, kind?: "extends" | "is") {
+          if (model.name === "Record") {
             context.reportDiagnostic({
               code: "arm-no-record",
-              target: model,
-              messageId: "is",
+              target: target,
+              messageId: kind || "default",
             });
+          } else if (model.baseModel !== undefined) {
+            checkModel(model.baseModel, model, "extends");
+          } else if (model.sourceModel !== undefined) {
+            checkModel(model.sourceModel, model, "is");
           }
-          for (const prop of model.properties.values()) {
-            if (prop.type.kind === "Model") {
-              if (isRecordType(prop.type)) {
-                context.reportDiagnostic({
-                  code: "arm-no-record",
-                  target: prop,
-                  messageId: "default",
-                });
-              } else {
-                checkModel(prop.type);
+          if (model?.properties !== undefined) {
+            for (const prop of model.properties.values()) {
+              if (prop.type.kind === "Model") {
+                checkModel(prop.type, prop);
               }
             }
           }
@@ -53,7 +40,7 @@ export const armNoRecordRule = createRule({
         // ensure only ARM resources and models they touch are checked
         const resources = getArmResources(context.program);
         for (const resource of resources) {
-          checkModel(resource.typespecType);
+          checkModel(resource.typespecType, resource.typespecType);
         }
       },
     };
diff --git a/packages/typespec-azure-resource-manager/test/rules/arm-no-record.test.ts b/packages/typespec-azure-resource-manager/test/rules/arm-no-record.test.ts
index dc9723731d..c543413846 100644
--- a/packages/typespec-azure-resource-manager/test/rules/arm-no-record.test.ts
+++ b/packages/typespec-azure-resource-manager/test/rules/arm-no-record.test.ts
@@ -201,6 +201,6 @@ it("emits a diagnostic if a deeply aliased model use Record type", async () => {
     .toEmitDiagnostics({
       code: "@azure-tools/typespec-azure-resource-manager/arm-no-record",
       message:
-        "Model properties or operation parameters should not be of type Record. ARM requires Resource provider teams to define types explicitly.",
+        "Models should not equate to type Record. ARM requires Resource provider teams to define types explicitly.",
     });
 });