lib/apiScenario/gen/restlerApiScenarioGenerator.ts

import * as path from "path";
import Heap from "heap";
import { inject, injectable } from "inversify";
import { dump } from "js-yaml";
import { pathJoin, pathResolve, urlParse } from "@azure-tools/openapi-tools-common";
import { cloneDeep } from "lodash";
import { inversifyGetInstance, TYPES } from "../../inversifyUtils";
import { FileLoader } from "../../swagger/fileLoader";
import { JsonLoader } from "../../swagger/jsonLoader";
import { SwaggerLoader } from "../../swagger/swaggerLoader";
import {
  SwaggerSpec,
  LowerHttpMethods,
  Schema,
  Parameter,
  Operation,
} from "../../swagger/swaggerTypes";
import { traverseSwagger } from "../../transform/traverseSwagger";
import { ApiScenarioLoaderOption } from "../apiScenarioLoader";
import {
  RawScenario,
  RawScenarioDefinition,
  RawStepExample,
  RawStepOperation,
  Variable,
  VarValue,
} from "../apiScenarioTypes";
import * as util from "../../generator/util";
import { setDefaultOpts } from "../../swagger/loader";
import Mocker from "../../generator/mocker";
import { logger } from ".././logger";
import { xmsExamples, xmsSkipUrlEncoding } from "../../util/constants";
import { ApiScenarioYamlLoader } from "../apiScenarioYamlLoader";
import { ArmResourceManipulator } from "./ApiTestRuleBasedGenerator";

export interface ApiScenarioGeneratorOption extends ApiScenarioLoaderOption {
  swaggerFilePaths: string[];
  dependencyPath: string;
  outputDir: string;
  useExample?: boolean;
  scope?: string;
}

interface Dependency {
  producer_endpoint: string;
  producer_method: string;
  producer_resource_name: string;
  consumer_param: string;
}

interface Dependencies {
  [path: string]: {
    [method: string]: {
      Path?: Dependency[];
      Query?: Dependency[];
    };
  };
}

interface Node {
  operationId: string;
  method: LowerHttpMethods;
  children: Map<string, Node>;
  inDegree: number;
  outDegree: number;
  visited: boolean;
  priority: number;
}

type ParameterNode = {
  operationIds: Set<string>;
  count: number;
  name: string;
  value: any;
};

const methodOrder: LowerHttpMethods[] = ["put", "get", "patch", "post", "delete"];

const envVariables = ["api-version", "subscriptionId", "resourceGroupName", "location"];

export const useRandom = {
  flag: true,
};

@injectable()
export class RestlerApiScenarioGenerator {
  private swaggers: SwaggerSpec[];
  private graph: Map<string, Node>;
  private parameters: Map<string, ParameterNode>;
  private mocker: any;
  private operations: Map<string, Operation>;

  public constructor(
    @inject(TYPES.opts) private opts: ApiScenarioGeneratorOption,
    private swaggerLoader: SwaggerLoader,
    private fileLoader: FileLoader,
    private jsonLoader: JsonLoader,
    private apiScenarioYamlLoader: ApiScenarioYamlLoader
  ) {
    this.swaggers = [];
    this.mocker = useRandom.flag
      ? new Mocker()
      : {
          mock: (paramSpec: any): any => {
            switch (paramSpec.type) {
              case "string":
                return "test";
              case "integer":
                return 1;
              case "number":
                return 1;
              case "boolean":
                return true;
              case "array":
                return [];
            }
          },
        };
  }

  public static create(opts: ApiScenarioGeneratorOption) {
    setDefaultOpts(opts, {
      swaggerFilePaths: [],
      outputDir: ".",
      dependencyPath: "",
      eraseXmsExamples: false,
      skipResolveRefKeys: [xmsExamples],
    });
    return inversifyGetInstance(RestlerApiScenarioGenerator, opts);
  }

  public async initialize() {
    this.opts.outputDir = pathResolve(this.opts.outputDir);
    this.opts.swaggerFilePaths = this.opts.swaggerFilePaths.map((p) => pathResolve(p));
    this.operations = new Map();
    for (const path of this.opts.swaggerFilePaths) {
      const swagger = await this.swaggerLoader.load(path);
      this.swaggers.push(swagger);
      traverseSwagger(swagger, {
        onOperation: (operation) => {
          this.operations.set(operation.operationId!, operation);
        },
      });
    }
    await this.generateGraph();
  }

  public async generateResourceDependency(res: ArmResourceManipulator): Promise<RawScenario> {
    const putOperation = res.getOperation("CreateOrUpdate")[0];
    if (!putOperation) {
      return { description: undefined, steps: [] };
    }
    const scenario = this.generateDependencySteps(res);
    this.updateStepExample(scenario);
    scenario.variables = this.getVariables(scenario);
    return scenario;
  }

  public updateStepExample(scenario: RawScenario) {
    if (this.opts.useExample) {
      scenario.steps.forEach((step) => {
        const operationId = (step as any).operationId;
        const operation = this.operations.get(operationId);
        const exampleObj = Object.values(operation?.["x-ms-examples"] || {})?.[0];
        const exampleFile =
          (step as RawStepExample).exampleFile || exampleObj?.$ref
            ? this.fileLoader.resolvePath(this.jsonLoader.getRealPath(exampleObj.$ref!))
            : null;
        if (exampleFile) {
          (step as RawStepExample).exampleFile = exampleFile;
        }
      });
    }
  }

  public async generate() {
    const definition: RawScenarioDefinition = {
      scope: this.opts.fileRoot?.includes("data-plane") ? "None" : "ResourceGroup",
      authentication: undefined,
      variables: undefined,
      scenarios: [],
    };

    if (this.opts.scope) {
      let scope;
      if (urlParse(this.opts.scope)) {
        definition.scope = this.opts.scope;
        scope = this.opts.scope;
      } else {
        definition.scope = path.relative(this.opts.outputDir, this.opts.scope);
        scope = path.resolve(process.cwd(), this.opts.scope);
      }
      const [scopeDef] = await this.apiScenarioYamlLoader.load(scope);
      definition.authentication = scopeDef.scenarios[0].authentication;
    }

    definition.scenarios.push(this.generateSteps());
    definition.variables = this.getVariables(definition.scenarios[0]);

    if (this.opts.useExample) {
      definition.scenarios[0].steps.forEach((step) => {
        const operationId = (step as any).operationId;
        const operation = this.operations.get(operationId);
        const examples = operation?.[xmsExamples];
        if (examples) {
          const example = Object.values(examples)[0];
          (step as RawStepExample).exampleFile = path.relative(
            this.opts.outputDir,
            this.fileLoader.resolvePath(this.jsonLoader.getRealPath(example.$ref!))
          );
        } else {
          console.warn(`${operationId} has no example.`);
        }
      });
    }

    return definition;
  }

  public async writeFile(definition: RawScenarioDefinition) {
    const fileContent =
      "# yaml-language-server: $schema=https://raw.githubusercontent.com/Azure/azure-rest-api-specs/main/documentation/api-scenario/references/v1.2/schema.json\n" +
      dump(definition);
    const filePath = pathJoin(this.opts.outputDir, "basic.yaml");
    await this.fileLoader.writeFile(filePath, fileContent);
    logger.info(`${filePath} is generated.`);
  }

  private getVariables(scenario: RawScenario) {
    const variables: { [name: string]: string | Variable } = {};
    if (!this.parameters) {
      const map = new Map();
      for (const swagger of this.swaggers) {
        traverseSwagger(swagger, {
          onOperation: (operation) => {
            for (let parameter of operation.parameters ?? []) {
              parameter = this.jsonLoader.resolveRefObj(parameter);
              if (
                !parameter.required ||
                envVariables.includes(parameter.name) ||
                (parameter.in === "path" && parameter[xmsSkipUrlEncoding])
              ) {
                continue;
              }

              let key: any = parameter.name;

              if (parameter.in === "body") {
                key = this.jsonLoader.resolveRefObj(parameter.schema!);
              }

              const value = map.get(key);

              if (value) {
                value.count++;
                value.operationIds.add(operation.operationId);
              } else {
                map.set(key, {
                  operationIds: new Set<string>([operation.operationId!]),
                  name: parameter.name,
                  count: 1,
                  value: this.generateVariable(parameter),
                });
              }
            }
          },
        });
      }
      this.parameters = map;
    }

    [...this.parameters.values()]
      .sort((a, b) => b.count - a.count)
      .forEach((v) => {
        const operation = this.operations.get([...v.operationIds.values()][0]);
        if (this.opts.useExample) {
          let p = operation?.parameters?.find((p) => {
            p = this.jsonLoader.resolveRefObj(p);
            return p.name === v.name;
          });
          if (p) {
            p = this.jsonLoader.resolveRefObj(p);
          }
          // for body,query parameter
          if (
            p?.in !== "path" &&
            operation?.[xmsExamples] &&
            Object.keys(operation?.[xmsExamples] ?? {}).length
          ) {
            return;
          }
        }
        const step = scenario.steps.find(
          (s) =>
            (s as RawStepOperation).operationId &&
            v.operationIds.has((s as RawStepOperation).operationId)
        );
        if (!step) {
          return;
        }
        if (!variables[v.name] && v.count > 1) {
          variables[v.name] = v.value;
          return;
        }
        if (!step?.variables) {
          step.variables = {};
        }
        step.variables[v.name] = v.value;
      });

    return variables;
  }

  private generateVariable(parameter: Parameter): Variable | string | undefined {
    const genValue = (name: string, schema: Schema): VarValue => {
      if (util.isObject(schema)) {
        const ret: VarValue = {};
        const allOf = [...(schema.allOf ?? []), ...(schema.oneOf ?? [])];
        const s = {
          required: cloneDeep(schema.required) ?? [],
          properties: cloneDeep(schema.properties) ?? {},
        };
        while (allOf.length > 0) {
          const item = this.jsonLoader.resolveRefObj(allOf.shift()!);
          allOf.push(...(item.allOf ?? []), ...(item.oneOf ?? []));
          s.required = [...s.required, ...(item.required ?? [])];
          s.properties = { ...s.properties, ...(item.properties ?? {}) };
        }

        for (const name of s.required) {
          const prop = this.jsonLoader.resolveRefObj(s.properties[name]);
          ret[name] = genValue(name, prop);
        }

        return ret;
      }

      if (schema.default) {
        return schema.default;
      }

      if (schema.enum) {
        return schema.enum[0];
      }

      if (schema.type === "string" && envVariables.includes(name)) {
        return `$(${name})`;
      }

      if (schema.type === "array") {
        const prop = this.jsonLoader.resolveRefObj(schema.items!) as Schema;
        return this.mocker.mock(schema, name, genValue("", prop));
      }

      return this.mocker.mock(schema, name);
    };

    if (parameter.in === "body") {
      const schema = this.jsonLoader.resolveRefObj(parameter.schema!);
      const value: Variable = {
        type: "object",
        value: genValue(parameter.name, schema) as { [key: string]: VarValue },
      };
      return value;
    }

    if (parameter.in === "path" && parameter.type === "string") {
      // set prefix length to 8, thus 8+6<15, which is the minimum max length of resource name
      return { type: "string", prefix: `${parameter.name.toLocaleLowerCase().substring(0, 8)}` };
    }

    switch (parameter.type) {
      case "string":
        return this.mocker.mock(parameter, parameter.name);
      case "integer":
      case "number":
        return { type: "int", value: this.mocker.mock(parameter, parameter.name) };
      case "boolean":
        return { type: "bool", value: this.mocker.mock(parameter, parameter.name) };
      case "array":
        return { type: "array", value: this.mocker.mock(parameter, parameter.name) };
      default:
        logger.warn(`Unsupported type ${parameter.type} of parameter ${parameter.name}`);
        return undefined;
    }
  }
  private generateDependencySteps(res: ArmResourceManipulator) {
    const scenario: RawScenario = {
      description: undefined,
      variables: undefined,
      steps: [],
    };
    function getPutOperationId() {
      return res.getOperation("CreateOrUpdate")?.[0]?.operationId || "";
    }
    const sortedNodes: Node[] = [];
    const cmp = (a: Node, b: Node) => {
      const degree = b.inDegree - a.inDegree;
      if (degree) {
        return degree;
      }
      return 0;
    };
    const graph = this.graph;
    function widthFirst(node: Node) {
      const heap = new Heap<Node>(cmp);
      for (const n of graph.values()) {
        if (n.method === "put" && n.children.get(node.operationId)) {
          heap.push(n);
        }
        sortedNodes.push(...heap.toArray());
      }
      heap.toArray().forEach((n) => widthFirst(n));
    }
    const node = this.getNode(getPutOperationId());
    sortedNodes.push(node);
    widthFirst(node);
    const uniqNodes = sortedNodes.reverse().reduce(function (a: Node[], b: Node) {
      if (a.indexOf(b) < 0) a.push(b);
      return a;
    }, []);
    uniqNodes.forEach((node) => scenario.steps.push({ operationId: node.operationId }));
    return scenario;
  }

  public addCleanupSteps(res: ArmResourceManipulator, scenario: RawScenario) {
    const dependencyPutOperations = this.generateDependencySteps(res)?.steps.reverse();
    const sortedNodes: Node[] = [];
    for (const dependency of dependencyPutOperations) {
      const node = this.getNode((dependency as RawStepOperation).operationId);
      if (node) {
        const deleteOperation = [...node.children.values()].find((n) => n.method === "delete");
        if (deleteOperation) {
          sortedNodes.push(deleteOperation);
        }
      }
    }
    sortedNodes.forEach((node) => scenario.steps.push({ operationId: node.operationId }));
    return scenario;
  }

  private generateSteps() {
    const scenario: RawScenario = {
      scenario: "GeneratedScenario",
      steps: [],
    };

    const heap = new Heap<Node>((a, b) => {
      const priority = b.priority - a.priority;
      if (priority) {
        return priority;
      }

      const degree = b.outDegree - a.outDegree;
      if (degree) {
        return degree;
      }
      return methodOrder.indexOf(a.method) - methodOrder.indexOf(b.method);
    });

    for (const node of this.graph.values()) {
      if (node.operationId.includes("CheckNameAvailability")) {
        scenario.steps.push({ operationId: node.operationId });
        node.visited = true;
        continue;
      }
      if (node.inDegree === 0 && node.method === "put") {
        heap.push(node);
        node.visited = true;
      }
    }

    const deleteStack: Node[] = [];

    while (!heap.empty()) {
      const node = heap.pop()!;

      scenario.steps.push({ operationId: node.operationId });
      const operation = node.operationId.split("_")[0];

      for (const n of node.children.values()) {
        n.inDegree--;
        if (n.inDegree === 0 && n.method === "put") {
          heap.push(n);
          n.visited = true;
        }
      }

      if (node.method !== "put") {
        continue;
      }

      for (const n of this.graph.values()) {
        if (n.inDegree === 0 && !n.visited && n.operationId.split("_")[0] === operation) {
          n.priority = 1;
          if (n.method === "delete") {
            n.visited = true;
            deleteStack.push(n);
          } else {
            heap.push(n);
            n.visited = true;
          }
        }
      }
    }

    for (const node of this.graph.values()) {
      if (!node.visited) {
        scenario.steps.push({ operationId: node.operationId });
        if (node.inDegree !== 0) {
          console.error("node inDegree is not 0 ", node.operationId, node.method);
        }
      }
    }

    while (deleteStack.length > 0) {
      const node = deleteStack.pop()!;
      scenario.steps.push({ operationId: node.operationId });
    }

    return scenario;
  }

  private async generateGraph() {
    this.graph = new Map<string, Node>();
    const dependencies = (await this.jsonLoader.load(this.opts.dependencyPath)) as Dependencies;
    for (const path of Object.keys(dependencies)) {
      if (!path.startsWith("/")) {
        continue;
      }
      for (const method of Object.keys(dependencies[path])) {
        const operationId =
          this.getOperationId(path, method) || this.getOperationId(path + "/", method);
        if (!operationId) {
          console.warn(`can't find operationId, ${path} ${method}`);
          continue;
        }
        const node = this.getNode(operationId);
        node.method = method.toLowerCase() as LowerHttpMethods;

        for (const dependency of [
          ...(dependencies[path][method].Path ?? []),
          ...(dependencies[path][method].Query ?? []),
        ]) {
          if (dependency.producer_endpoint && dependency.producer_method) {
            const producerOperationId = this.getOperationId(
              dependency.producer_endpoint,
              dependency.producer_method
            );
            this.addDependency(operationId, producerOperationId);
          }
        }
      }
    }
  }

  private addDependency(operationId: string, producerOperationId: string) {
    const node = this.getNode(operationId);
    const dependNode = this.getNode(producerOperationId);

    node.inDegree++;
    dependNode.children.set(operationId, node);
    dependNode.outDegree++;
  }

  private getNode(operationId: string) {
    if (!this.graph.has(operationId)) {
      const node: Node = {
        operationId: operationId,
        children: new Map<string, Node>(),
        inDegree: 0,
        outDegree: 0,
        visited: false,
        method: "get",
        priority: 0,
      };
      this.graph.set(operationId, node);
    }

    return this.graph.get(operationId)!;
  }

  private getOperationId(path: string, method: string) {
    const m = method.toLowerCase() as LowerHttpMethods;
    for (const spec of this.swaggers) {
      const operationId = spec.paths[path]?.[m]?.operationId;
      if (operationId) {
        return operationId;
      }
    }
    return "";
  }
}