feat(probes) add finalize

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/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;
@@ -137,6 +139,9 @@ export class ProbeRunner {
 
       try {
         const result = this.#runProbe(probe, node);
+        if (probe.finalize) {
+          probe.finalize(this.sourceFile);
+        }
         if (result === null) {
           continue;
         }

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

@@ -145,4 +145,57 @@ 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 fakeProbeNeverValid = {
+        validateNode: (_: ESTree.Node) => [false],
+        main: () => ProbeSignals.Skip,
+        finalize: mock.fn()
+      };
+
+      const probes = [fakeProbe, fakeProbeBreak, fakeProbeSkip, fakeProbeNeverValid];
+
+      for (const probe of probes) {
+        const sourceFile = new SourceFile("");
+
+        const pr = new ProbeRunner(
+          sourceFile,
+          // @ts-expect-error
+          [probe]
+        );
+
+        const astNode: ESTree.Node = {
+          type: "Literal",
+          value: "test"
+        };
+        pr.walk(astNode);
+
+        assert.strictEqual(probe.finalize.mock.calls.length, 1);
+        assert.deepEqual(probe.finalize.mock.calls.at(0)?.arguments, [
+          sourceFile
+        ]);
+      }
+    });
+  });
 });