learning_ai_common_plat/packages/testing/src/cosmos-mocks.ts

/**
 * Shared Cosmos DB mock factories for Vitest.
 *
 * Usage:
 * ```ts
 * import { createCosmosMocks } from '@bytelyst/testing';
 * const { mockContainer, mockDatabase, MockCosmosClient, resetMocks } = createCosmosMocks();
 * vi.mock('@azure/cosmos', () => ({ CosmosClient: MockCosmosClient }));
 * ```
 */

import { vi } from 'vitest';

export interface MockItem {
  id: string;
  [key: string]: unknown;
}

export interface CosmosMocks {
  mockContainer: {
    id: string;
    items: {
      create: ReturnType<typeof vi.fn>;
      query: ReturnType<typeof vi.fn>;
      upsert: ReturnType<typeof vi.fn>;
    };
    item: ReturnType<typeof vi.fn>;
  };
  mockDatabase: {
    container: ReturnType<typeof vi.fn>;
    containers: { createIfNotExists: ReturnType<typeof vi.fn> };
  };
  mockDatabases: {
    createIfNotExists: ReturnType<typeof vi.fn>;
  };
  MockCosmosClient: ReturnType<typeof vi.fn>;
  resetMocks: () => void;
}

/**
 * Create a full set of Cosmos DB mocks for unit testing.
 *
 * Returns mock objects for container, database, and the CosmosClient constructor,
 * plus a `resetMocks()` function to clear all mock state between tests.
 */
export function createCosmosMocks(): CosmosMocks {
  const mockItem = {
    read: vi.fn().mockResolvedValue({ resource: undefined }),
    replace: vi.fn().mockResolvedValue({ resource: undefined }),
    delete: vi.fn().mockResolvedValue({}),
    patch: vi.fn().mockResolvedValue({ resource: undefined }),
  };

  const mockContainer = {
    id: 'test-container',
    items: {
      create: vi.fn().mockResolvedValue({ resource: {} }),
      query: vi.fn().mockReturnValue({
        fetchAll: vi.fn().mockResolvedValue({ resources: [] }),
      }),
      upsert: vi.fn().mockResolvedValue({ resource: {} }),
    },
    item: vi.fn().mockReturnValue(mockItem),
  };

  const mockDatabase = {
    container: vi.fn().mockReturnValue(mockContainer),
    containers: {
      createIfNotExists: vi.fn().mockResolvedValue({ container: mockContainer }),
    },
  };

  const mockDatabases = {
    createIfNotExists: vi.fn().mockResolvedValue({ database: mockDatabase }),
  };

  const MockCosmosClient = vi.fn().mockImplementation(() => ({
    database: vi.fn().mockReturnValue(mockDatabase),
    databases: mockDatabases,
  }));

  function resetMocks() {
    vi.clearAllMocks();
  }

  return { mockContainer, mockDatabase, mockDatabases, MockCosmosClient, resetMocks };
}