memind/scripts/backfill-memory-v2-pgvector.test.mjs

import assert from 'node:assert/strict';
import test from 'node:test';
import {
  parseMemoryPgvectorBackfillArgs,
  runMemoryPgvectorBackfillCli,
} from './backfill-memory-v2-pgvector.mjs';

function createMysqlPool(rows, closed) {
  return {
    async query() {
      return [rows];
    },
    async end() {
      closed.mysql = true;
    },
  };
}

function createPgPool(calls, closed) {
  return {
    async query(sql, params) {
      calls.push({ sql, params });
      return { rows: [] };
    },
    async end() {
      closed.pg = true;
    },
  };
}

const rows = [
  {
    id: 'mem-1',
    user_id: 'user-1',
    label: 'fact',
    memory_text: '用户关注 Memory V2',
    status: 'active',
    created_at: 1000,
    updated_at: 2000,
  },
];

test('parseMemoryPgvectorBackfillArgs defaults to safe dry-run settings', () => {
  assert.deepEqual(parseMemoryPgvectorBackfillArgs([]), {
    apply: false,
    limit: 100,
    cursor: { updatedAt: 0, id: '' },
    tableName: 'memory_embeddings',
    mysqlUrlEnv: 'MEMORY_BACKFILL_MYSQL_URL',
    pgUrlEnv: 'MEMORY_PGVECTOR_DATABASE_URL',
    embeddingModule: null,
    help: false,
  });
});

test('parseMemoryPgvectorBackfillArgs maps checkpoint and apply options', () => {
  assert.deepEqual(
    parseMemoryPgvectorBackfillArgs([
      '--apply',
      '--limit',
      '25',
      '--cursor-updated-at',
      '12345',
      '--cursor-id',
      'mem-9',
      '--table',
      'memory_embeddings_local',
      '--mysql-url-env',
      'LOCAL_MYSQL_URL',
      '--pg-url-env',
      'LOCAL_PG_URL',
      '--embedding-module',
      './embed-memory.mjs',
    ]),
    {
      apply: true,
      limit: 25,
      cursor: { updatedAt: 12345, id: 'mem-9' },
      tableName: 'memory_embeddings_local',
      mysqlUrlEnv: 'LOCAL_MYSQL_URL',
      pgUrlEnv: 'LOCAL_PG_URL',
      embeddingModule: './embed-memory.mjs',
      help: false,
    },
  );
});

test('runMemoryPgvectorBackfillCli dry-run scans MySQL and does not open PostgreSQL', async () => {
  const closed = { mysql: false, pg: false };
  let pgOpened = false;
  const lines = [];
  const originalLog = console.log;
  console.log = (line = '') => {
    lines.push(String(line));
  };
  try {
    const result = await runMemoryPgvectorBackfillCli(
      ['--limit', '10'],
      { MEMORY_BACKFILL_MYSQL_URL: 'mysql://local' },
      {
        createMysqlPool: async () => createMysqlPool(rows, closed),
        createPgPool: async () => {
          pgOpened = true;
          return createPgPool([], closed);
        },
      },
    );
    assert.equal(result.mode, 'dry-run');
    assert.equal(result.scanned, 1);
    assert.equal(result.inserted, 0);
    assert.equal(pgOpened, false);
    assert.equal(closed.mysql, true);
    assert.match(lines.join('\n'), /"mode": "dry-run"/);
  } finally {
    console.log = originalLog;
  }
});

test('runMemoryPgvectorBackfillCli apply requires pg url and embedding module', async () => {
  await assert.rejects(
    () => runMemoryPgvectorBackfillCli(
      ['--apply'],
      { MEMORY_BACKFILL_MYSQL_URL: 'mysql://local' },
      { createMysqlPool: async () => createMysqlPool(rows, {}) },
    ),
    /MEMORY_PGVECTOR_DATABASE_URL/,
  );
  await assert.rejects(
    () => runMemoryPgvectorBackfillCli(
      ['--apply'],
      {
        MEMORY_BACKFILL_MYSQL_URL: 'mysql://local',
        MEMORY_PGVECTOR_DATABASE_URL: 'postgresql://local',
      },
      { createMysqlPool: async () => createMysqlPool(rows, {}) },
    ),
    /--embedding-module/,
  );
});

test('runMemoryPgvectorBackfillCli apply embeds and closes both pools', async () => {
  const pgCalls = [];
  const closed = { mysql: false, pg: false };
  const originalLog = console.log;
  console.log = () => {};
  try {
    const result = await runMemoryPgvectorBackfillCli(
      ['--apply', '--embedding-module', './embed-memory.mjs'],
      {
        MEMORY_BACKFILL_MYSQL_URL: 'mysql://local',
        MEMORY_PGVECTOR_DATABASE_URL: 'postgresql://local',
      },
      {
        createMysqlPool: async () => createMysqlPool(rows, closed),
        createPgPool: async () => createPgPool(pgCalls, closed),
        loadEmbedMemory: async () => async () => [0.1, 0.2, 0.3],
      },
    );
    assert.equal(result.mode, 'apply');
    assert.equal(result.inserted, 1);
    assert.equal(pgCalls.length, 1);
    assert.equal(closed.mysql, true);
    assert.equal(closed.pg, true);
  } finally {
    console.log = originalLog;
  }
});