Extract utility for rewriting

AikidoSec · Nov 28, 2024 · 91887b3 · 91887b3
1 parent a44c859
commit 91887b3
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 17 deletions.
diff --git a/library/agent/hooks/wrapRequire.ts b/library/agent/hooks/wrapRequire.ts
@@ -306,7 +306,10 @@ export function getOriginalRequire() {
   return originalRequire;
 }
 
-export function __internalRewritePackageName(
+// In order to support multiple versions of the same package, we need to rewrite the package name
+// e.g. In our sources and sinks, we use the real package name `hooks.addPackage("undici")`
+// but in the tests we want to `require("undici-v6")` instead of `require("undici")`
+export function rewritePackageName(
   packageName: string,
   aliasForTesting: string
 ) {

diff --git a/library/helpers/startTestAgent.ts b/library/helpers/startTestAgent.ts
@@ -0,0 +1,35 @@
+import type { ReportingAPI } from "../agent/api/ReportingAPI";
+import type { Token } from "../agent/api/Token";
+import { rewritePackageName } from "../agent/hooks/wrapRequire";
+import type { Logger } from "../agent/logger/Logger";
+import { Wrapper } from "../agent/Wrapper";
+import { createTestAgent } from "./createTestAgent";
+
+type PackageName = string;
+type AliasToRequire = string;
+
+/**
+ * Start a test agent for testing purposes
+ */
+export function startTestAgent(opts: {
+  block?: boolean;
+  logger?: Logger;
+  api?: ReportingAPI;
+  token?: Token;
+  serverless?: string;
+  wrappers: Wrapper[];
+  rewrite: Record<PackageName, AliasToRequire>;
+}) {
+  const agent = createTestAgent(opts);
+  agent.start(opts.wrappers);
+
+  // In order to support multiple versions of the same package, we need to rewrite the package name
+  // e.g. In our sources and sinks, we use the real package name `hooks.addPackage("undici")`
+  // but in the tests we want to `require("undici-v6")` instead of `require("undici")`
+  // The `__internalRewritePackageName` function allows us to do this
+  Object.keys(opts.rewrite).forEach((packageName) => {
+    rewritePackageName(packageName, opts.rewrite[packageName]);
+  });
+
+  return agent;
+}
diff --git a/library/sinks/Undici.custom.test.ts b/library/sinks/Undici.custom.test.ts
@@ -3,11 +3,10 @@ import * as dns from "dns";
 import * as t from "tap";
 import { Token } from "../agent/api/Token";
 import { Context, runWithContext } from "../agent/Context";
-import { __internalRewritePackageName } from "../agent/hooks/wrapRequire";
+import { startTestAgent } from "../helpers/startTestAgent";
 import { wrap } from "../helpers/wrap";
 import { getMajorNodeVersion } from "../helpers/getNodeVersion";
 import { Undici } from "./Undici";
-import { createTestAgent } from "../helpers/createTestAgent";
 
 function createContext(): Context {
   return {
@@ -43,11 +42,11 @@ t.test(
       getMajorNodeVersion() <= 16 ? "ReadableStream is not available" : false,
   },
   async (t) => {
-    const agent = createTestAgent({
+    startTestAgent({
       token: new Token("123"),
+      wrappers: [new Undici()],
+      rewrite: { undici: "undici-v6" },
     });
-    agent.start([new Undici()]);
-    __internalRewritePackageName("undici", "undici-v6");
 
     const { request, Dispatcher, setGlobalDispatcher, getGlobalDispatcher } =
       require("undici-v6") as typeof import("undici-v6");

diff --git a/library/sinks/Undici.localhost.test.ts b/library/sinks/Undici.localhost.test.ts
@@ -3,9 +3,8 @@ import * as t from "tap";
 import { createServer, Server } from "http";
 import { Token } from "../agent/api/Token";
 import { Context, runWithContext } from "../agent/Context";
-import { __internalRewritePackageName } from "../agent/hooks/wrapRequire";
-import { createTestAgent } from "../helpers/createTestAgent";
 import { getMajorNodeVersion } from "../helpers/getNodeVersion";
+import { startTestAgent } from "../helpers/startTestAgent";
 import { Undici } from "./Undici";
 
 function createContext({
@@ -51,13 +50,12 @@ function createContext({
   };
 }
 
-const agent = createTestAgent({
+startTestAgent({
   token: new Token("123"),
+  wrappers: [new Undici()],
+  rewrite: { undici: "undici-v6" },
 });
 
-agent.start([new Undici()]);
-__internalRewritePackageName("undici", "undici-v6");
-
 const port = 1346;
 const serverUrl = `http://localhost:${port}`;
 const hostHeader = `localhost:${port}`;

diff --git a/library/sinks/Undici.tests.ts b/library/sinks/Undici.tests.ts
@@ -5,11 +5,10 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Token } from "../agent/api/Token";
 import { Context, runWithContext } from "../agent/Context";
 import { LoggerForTesting } from "../agent/logger/LoggerForTesting";
+import { startTestAgent } from "../helpers/startTestAgent";
 import { wrap } from "../helpers/wrap";
 import { getMajorNodeVersion } from "../helpers/getNodeVersion";
 import { Undici } from "./Undici";
-import { createTestAgent } from "../helpers/createTestAgent";
-import { __internalRewritePackageName } from "../agent/hooks/wrapRequire";
 
 export function createUndiciTests(undiciPkgName: string, port: number) {
   const calls: Record<string, number> = {};
@@ -92,13 +91,15 @@ export function createUndiciTests(undiciPkgName: string, port: number) {
         block: true,
         receivedAnyStats: false,
       });
-      const agent = createTestAgent({
+      const agent = startTestAgent({
         api,
         logger,
         token: new Token("123"),
+        wrappers: [new Undici()],
+        rewrite: {
+          undici: undiciPkgName,
+        },
       });
-      agent.start([new Undici()]);
-      __internalRewritePackageName("undici", undiciPkgName);
 
       const {
         request,