API: Minor followup changes for #1059 (#1135)

src/NetscriptFunctions.ts

@@ -108,6 +108,7 @@ import { getRamCost } from "./Netscript/RamCostGenerator";
 import { getEnumHelper } from "./utils/EnumHelper";
 import { setDeprecatedProperties, deprecationWarning } from "./utils/DeprecationHelper";
 import { ServerConstants } from "./Server/data/Constants";
+import { assertFunction } from "./Netscript/TypeAssertion";
 
 export const enums: NSEnums = {
   CityName,
@@ -1708,21 +1709,11 @@ export const ns: InternalAPI<NSFull> = {
     sinceInstall: Object.assign({}, Player.moneySourceA),
     sinceStart: Object.assign({}, Player.moneySourceB),
   }),
-  atExit:
-    (ctx) =>
-    (f, id = "default") => {
-      if (typeof f !== "function") {
-        throw helpers.errorMessage(ctx, "argument should be function");
-      }
-
-      if (typeof id !== "string") {
-        throw helpers.errorMessage(ctx, "id should be a string");
-      }
-
-      ctx.workerScript.atExit.set(id, () => {
-        f();
-      }); // Wrap the user function to prevent WorkerScript leaking as 'this'
-    },
+  atExit: (ctx) => (callback, _id) => {
+    const id = _id ? helpers.string(ctx, "id", _id) : "default";
+    assertFunction(ctx, "callback", callback);
+    ctx.workerScript.atExit.set(id, callback);
+  },
   mv: (ctx) => (_host, _source, _destination) => {
     const hostname = helpers.string(ctx, "host", _host);
     const server = helpers.getServer(ctx, hostname);