Files
omniai-server/scripts/keyManagerReleaseContract.test.js
T

74 lines
2.1 KiB
JavaScript

const assert = require("node:assert/strict");
const { createRequire } = require("node:module");
const nodeRequire = createRequire(__filename);
function loadKeyManagerWithPool(pool) {
const dbPath = nodeRequire.resolve("../src/db");
const keyManagerPath = nodeRequire.resolve("../src/keyManager");
const originalDbModule = nodeRequire.cache[dbPath];
const originalKeyManagerModule = nodeRequire.cache[keyManagerPath];
delete nodeRequire.cache[keyManagerPath];
nodeRequire.cache[dbPath] = {
id: dbPath,
filename: dbPath,
loaded: true,
exports: {
pool,
withTransaction: async (fn) => fn(pool),
},
};
return {
keyManager: nodeRequire("../src/keyManager"),
restore() {
delete nodeRequire.cache[keyManagerPath];
if (originalKeyManagerModule) nodeRequire.cache[keyManagerPath] = originalKeyManagerModule;
if (originalDbModule) nodeRequire.cache[dbPath] = originalDbModule;
else delete nodeRequire.cache[dbPath];
},
};
}
function createReleasePool() {
const calls = [];
return {
calls,
async query(sql, params) {
calls.push({ sql, params });
if (/WITH candidate AS/i.test(sql)) {
return {
rows: [{
id: 10,
key_id: 20,
lease_user_id: 30,
lease_enterprise_id: 40,
provider: "dashscope",
}],
};
}
if (/UPDATE api_keys SET active_count/i.test(sql)) return { rows: [] };
if (/INSERT INTO usage_logs/i.test(sql)) return { rows: [] };
throw new Error(`Unexpected SQL: ${sql}`);
},
};
}
(async () => {
const pool = createReleasePool();
const { keyManager, restore } = loadKeyManagerWithPool(pool);
try {
const result = await keyManager.releaseKey("lease-token-without-user-context");
assert.equal(result.released, true);
const usageLogCall = pool.calls.find((call) => /INSERT INTO usage_logs/i.test(call.sql));
assert.deepEqual(usageLogCall.params, [30, 40, 20, 20, "release"]);
} finally {
restore();
}
})().catch((error) => {
console.error(error);
process.exitCode = 1;
});