fix(scripts): Improve code sample selection logic (#6018)

Author: MarioAlexandruDan

scripts/config.ts

@@ -49,6 +49,10 @@ export function getSnippetFile(gen: Generator): string {
   return `docs/snippets/${gen.language}/${gen.snippets.outputFolder}/${createClientName(gen.client, gen.language)}${gen.snippets.extension}`;
 }
 
+export function getCTSRequestDir(clientName: string): string {
+  return `tests/CTS/requests/${clientName}`;
+}
+
 export function getDockerImage(language?: Language): string | undefined {
   if (CI || !language) {
     return undefined;

scripts/specs/format.ts

@@ -10,7 +10,7 @@ import { GENERATORS, run, toAbsolutePath } from '../common.ts';
 import { createSpinner } from '../spinners.ts';
 import type { Spec } from '../types.ts';
 
-import { getCodeSampleLabel, transformGeneratedSnippetsToCodeSamples } from './snippets.ts';
+import { CODE_SAMPLE_KEY, getCodeSampleLabel, transformGeneratedSnippetsToCodeSamples } from './snippets.ts';
 
 export async function lintCommon(useCache: boolean): Promise<void> {
   const spinner = createSpinner('linting common spec');
@@ -70,11 +70,15 @@ export async function bundleSpecsForDoc(bundledPath: string, clientName: string)
           specMethod['x-codeSamples'] = [];
         }
 
+        // if a CTS test is marked with isCodeSample: true, it takes priority; otherwise fall back to the first snippet
         if (codeSamples[gen.language][specMethod.operationId]) {
           specMethod['x-codeSamples'].push({
             lang: gen.language,
             label: getCodeSampleLabel(gen.language),
-            source: Object.values(codeSamples[gen.language][specMethod.operationId])[0],
+            source:
+              (Object.hasOwn(codeSamples[gen.language][specMethod.operationId], CODE_SAMPLE_KEY)
+                ? codeSamples[gen.language][specMethod.operationId][CODE_SAMPLE_KEY]
+                : undefined) || Object.values(codeSamples[gen.language][specMethod.operationId])[0],
           });
         }
       }

scripts/specs/snippets.ts

@@ -1,11 +1,15 @@
 import fsp from 'fs/promises';
+import path from 'path';
 
 import { GENERATORS, capitalize, exists, toAbsolutePath } from '../common.ts';
 import type { Language } from '../types.ts';
 
-import { getSnippetFile } from '../config.ts';
+import { getCTSRequestDir, getSnippetFile } from '../config.ts';
 import type { CodeSamples, OpenAPICodeSample, SampleForOperation } from './types.ts';
 
+// Reserved key used to store the preferred code sample (marked with isCodeSample: true in CTS JSON).
+export const CODE_SAMPLE_KEY = 'PREFERRED_SNIPPET';
+
 export function getCodeSampleLabel(language: Language): OpenAPICodeSample['label'] {
   switch (language) {
     case 'csharp':
@@ -52,6 +56,53 @@ export function generateSnippetsJSON(codeSamples: CodeSamples): CodeSamples {
   return codeSamples;
 }
 
+// Reads CTS request JSON files and, for tests marked with `isCodeSample: true`,
+// tags the matching snippet with the reserved CODE_SAMPLE_KEY.
+// Throws if more than one test per operationId is marked as a code sample.
+export async function tagCodeSamples(clientName: string, codeSamples: CodeSamples): Promise<void> {
+  const ctsDir = toAbsolutePath(getCTSRequestDir(clientName));
+
+  if (!(await exists(ctsDir))) {
+    return;
+  }
+
+  const files = await fsp.readdir(ctsDir);
+
+  for (const file of files) {
+    // Use basename to strip any accidental path components returned by readdir.
+    const safeFile = path.basename(file);
+    if (!safeFile.endsWith('.json')) {
+      continue;
+    }
+
+    const operationId = path.basename(safeFile, '.json');
+    const tests: Array<{ testName: string; isCodeSample?: boolean }> = JSON.parse(
+      await fsp.readFile(path.join(ctsDir, safeFile), 'utf8'),
+    );
+
+    const codeSampleTests = tests.filter((t) => t.isCodeSample === true);
+
+    if (codeSampleTests.length > 1) {
+      throw new Error(
+        `Found ${codeSampleTests.length} tests with isCodeSample: true for operationId "${operationId}" in ${clientName}. Only one is allowed.`,
+      );
+    }
+
+    if (codeSampleTests.length === 0) {
+      continue;
+    }
+
+    const testName = codeSampleTests[0].testName;
+
+    for (const lang of Object.keys(codeSamples) as Language[]) {
+      const samplesForOp = codeSamples[lang][operationId];
+      if (samplesForOp !== undefined && Object.hasOwn(samplesForOp, testName)) {
+        samplesForOp[CODE_SAMPLE_KEY] = samplesForOp[testName];
+      }
+    }
+  }
+}
+
 // Reads the generated `docs/snippets/` file for every languages of the given `clientName` and builds an hashmap of snippets per operationId per language.
 export async function transformGeneratedSnippetsToCodeSamples(clientName: string): Promise<CodeSamples> {
   const codeSamples = Object.values(GENERATORS).reduce<CodeSamples>(
@@ -126,5 +177,7 @@ export async function transformGeneratedSnippetsToCodeSamples(clientName: string
     }
   }
 
+  await tagCodeSamples(clientName, codeSamples);
+
   return codeSamples;
 }