Normalize processing of links and references

markcowl · Sep 6, 2024 · 7b9a115 · 7b9a115
1 parent d002fb2
commit 7b9a115
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 50 deletions.
diff --git a/packages/typespec-azure-core/src/lro-helpers.ts b/packages/typespec-azure-core/src/lro-helpers.ts
@@ -313,15 +313,18 @@ export function getLroMetadata(program: Program, operation: Operation): LroMetad
   if (context === undefined) return undefined;
   processFinalReference(program, operation, context);
   processFinalLink(program, operation, context);
-  const nextReference: NextOperationReference | undefined = processStatusMonitorReference(
-    program,
-    operation,
-    context
-  );
+  const nextReference:
+    | NextOperationReference
+    | (NextOperationLink & { operation: Operation })
+    | undefined = processStatusMonitorReference(program, operation, context);
   if (nextReference !== undefined && nextReference.responseModel.kind === "Model") {
     context.statusMonitorStep = nextReference;
-    processFinalReference(program, nextReference.target.operation, context);
-    processFinalLink(program, nextReference.target.operation, context);
+    const linkedOperation =
+      nextReference.kind === "nextOperationLink"
+        ? nextReference.operation
+        : nextReference.target.operation;
+    processFinalReference(program, linkedOperation, context);
+    processFinalLink(program, linkedOperation, context);
     context.pollingStep = getPollingStep(program, nextReference.responseModel, context);
     return createLroMetadata(program, operation, context);
   }
@@ -466,7 +469,9 @@ function createOperationLink(program: Program, modelProperty: ModelProperty): Op
   };
 }
 
-function createOperationReference(metadata: OperationLinkMetadata): OperationReference | undefined {
+function createOperationReferenceOrLink(
+  metadata: OperationLinkMetadata
+): OperationReference | OperationLink | undefined {
   if (!metadata.parameterMap && !metadata.link) return undefined;
   const map = new Map<string, ParameterSource>();
   if (metadata.parameterMap) {
@@ -492,9 +497,9 @@ function createOperationReference(metadata: OperationLinkMetadata): OperationRef
   }
   if (metadata.link) {
     return {
-      kind: "reference",
-      operation: metadata.linkedOperation,
-      link: metadata.link,
+      kind: "link",
+      location: metadata.link.location,
+      property: metadata.link.property,
     };
   }
 
@@ -992,13 +997,24 @@ function processFinalReference(program: Program, operation: Operation, context:
   )
     return;
   context.finalOperationLink = link;
-  const reference = createOperationReference(link);
+  const reference = createOperationReferenceOrLink(link);
   if (reference === undefined) return;
-  context.finalStep = {
-    kind: "finalOperationReference",
-    responseModel: link.result?.type,
-    target: reference,
-  };
+  switch (reference.kind) {
+    case "reference":
+      context.finalStep = {
+        kind: "finalOperationReference",
+        responseModel: link.result?.type,
+        target: reference,
+      };
+      break;
+    case "link":
+      context.finalStep = {
+        kind: "finalOperationLink",
+        responseModel: link.result?.type,
+        target: reference,
+      };
+      break;
+  }
 }
 
 function createStatusMonitorPollingData(data: StatusMonitorMetadata): StatusMonitorInfo {
@@ -1081,7 +1097,7 @@ function processStatusMonitorReference(
   program: Program,
   referencedOperation: Operation,
   context: LroContext
-): NextOperationReference | undefined {
+): NextOperationReference | (NextOperationLink & { operation: Operation }) | undefined {
   const references: Map<string, OperationLinkMetadata> | undefined = getOperationLinks(
     program,
     referencedOperation
@@ -1091,19 +1107,29 @@ function processStatusMonitorReference(
   const pollingData: OperationLinkMetadata | undefined = references.get(PollingOperationKey);
   if (pollingData === undefined) return undefined;
   context.pollingOperationLink = pollingData;
-  const pollingReference = createOperationReference(pollingData);
+  const pollingReference = createOperationReferenceOrLink(pollingData);
   if (pollingReference === undefined) return undefined;
   context.statusMonitorInfo = pollingData.result?.statusMonitor;
   const finalData: OperationLinkMetadata | undefined = references.get(FinalOperationKey);
   if (context.finalStep === undefined && finalData !== undefined) {
-    const finalReference = createOperationReference(finalData);
+    const finalReference = createOperationReferenceOrLink(finalData);
     const finalModel = finalData?.result?.type;
     if (finalReference !== undefined && finalModel !== undefined) {
-      context.finalStep = {
-        kind: "finalOperationReference",
-        responseModel: finalModel,
-        target: finalReference,
-      };
+      switch (finalReference.kind) {
+        case "reference":
+          context.finalStep = {
+            kind: "finalOperationReference",
+            responseModel: finalModel,
+            target: finalReference,
+          };
+          break;
+        case "link":
+          context.finalStep = {
+            kind: "finalOperationLink",
+            responseModel: finalModel,
+            target: finalReference,
+          };
+      }
     }
   }
   if (
@@ -1120,11 +1146,21 @@ function processStatusMonitorReference(
   }
   const responseModel = pollingData.result?.type;
   if (responseModel === undefined) return undefined;
-  return {
-    kind: "nextOperationReference",
-    responseModel: responseModel,
-    target: pollingReference,
-  };
+  switch (pollingReference.kind) {
+    case "reference":
+      return {
+        kind: "nextOperationReference",
+        responseModel: responseModel,
+        target: pollingReference,
+      };
+    case "link":
+      return {
+        kind: "nextOperationLink",
+        responseModel: responseModel,
+        target: pollingReference,
+        operation: pollingData.linkedOperation,
+      };
+  }
 }
 
 function resolveOperationLocation(

diff --git a/packages/typespec-azure-core/test/operations.test.ts b/packages/typespec-azure-core/test/operations.test.ts
@@ -891,7 +891,7 @@ describe("typespec-azure-core: operation templates", () => {
         "createOrUpdate"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "ResourceOperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -943,7 +943,7 @@ describe("typespec-azure-core: operation templates", () => {
       );
 
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "ResourceOperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1015,7 +1015,7 @@ op startJobAsync(
       );
 
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "OperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1091,7 +1091,7 @@ op createJob(
       );
 
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "OperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1146,7 +1146,7 @@ op createJob(
         "delete"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "ResourceOperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1205,7 +1205,7 @@ op createJob(
         "update"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "ResourceOperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1273,7 +1273,7 @@ op createJob(
         "doAction"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "ResourceOperationStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1315,7 +1315,7 @@ op createJob(
         "lrRpcOp"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationReference");
+      deepStrictEqual(metadata.statusMonitorStep?.kind, "nextOperationLink");
       deepStrictEqual(metadata.statusMonitorStep?.responseModel.name, "PollingStatus");
 
       deepStrictEqual(metadata.pollingInfo.kind, "pollingOperationStep");
@@ -1510,9 +1510,9 @@ op createJob(
         "createWidget"
       );
       ok(metadata);
-      deepStrictEqual(metadata.statusMonitorStep?.target.kind, "reference");
-      deepStrictEqual((metadata.statusMonitorStep?.target.link as any).location, "ResponseHeader");
-      deepStrictEqual((metadata.statusMonitorStep?.target.link as any).property.name, "opLink");
+      deepStrictEqual(metadata.statusMonitorStep?.target.kind, "link");
+      deepStrictEqual((metadata.statusMonitorStep?.target as any).location, "ResponseHeader");
+      deepStrictEqual((metadata.statusMonitorStep?.target as any).property.name, "opLink");
 
       deepStrictEqual(metadata.pollingInfo?.responseModel.name, "PollingStatus");
       deepStrictEqual(metadata.pollingInfo?.terminationStatus.kind, "model-property");
@@ -2140,9 +2140,9 @@ op createJob(
       deepStrictEqual(metadata?.finalResultPath, undefined);
 
       ok(metadata.statusMonitorStep);
-      deepStrictEqual(metadata.statusMonitorStep.target.kind, "reference");
-      deepStrictEqual((metadata.statusMonitorStep.target.link as any)?.location, "ResponseHeader");
-      deepStrictEqual((metadata.statusMonitorStep.target.link as any)?.property.name, "opLink");
+      deepStrictEqual(metadata.statusMonitorStep.target.kind, "link");
+      deepStrictEqual((metadata.statusMonitorStep.target as any)?.location, "ResponseHeader");
+      deepStrictEqual((metadata.statusMonitorStep.target as any)?.property.name, "opLink");
 
       ok(metadata.pollingInfo);
       deepStrictEqual(metadata.pollingInfo.responseModel.name, "PollingStatus");
@@ -2200,9 +2200,9 @@ op createJob(
       deepStrictEqual(metadata?.finalResultPath, undefined);
 
       ok(metadata.statusMonitorStep);
-      deepStrictEqual(metadata.statusMonitorStep.target.kind, "reference");
-      deepStrictEqual((metadata.statusMonitorStep.target.link as any)?.location, "ResponseHeader");
-      deepStrictEqual((metadata.statusMonitorStep.target.link as any)?.property.name, "opLink");
+      deepStrictEqual(metadata.statusMonitorStep.target.kind, "link");
+      deepStrictEqual((metadata.statusMonitorStep.target as any)?.location, "ResponseHeader");
+      deepStrictEqual((metadata.statusMonitorStep.target as any)?.property.name, "opLink");
 
       ok(metadata.pollingInfo);
       deepStrictEqual(metadata.pollingInfo.responseModel.name, "PollingStatus");
@@ -2365,7 +2365,7 @@ op createJob(
       deepStrictEqual(metadata?.logicalPath, undefined);
 
       ok(metadata.statusMonitorStep);
-      deepStrictEqual(metadata.statusMonitorStep.target.kind, "reference");
+      deepStrictEqual(metadata.statusMonitorStep.target.kind, "link");
 
       ok(metadata.finalStep);
       deepStrictEqual(metadata.finalStep.kind, "noPollingResult");
@@ -2422,7 +2422,7 @@ op createJob(
       deepStrictEqual(metadata.finalStep.target?.kind, "reference");
       deepStrictEqual(metadata.finalStep.target?.operation.name, "getWidget");
       ok(metadata.statusMonitorStep);
-      deepStrictEqual(metadata.statusMonitorStep.target.kind, "reference");
+      deepStrictEqual(metadata.statusMonitorStep.target.kind, "link");
       ok(metadata.pollingInfo);
       deepStrictEqual(metadata.pollingInfo.responseModel.name, "PollingStatus");
       deepStrictEqual(metadata.pollingInfo.terminationStatus.kind, "model-property");
@@ -2484,7 +2484,7 @@ op createJob(
       deepStrictEqual(metadata.finalStep.target?.kind, "reference");
       deepStrictEqual(metadata.finalStep.target?.operation.name, "getWidget");
       ok(metadata.statusMonitorStep);
-      deepStrictEqual(metadata.statusMonitorStep.target.kind, "reference");
+      deepStrictEqual(metadata.statusMonitorStep.target.kind, "link");
       ok(metadata.pollingInfo);
       deepStrictEqual(metadata.pollingInfo.responseModel.name, "PollingStatus");
       deepStrictEqual(metadata.pollingInfo.terminationStatus.kind, "model-property");