From a78a84c5b5cdd4e5ab2e6a5a0505b5ec59f9b45c Mon Sep 17 00:00:00 2001
From: Snarling <84951833+Snarling@users.noreply.github.com>
Date: Wed, 5 Oct 2022 13:14:24 -0400
Subject: [PATCH] Make ram checks more robust

* Instead of hardcoded categories, automatically walk through all layers of ns, check for their associated costs, and check that ingame static and dynamic costs match the expected assigned costs.
---
 test/jest/Netscript/RamCalculation.test.ts | 244 +++++++--------------
 1 file changed, 81 insertions(+), 163 deletions(-)

diff --git a/test/jest/Netscript/RamCalculation.test.ts b/test/jest/Netscript/RamCalculation.test.ts
index 569abec7b..c29a10cde 100644
--- a/test/jest/Netscript/RamCalculation.test.ts
+++ b/test/jest/Netscript/RamCalculation.test.ts
@@ -7,8 +7,23 @@ import { Script } from "../../../src/Script/Script";
 import { WorkerScript } from "../../../src/Netscript/WorkerScript";
 import { calculateRamUsage } from "../../../src/Script/RamCalculations";
 
+type PotentiallyAsyncFunction = (arg?: unknown) => { catch?: PotentiallyAsyncFunction };
+type NSLayer = {
+  [key: string]: NSLayer | PotentiallyAsyncFunction;
+};
+type RamLayer = {
+  [key: string]: number | (() => number) | RamLayer;
+};
+function grabCost(ramLayer: RamLayer, fullPath: string[]) {
+  const ramEntry = ramLayer[fullPath[fullPath.length - 1]];
+  const expectedRam = typeof ramEntry === "function" ? ramEntry() : ramEntry;
+  if (typeof expectedRam !== "number") throw new Error(`There was no defined ram cost for ${fullPath.join(".")}().`);
+  return expectedRam;
+}
+
 describe("Netscript RAM Calculation/Generation Tests", function () {
   Player.sourceFiles[0] = { n: 4, lvl: 3 };
+  const sf4 = Player.sourceFiles[0];
   // For simulating costs of singularity functions.
   const ScriptBaseCost = RamCostConstants.ScriptBaseRamCost;
   const script = new Script();
@@ -20,15 +35,14 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
     return runningScript;
   }
 
-  type potentiallyAsyncFunction = (arg?: unknown) => { catch?: potentiallyAsyncFunction };
   /** Runs a Netscript function and properly catches an error even if it returns promise. */
-  function tryFunction(fn: potentiallyAsyncFunction) {
+  function tryFunction(fn: PotentiallyAsyncFunction) {
     try {
       fn()?.catch?.(() => undefined);
     } catch {}
   }
 
-  let runningScript = createRunningScript("");
+  let scriptRef = createRunningScript("");
   //Since it is expensive to create a workerscript and wrap the ns API, this is done once
   const workerScript = {
     args: [] as string[],
@@ -37,28 +51,40 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
     dynamicLoadedFns: {},
     dynamicRamUsage: RamCostConstants.ScriptBaseRamCost,
     env: new Environment(),
-    ramUsage: runningScript.ramUsage,
-    scriptRef: runningScript,
+    ramUsage: scriptRef.ramUsage,
+    scriptRef,
   };
   const ns = NetscriptFunctions(workerScript as WorkerScript);
 
-  function dynamicCheck(fnPath: string[], expectedRamCost: number, extraLayerCost = 0) {
+  function combinedRamCheck(
+    fn: PotentiallyAsyncFunction,
+    fnPath: string[],
+    expectedRamCost: number,
+    extraLayerCost = 0,
+  ) {
     const code = `${fnPath.join(".")}();\n`.repeat(3);
     const fnName = fnPath[fnPath.length - 1];
 
-    // update our existing WorkerScript
-    runningScript = createRunningScript(code);
-    workerScript.code = code;
-    workerScript.scriptRef = runningScript;
-    workerScript.ramUsage = workerScript.scriptRef.ramUsage;
-    workerScript.dynamicRamUsage = ScriptBaseCost;
-    workerScript.env = new Environment();
-    workerScript.dynamicLoadedFns = {};
+    //check imported getRamCost fn vs. expected ram from test
+    expect(getRamCost(...fnPath)).toEqual(expectedRamCost);
+
+    // Static ram check
+    const staticCost = calculateRamUsage(code, []).cost;
+    expect(staticCost).toBeCloseTo(ScriptBaseCost + expectedRamCost + extraLayerCost);
+
+    // reset workerScript for dynamic check
+    scriptRef = createRunningScript(code);
+    Object.assign(workerScript, {
+      code,
+      scriptRef,
+      ramUsage: scriptRef.ramUsage,
+      dynamicRamUsage: ScriptBaseCost,
+      env: new Environment(),
+      dynamicLoadedFns: {},
+    });
     workerScript.env.vars = ns;
 
     // Run the function through the workerscript's args
-    const fn = fnPath.reduce((prev, curr) => prev[curr], ns as any);
-
     if (typeof fn === "function") {
       tryFunction(fn);
       tryFunction(fn);
@@ -69,163 +95,55 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
 
     expect(workerScript.dynamicLoadedFns).toHaveProperty(fnName);
     expect(workerScript.dynamicRamUsage - ScriptBaseCost).toBeCloseTo(expectedRamCost, 5);
-    expect(workerScript.dynamicRamUsage).toBeCloseTo(runningScript.ramUsage - extraLayerCost, 5);
+    expect(workerScript.dynamicRamUsage).toBeCloseTo(scriptRef.ramUsage - extraLayerCost, 5);
   }
 
-  function testFunctionExpectZero(fnPath: string[], extraLayerCost = 0) {
-    const wholeFn = `${fnPath.join(".")}()`;
-    describe(wholeFn, () => {
-      it("Zero Ram Static Check", () => {
-        const ramCost = getRamCost(...fnPath);
-        expect(ramCost).toEqual(0);
-        const code = wholeFn;
-        const staticCost = calculateRamUsage(code, []).cost;
-        expect(staticCost).toEqual(ScriptBaseCost + extraLayerCost);
-      });
-      it("Zero Ram Dynamic check", () => dynamicCheck(fnPath, 0, extraLayerCost));
+  describe("ns", () => {
+    Object.entries(ns as unknown as NSLayer).forEach(([key, val]) => {
+      if (key === "args" || key === "enums") return;
+      if (typeof val === "function") {
+        const expectedRam = grabCost(RamCosts, [key]);
+        it(`${key}()`, () => combinedRamCheck(val, [key], expectedRam));
+      }
+      //The only other option should be an NSLayer
+      const extraLayerCost = { corporation: 1022.4, hacknet: 4 }[key] ?? 0;
+      testLayer(val as NSLayer, RamCosts[key as keyof typeof RamCosts] as RamLayer, [key], extraLayerCost);
     });
-  }
+  });
 
-  function testFunctionExpectNonzero(fnPath: string[], extraLayerCost = 0) {
-    const wholeFn = `${fnPath.join(".")}()`;
-    const ramCost = getRamCost(...fnPath);
-    describe(wholeFn, () => {
-      it("Static Check", () => {
-        expect(ramCost).toBeGreaterThan(0);
-        const code = wholeFn;
-        const staticCost = calculateRamUsage(code, []).cost;
-        expect(staticCost).toBeCloseTo(ramCost + ScriptBaseCost + extraLayerCost, 5);
-      });
-      it("Dynamic Check", () => dynamicCheck(fnPath, ramCost, extraLayerCost));
-    });
-  }
-
-  //input type for testSingularityFunctions
-  type singularityData = { fnPath: string[]; baseCost: number };
-
-  function testSingularityFunctions(data: singularityData[]) {
-    const sf4 = Player.sourceFiles[0];
-    data.forEach(({ fnPath, baseCost }) => {
-      const wholeFn = `${fnPath.join(".")}()`;
-      describe(wholeFn, () => {
-        it("SF4.3", () => {
-          sf4.lvl = 3;
-          const ramCost = getRamCost(...fnPath);
-          expect(ramCost).toEqual(baseCost);
-          const code = wholeFn;
-          const staticCost = calculateRamUsage(code, []).cost;
-          expect(staticCost).toBeCloseTo(ramCost + ScriptBaseCost);
-          dynamicCheck(fnPath, baseCost);
-        });
-        it("SF4.2", () => {
-          sf4.lvl = 2;
-          const ramCost = getRamCost(...fnPath);
-          expect(ramCost).toBeCloseTo(baseCost * 4, 5);
-          const code = wholeFn;
-          const staticCost = calculateRamUsage(code, []).cost;
-          expect(staticCost).toBeCloseTo(ramCost + ScriptBaseCost);
-          dynamicCheck(fnPath, ramCost);
-        });
-        it("SF4.1", () => {
-          sf4.lvl = 1;
-          const ramCost = getRamCost(...fnPath);
-          expect(ramCost).toBeCloseTo(baseCost * 16, 5);
-          const code = wholeFn;
-          const staticCost = calculateRamUsage(code, []).cost;
-          expect(staticCost).toBeCloseTo(ramCost + ScriptBaseCost);
-          dynamicCheck(fnPath, ramCost);
-        });
+  function testLayer(nsLayer: NSLayer, ramLayer: RamLayer, path: string[], extraLayerCost: number) {
+    describe(path[path.length - 1], () => {
+      Object.entries(nsLayer).forEach(([key, val]) => {
+        const newPath = [...path, key];
+        if (typeof val === "function") {
+          const fnName = newPath.join(".");
+          const expectedRam = grabCost(ramLayer, newPath);
+          it(`${fnName}()`, () => combinedRamCheck(val, newPath, expectedRam, extraLayerCost));
+        }
+        //A layer should be the only other option.
+        else testLayer(val, ramLayer[key] as RamLayer, newPath, 0);
       });
     });
   }
 
-  type NSLayer = {
-    [key: string]: NSLayer | unknown[] | (() => unknown);
-  };
-  type RamLayer = {
-    [key: string]: RamLayer | number | (() => number);
-  };
-  function testLayer(nsLayer: NSLayer, ramLayer: RamLayer, path: string[], extraLayerCost = 0) {
-    const zeroCostFunctions = Object.entries(nsLayer)
-      .filter(([key, val]) => ramLayer[key] === 0 && typeof val === "function")
-      .map(([key]) => [...path, key]);
-    zeroCostFunctions.forEach((fn) => testFunctionExpectZero(fn, extraLayerCost));
-
-    const nonzeroCostFunctions = Object.entries(nsLayer)
-      .filter(([key, val]) => ramLayer[key] > 0 && typeof val === "function")
-      .map(([key]) => [...path, key]);
-    nonzeroCostFunctions.forEach((fn) => testFunctionExpectNonzero(fn, extraLayerCost));
-  }
-
-  describe("Top level ns functions", () => {
-    const nsScope = ns as unknown as NSLayer;
-    const ramScope = RamCosts;
-    testLayer(nsScope, ramScope, []);
-  });
-
-  describe("Bladeburner API (bladeburner) functions", () => {
-    const nsScope = ns.bladeburner as unknown as NSLayer;
-    const ramScope = RamCosts.bladeburner;
-    testLayer(nsScope, ramScope, ["bladeburner"]);
-  });
-
-  describe("Corporation API (corporation) functions", () => {
-    const nsScope = ns.corporation as unknown as NSLayer;
-    const ramScope = RamCosts.corporation;
-    testLayer(nsScope, ramScope, ["corporation"], 1024 - ScriptBaseCost);
-  });
-
-  describe("TIX API (stock) functions", () => {
-    const nsScope = ns.stock as unknown as NSLayer;
-    const ramScope = RamCosts.stock;
-    testLayer(nsScope, ramScope, ["stock"]);
-  });
-
-  describe("Gang API (gang) functions", () => {
-    const nsScope = ns.gang as unknown as NSLayer;
-    const ramScope = RamCosts.gang;
-    testLayer(nsScope, ramScope, ["gang"]);
-  });
-
-  describe("Coding Contract API (codingcontract) functions", () => {
-    const nsScope = ns.codingcontract as unknown as NSLayer;
-    const ramScope = RamCosts.codingcontract;
-    testLayer(nsScope, ramScope, ["codingcontract"]);
-  });
-
-  describe("Sleeve API (sleeve) functions", () => {
-    const nsScope = ns.sleeve as unknown as NSLayer;
-    const ramScope = RamCosts.sleeve;
-    testLayer(nsScope, ramScope, ["sleeve"]);
-  });
-
-  //Singularity functions are tested in a different way because they also check SF4 level effect
-  describe("ns.singularity functions", () => {
-    const singularityFunctions = Object.entries(RamCosts.singularity).map(([fnName, ramFn]) => {
+  describe("Singularity multiplier checks", () => {
+    sf4.lvl = 3;
+    const singFunctions = Object.entries(ns.singularity).filter(([key, val]) => typeof val === "function");
+    const singObjects = singFunctions.map(([key, val]) => {
       return {
-        fnPath: ["singularity", fnName],
-        // This will error if a singularity function is assigned a flat cost instead of using the SF4 function
-        baseCost: (ramFn as () => number)(),
+        name: key,
+        fn: val,
+        baseRam: grabCost(RamCosts.singularity, ["singularity", key]),
       };
     });
-    testSingularityFunctions(singularityFunctions);
-  });
-
-  //Formulas requires deeper layer penetration
-  function formulasTest(
-    newLayer = "formulas",
-    oldNSLayer = ns as unknown as NSLayer,
-    oldRamLayer = RamCosts as unknown as RamLayer,
-    path = ["formulas"],
-    nsLayer = oldNSLayer[newLayer] as NSLayer,
-    ramLayer = oldRamLayer[newLayer] as RamLayer,
-  ) {
-    testLayer(nsLayer, ramLayer, path);
-    for (const [key, val] of Object.entries(nsLayer)) {
-      if (Array.isArray(val) || typeof val === "function" || key === "enums") continue;
-      // Only other option is an object / new layer
-      describe(key, () => formulasTest(key, nsLayer, ramLayer, [...path, key]));
+    const lvlToMult: Record<number, number> = { 0: 16, 1: 16, 2: 4 };
+    for (const lvl of [0, 1, 2]) {
+      it(`SF4.${lvl} check for x${lvlToMult[lvl]} costs`, () => {
+        sf4.lvl = lvl;
+        singObjects.forEach((obj) =>
+          combinedRamCheck(obj.fn, ["singularity", obj.name], obj.baseRam * lvlToMult[lvl], 0),
+        );
+      });
     }
-  }
-  describe("ns.formulas functions", formulasTest);
+  });
 });