feat: Allow specifying only credentials to pipe mode

rekmarks · Dec 20, 2024 · 1de66d2 · 1de66d2
1 parent 90f93b8
commit 1de66d2
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 73 deletions.
diff --git a/src/commands/migrate/pipe.test.ts b/src/commands/migrate/pipe.test.ts
@@ -3,24 +3,26 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 
 import { handlePipe } from './pipe.js';
 import { makeMockCredentials } from '../../../test/utils.js';
-import type { MockMigration } from '../../../test/utils.js';
-import { Migration } from '../../migration/index.js';
+import { Migration, operations } from '../../migration/index.js';
 import type { SerializedMigration } from '../../migration/types.js';
 
-vi.mock('../../migration/index.js', () => ({
-  Migration: {
-    deserialize: vi.fn(),
-  },
-}));
+vi.mock('../../migration/operations/index.js', async () => {
+  const { makeMockOperations } = await import('../../../test/utils.js');
+  return makeMockOperations({
+    initializeAgents: vi.fn().mockResolvedValue({
+      oldAgent: {},
+      newAgent: {},
+      accountDid: 'foo',
+    }),
+  });
+});
 
 type Stdin = typeof process.stdin;
 type Stdout = typeof process.stdout;
 
 describe('handlePipe', () => {
   let mockStdin: Readable;
   let mockStdout: Writable;
-  let mockMigration: MockMigration;
-  let mockSerializedMigration: SerializedMigration;
   let writtenOutput: string;
 
   beforeEach(() => {
@@ -38,99 +40,85 @@ describe('handlePipe', () => {
       },
     );
     vi.spyOn(process, 'stdout', 'get').mockReturnValue(mockStdout as Stdout);
-
-    mockMigration = {
-      run: vi.fn(),
-      serialize: vi.fn(),
-      deserialize: vi.fn(),
-      state: 'Ready',
-      confirmationToken: undefined,
-      newPrivateKey: undefined,
-      credentials: makeMockCredentials(),
-    };
-
-    mockSerializedMigration = {
-      state: 'Ready',
-      credentials: makeMockCredentials(),
-    };
-
-    vi.mocked(Migration.deserialize).mockResolvedValue(
-      mockMigration as unknown as Migration,
-    );
   });
 
   it('runs migration until RequestedPlcOperation when no confirmation token provided', async () => {
-    const expectedState = 'RequestedPlcOperation';
-    vi.mocked(mockMigration.run).mockResolvedValue(expectedState);
-
-    const serializedResult: SerializedMigration = {
-      ...mockSerializedMigration,
-      state: expectedState,
+    const inputData = {
+      state: 'Ready',
+      credentials: makeMockCredentials(),
     };
-    vi.mocked(mockMigration.serialize).mockReturnValue(serializedResult);
 
-    mockStdin.push(JSON.stringify(mockSerializedMigration));
+    mockStdin.push(JSON.stringify(inputData));
     mockStdin.push(null);
 
     await handlePipe();
 
-    expect(Migration.deserialize).toHaveBeenCalledOnce();
-    expect(Migration.deserialize).toHaveBeenCalledWith(mockSerializedMigration);
-
-    expect(mockMigration.run).toHaveBeenCalledOnce();
-
-    expect(writtenOutput).toBe(JSON.stringify(serializedResult));
+    expect(JSON.parse(writtenOutput)).toStrictEqual({
+      state: 'RequestedPlcOperation',
+      credentials: makeMockCredentials(),
+    });
   });
 
   it('runs migration to completion when confirmation token provided', async () => {
     const inputData = {
-      ...mockSerializedMigration,
+      state: 'Ready',
+      credentials: makeMockCredentials(),
       confirmationToken: 'test-token',
     };
+
+    mockStdin.push(JSON.stringify(inputData));
+    mockStdin.push(null);
+
     const expectedState = 'Finalized';
     const expectedPrivateKey = 'test-private-key';
 
-    mockMigration.confirmationToken = 'test-token';
-    mockMigration.newPrivateKey = expectedPrivateKey;
-    vi.mocked(mockMigration.run).mockResolvedValue(expectedState);
+    vi.mocked(operations.migrateIdentity).mockResolvedValue(expectedPrivateKey);
 
-    const serializedResult: SerializedMigration = {
+    await handlePipe();
+
+    expect(JSON.parse(writtenOutput)).toStrictEqual({
       ...inputData,
       state: expectedState,
       newPrivateKey: expectedPrivateKey,
+    });
+  });
+
+  it('handles partial migration', async () => {
+    const inputData = {
+      // No state
+      credentials: makeMockCredentials(),
     };
-    vi.mocked(mockMigration.serialize).mockReturnValue(serializedResult);
 
     mockStdin.push(JSON.stringify(inputData));
     mockStdin.push(null);
 
     await handlePipe();
 
-    expect(Migration.deserialize).toHaveBeenCalledOnce();
-    expect(Migration.deserialize).toHaveBeenCalledWith(inputData);
-
-    expect(mockMigration.run).toHaveBeenCalledOnce();
-
-    expect(writtenOutput).toBe(JSON.stringify(serializedResult));
+    expect(JSON.parse(writtenOutput)).toStrictEqual({
+      state: 'RequestedPlcOperation',
+      credentials: makeMockCredentials(),
+    });
   });
 
   it('outputs current state even when migration fails', async () => {
-    vi.mocked(mockMigration.run).mockRejectedValue(
-      new Error('Migration failed'),
-    );
-
-    const serializedResult: SerializedMigration = {
-      ...mockSerializedMigration,
+    const inputData: SerializedMigration = {
+      credentials: makeMockCredentials(),
       state: 'Ready',
     };
-    vi.mocked(mockMigration.serialize).mockReturnValue(serializedResult);
 
-    mockStdin.push(JSON.stringify(mockSerializedMigration));
+    mockStdin.push(JSON.stringify(inputData));
     mockStdin.push(null);
 
+    vi.mocked(operations.createNewAccount).mockRejectedValue(
+      new Error('Migration failed'),
+    );
+
     await expect(handlePipe()).rejects.toThrow('Migration failed');
 
-    expect(writtenOutput).toBe(JSON.stringify(serializedResult));
+    expect(JSON.parse(writtenOutput)).toStrictEqual({
+      ...inputData,
+      state: 'Initialized',
+    });
   });
 
   it('handles invalid JSON input', async () => {
@@ -150,16 +138,17 @@ describe('handlePipe', () => {
   });
 
   it('handles unexpected migration state', async () => {
-    const unexpectedState = 'Initialized';
-    vi.mocked(mockMigration.run).mockResolvedValue(unexpectedState);
-    vi.mocked(mockMigration.serialize).mockReturnValue({
-      ...mockSerializedMigration,
-      state: unexpectedState,
-    });
+    const inputData: SerializedMigration = {
+      credentials: makeMockCredentials(),
+      state: 'Ready',
+    };
 
-    mockStdin.push(JSON.stringify(mockSerializedMigration));
+    mockStdin.push(JSON.stringify(inputData));
     mockStdin.push(null);
 
+    const unexpectedState = 'Finalized';
+    vi.spyOn(Migration.prototype, 'run').mockResolvedValue(unexpectedState);
+
     await expect(handlePipe()).rejects.toThrow(
       `Fatal: Unexpected migration state "${unexpectedState}" after initial run`,
     );

diff --git a/src/commands/migrate/pipe.ts b/src/commands/migrate/pipe.ts
@@ -1,4 +1,7 @@
-import { Migration } from '../../migration/index.js';
+import {
+  isPartialSerializedMigration,
+  Migration,
+} from '../../migration/index.js';
 import type { MigrationState } from '../../migration/types.js';
 import { handleUnknownError, isPlainObject } from '../../utils/index.js';
 
@@ -21,7 +24,12 @@ export async function handlePipe(): Promise<void> {
 
   let migration: Migration;
   try {
-    migration = await Migration.deserialize(rawCredentials);
+    migration = isPartialSerializedMigration(rawCredentials)
+      ? await Migration.deserialize({
+          ...rawCredentials,
+          state: 'Ready',
+        })
+      : await Migration.deserialize(rawCredentials);
   } catch (error) {
     throw handleUnknownError('Invalid migration arguments', error);
   }

diff --git a/src/migration/types.ts b/src/migration/types.ts
@@ -83,6 +83,23 @@ export const SerializedMigrationSchema = union([
 
 export type SerializedMigration = TypeOf<typeof SerializedMigrationSchema>;
 
+const PartialMigrationSchema = object({
+  credentials: object({}).passthrough(),
+  confirmationToken: string().optional(),
+}).strict();
+
+/**
+ * A "partial" migration is a migration with only the credentials and, optionally,
+ * a confirmation token set. We will assume that such a migration is in the
+ * "Ready" state.
+ */
+export type PartialSerializedMigration = TypeOf<typeof PartialMigrationSchema>;
+
+export const isPartialSerializedMigration = (
+  value: unknown,
+): value is PartialSerializedMigration =>
+  PartialMigrationSchema.safeParse(value).success;
+
 export type MigrationCredentials = TypeOf<typeof MigrationCredentialsSchema>;
 
 export type AgentPair = {

diff --git a/test/utils.ts b/test/utils.ts
@@ -3,7 +3,7 @@ import type { HeadersMap } from '@atproto/xrpc';
 import type { Mock, Mocked } from 'vitest';
 import { vi } from 'vitest';
 
-import type { Migration } from '../src/migration/index.js';
+import type { Migration, operations } from '../src/migration/index.js';
 import type {
   MigrationCredentials,
   MigrationState,
@@ -31,6 +31,18 @@ export const makeMockCredentials = (): MigrationCredentials => ({
   inviteCode: 'invite-123',
 });
 
+export const makeMockOperations = (
+  mocks: Partial<typeof operations> = {},
+): typeof operations => ({
+  initializeAgents: vi.fn(),
+  createNewAccount: vi.fn(),
+  migrateData: vi.fn(),
+  requestPlcOperation: vi.fn(),
+  migrateIdentity: vi.fn(),
+  finalize: vi.fn(),
+  ...mocks,
+});
+
 export const mockAccountDid = 'did:plc:testuser123';
 
 export const makeXrpcResponse = <Data>(