feat(probes) add finalize callback (#382)

* feat(probes) add finalize

* Update workspaces/js-x-ray/test/AstAnalyser.spec.ts

Co-authored-by: PierreDemailly <[email protected]>

---------

Co-authored-by: PierreDemailly <[email protected]>

Author: clemgbld

.changeset/lazy-weeks-roll.md

@@ -0,0 +1,5 @@
+---
+"@nodesecure/js-x-ray": minor
+---
+
+feat(probes): add finalize callback

workspaces/js-x-ray/src/AstAnalyser.ts

@@ -137,6 +137,7 @@ export class AstAnalyser {
       }
       finalize(source);
     }
+    source.probesRunner.finalize();
 
     // Add oneline-require flag if this is a one-line require expression
     if (isOneLineExpressionExport(body)) {

workspaces/js-x-ray/src/ProbeRunner.ts

@@ -25,6 +25,7 @@ import type { OptionalWarningName } from "./warnings.js";
 
 export type ProbeReturn = void | null | symbol;
 export type ProbeInitializeCallback = (sourceFile: SourceFile) => void;
+export type ProbeFinalizeCallback = (sourceFile: SourceFile) => void;
 export type ProbeMainCallback = (
   node: any,
   options: { sourceFile: SourceFile; data?: any; }
@@ -35,6 +36,7 @@ export type ProbeValidationCallback = (node: ESTree.Node, sourceFile: SourceFile
 export interface Probe {
   name: string;
   initialize?: ProbeInitializeCallback;
+  finalize?: ProbeFinalizeCallback;
   validateNode: ProbeValidationCallback | ProbeValidationCallback[];
   main: ProbeMainCallback;
   teardown?: ProbeTeardownCallback;
@@ -164,4 +166,12 @@ export class ProbeRunner {
 
     return null;
   }
+
+  finalize(): void {
+    for (const probe of this.probes) {
+      if (probe.finalize) {
+        probe.finalize(this.sourceFile);
+      }
+    }
+  }
 }

workspaces/js-x-ray/test/AstAnalyser.spec.ts

@@ -355,6 +355,33 @@ describe("AstAnalyser", () => {
       assert.equal(result.warnings.length, 1);
     });
 
+    it("should call initialize and finalize of every probes at the end", async(t) => {
+      const calls: string[] = [];
+      await new AstAnalyser(
+        {
+          customParser: new JsSourceParser(),
+          customProbes: [
+            {
+              name: "name",
+              initialize: () => calls.push("initialize"),
+              validateNode: () => [true],
+              main: t.mock.fn(),
+              finalize: () => calls.push("finalize")
+            },
+            {
+              name: "classic probe",
+              validateNode: () => [true],
+              main: t.mock.fn()
+            }
+
+          ],
+          skipDefaultProbes: true
+        }
+      ).analyseFile(new URL("customProbe.js", kFixtureURL));
+
+      assert.deepEqual(calls, ["initialize", "finalize"]);
+    });
+
     describe("hooks", () => {
       const analyser = new AstAnalyser();
       const url = new URL("depName.js", kFixtureURL);

workspaces/js-x-ray/test/ProbeRunner.spec.ts

@@ -145,4 +145,47 @@ describe("ProbeRunner", () => {
       assert.strictEqual(fakeProbe.teardown.mock.calls.length, 1);
     });
   });
+
+  describe("finalize", () => {
+    it("should call the finalize methods", () => {
+      const fakeProbe = {
+        validateNode: (_: ESTree.Node) => [true],
+        main: () => ProbeSignals.Skip,
+        finalize: mock.fn()
+      };
+
+      const fakeProbeSkip = {
+        validateNode: (_: ESTree.Node) => [true],
+        main: () => ProbeSignals.Skip,
+        teardown: mock.fn(),
+        finalize: mock.fn()
+      };
+
+      const fakeProbeBreak = {
+        validateNode: (_: ESTree.Node) => [true],
+        main: () => ProbeSignals.Break,
+        teardown: mock.fn(),
+        finalize: mock.fn()
+      };
+
+      const probes = [fakeProbe, fakeProbeBreak, fakeProbeSkip];
+
+      const sourceFile = new SourceFile("");
+
+      const pr = new ProbeRunner(
+        sourceFile,
+        // @ts-expect-error
+        probes
+      );
+
+      pr.finalize();
+
+      probes.forEach((probe) => {
+        assert.strictEqual(probe.finalize.mock.calls.length, 1);
+        assert.deepEqual(probe.finalize.mock.calls.at(0)?.arguments, [
+          sourceFile
+        ]);
+      });
+    });
+  });
 });