Mocking and Spying

/**
 * Create a mock function with optional implementation
 * @param implementation - Optional function implementation
 * @returns Mock function with tracking capabilities
 */
function fn<T extends Procedure = Procedure>(implementation?: T): Mock<T>;

// Legacy overloads (deprecated in v9.0)
function fn<TArgs extends any[] = any, R = any>(): Mock<(...args: TArgs) => R>;
function fn<TArgs extends any[] = any[], R = any>(
  implementation: (...args: TArgs) => R
): Mock<(...args: TArgs) => R>;

type Procedure = (...args: any[]) => any;

type Mock<T extends Procedure | any[] = any[], R = any> = T extends Procedure
  ? MockV2<T>
  : T extends any[]
    ? MockV2<(...args: T) => R>
    : never;

import { fn, expect } from '@storybook/test';

// Create a basic mock
const mockCallback = fn();
mockCallback('arg1', 'arg2');
expect(mockCallback).toHaveBeenCalledWith('arg1', 'arg2');

// Create a mock with implementation
const mockAdd = fn((a: number, b: number) => a + b);
const result = mockAdd(2, 3);
expect(result).toBe(5);
expect(mockAdd).toHaveBeenCalledWith(2, 3);

// Mock with custom name for debugging
const namedMock = fn().mockName('customName');
expect(namedMock.getMockName()).toBe('customName');

/**
 * Create a spy on an object method
 * @param object - Target object
 * @param method - Method name to spy on
 * @returns Mock instance that wraps the original method
 */
function spyOn<T, K extends keyof T>(
  object: T,
  method: K
): T[K] extends (...args: any[]) => any ? MockInstance<Parameters<T[K]>, ReturnType<T[K]>> : never;

import { spyOn, expect } from '@storybook/test';

class Calculator {
  add(a: number, b: number) {
    return a + b;
  }
}

const calc = new Calculator();

export const SpyStory = {
  play: async () => {
    // Spy on the method while preserving original behavior
    const addSpy = spyOn(calc, 'add');
    
    const result = calc.add(2, 3);
    expect(result).toBe(5); // Original behavior preserved
    expect(addSpy).toHaveBeenCalledWith(2, 3);
    
    // Override implementation
    addSpy.mockImplementation((a, b) => a * b);
    const multiplied = calc.add(2, 3);
    expect(multiplied).toBe(6); // Now it multiplies
    
    // Restore original behavior
    addSpy.mockRestore();
    const restored = calc.add(2, 3);
    expect(restored).toBe(5); // Back to addition
  },
};

interface MockInstance<TArgs extends any[] = any[], TReturns = any> {
  /**
   * Get the current mock name for debugging
   */
  getMockName(): string;
  
  /**
   * Set a name for the mock for debugging purposes
   * @param name - Name to assign to the mock
   */
  mockName(name: string): this;
  
  /**
   * Clear all call history but preserve implementation
   */
  mockClear(): this;
  
  /**
   * Reset mock to initial state (clear calls and implementation)
   */
  mockReset(): this;
  
  /**
   * Restore original implementation (for spies)
   */
  mockRestore(): void;
  
  /**
   * Set a custom implementation for the mock
   * @param fn - Implementation function
   */
  mockImplementation(fn?: (...args: TArgs) => TReturns): this;
  
  /**
   * Set implementation for the next single call only
   * @param fn - Implementation function for one call
   */
  mockImplementationOnce(fn: (...args: TArgs) => TReturns): this;
  
  /**
   * Set return value for all calls
   * @param value - Value to return
   */
  mockReturnValue(value: TReturns): this;
  
  /**
   * Set return value for the next single call only
   * @param value - Value to return for one call
   */
  mockReturnValueOnce(value: TReturns): this;
  
  /**
   * Set resolved value for async mocks
   * @param value - Value to resolve with
   */
  mockResolvedValue(value: Awaited<TReturns>): this;
  
  /**
   * Set resolved value for the next single call only
   * @param value - Value to resolve with for one call
   */
  mockResolvedValueOnce(value: Awaited<TReturns>): this;
  
  /**
   * Set rejected value for async mocks
   * @param value - Value to reject with
   */
  mockRejectedValue(value: any): this;
  
  /**
   * Set rejected value for the next single call only
   * @param value - Value to reject with for one call
   */
  mockRejectedValueOnce(value: any): this;
  
  // Call tracking properties
  mock: {
    calls: TArgs[];
    results: Array<{ type: 'return' | 'throw'; value: TReturns }>;
    instances: any[];
    contexts: any[];
    lastCall?: TArgs;
  };
}

import { fn, expect } from '@storybook/test';

export const MockInstanceStory = {
  play: async () => {
    const mock = fn<[number, number], number>();
    
    // Configure return values
    mock.mockReturnValue(42);
    expect(mock(1, 2)).toBe(42);
    
    mock.mockReturnValueOnce(100);
    expect(mock(3, 4)).toBe(100); // One-time return
    expect(mock(5, 6)).toBe(42);  // Back to default
    
    // Configure implementation
    mock.mockImplementation((a, b) => a + b);
    expect(mock(10, 20)).toBe(30);
    
    // One-time implementation
    mock.mockImplementationOnce((a, b) => a * b);
    expect(mock(3, 4)).toBe(12);  // Multiply once
    expect(mock(3, 4)).toBe(7);   // Back to addition
    
    // Async mocks
    const asyncMock = fn<[], Promise<string>>();
    asyncMock.mockResolvedValue('success');
    await expect(asyncMock()).resolves.toBe('success');
    
    asyncMock.mockRejectedValueOnce('error');
    await expect(asyncMock()).rejects.toBe('error');
    await expect(asyncMock()).resolves.toBe('success');
  },
};

/**
 * Clear call history for all active mocks
 * Calls .mockClear() on every mock
 */
function clearAllMocks(): void;

/**
 * Reset all active mocks to initial state
 * Calls .mockReset() on every mock
 */
function resetAllMocks(): void;

/**
 * Restore all spied methods to original implementations
 * Calls .mockRestore() on every mock
 */
function restoreAllMocks(): void;

/**
 * Check if a value is a mock function
 * @param value - Value to check
 * @returns True if value is a mock function
 */
function isMockFunction(value: unknown): boolean;

/**
 * Set of all active mock instances
 */
const mocks: Set<MockInstance>;

import { fn, spyOn, clearAllMocks, resetAllMocks, restoreAllMocks, isMockFunction } from '@storybook/test';

const originalConsole = console.log;

export const LifecycleStory = {
  play: async () => {
    // Create various mocks
    const mockFn = fn();
    const consoleSpy = spyOn(console, 'log');
    
    mockFn('test');
    console.log('test message');
    
    // Check if something is a mock
    expect(isMockFunction(mockFn)).toBe(true);
    expect(isMockFunction(console.log)).toBe(true);
    expect(isMockFunction(originalConsole)).toBe(false);
    
    // Clear all call history
    clearAllMocks();
    expect(mockFn).not.toHaveBeenCalled();
    expect(consoleSpy).not.toHaveBeenCalled();
    
    // Reset all mocks (clears history and implementations)
    mockFn.mockReturnValue('custom');
    resetAllMocks();
    expect(mockFn()).toBeUndefined(); // Implementation reset
    
    // Restore original implementations
    restoreAllMocks();
    console.log('this works normally again');
  },
};

/**
 * Register a listener for all mock function calls
 * @param callback - Function called when any mock is invoked
 * @returns Cleanup function to remove the listener
 */
function onMockCall(callback: Listener): () => void;

type Listener = (mock: MockInstance, args: unknown[]) => void;

import { fn, onMockCall } from '@storybook/test';

export const ListenerStory = {
  play: async () => {
    const calls: Array<{ mock: MockInstance; args: unknown[] }> = [];
    
    // Listen to all mock calls
    const unsubscribe = onMockCall((mock, args) => {
      calls.push({ mock, args });
      console.log(`Mock ${mock.getMockName()} called with:`, args);
    });
    
    const mockA = fn().mockName('mockA');
    const mockB = fn().mockName('mockB');
    
    mockA('arg1');
    mockB('arg2', 'arg3');
    
    expect(calls).toHaveLength(2);
    expect(calls[0].args).toEqual(['arg1']);
    expect(calls[1].args).toEqual(['arg2', 'arg3']);
    
    // Clean up listener
    unsubscribe();
    
    mockA('this will not be logged');
    expect(calls).toHaveLength(2); // No new calls logged
  },
};

/**
 * Type helper for objects that may have mocked methods
 * @param item - Object to type as potentially mocked
 * @param deep - Whether to deeply mock nested objects (default: false)
 * @param options - Configuration for mocking behavior
 * @returns The same object with mock typing
 */
function mocked<T>(item: T, deep?: false): MaybeMocked<T>;
function mocked<T>(item: T, deep: true): MaybeMockedDeep<T>;
function mocked<T>(item: T, options: { partial?: false; deep?: false }): MaybeMocked<T>;
function mocked<T>(item: T, options: { partial?: false; deep: true }): MaybeMockedDeep<T>;  
function mocked<T>(item: T, options: { partial: true; deep?: false }): MaybePartiallyMocked<T>;
function mocked<T>(item: T, options: { partial: true; deep: true }): MaybePartiallyMockedDeep<T>;

type MaybeMocked<T> = T & {
  [K in keyof T]: T[K] extends (...args: any[]) => any
    ? MockInstance<Parameters<T[K]>, ReturnType<T[K]>>
    : T[K];
};

type MaybeMockedDeep<T> = T & {
  [K in keyof T]: T[K] extends (...args: any[]) => any
    ? MockInstance<Parameters<T[K]>, ReturnType<T[K]>>
    : MaybeMockedDeep<T[K]>;
};

type MaybePartiallyMocked<T> = {
  [K in keyof T]?: T[K] extends (...args: any[]) => any
    ? MockInstance<Parameters<T[K]>, ReturnType<T[K]>>
    : T[K];
};

type MaybePartiallyMockedDeep<T> = {
  [K in keyof T]?: T[K] extends (...args: any[]) => any
    ? MockInstance<Parameters<T[K]>, ReturnType<T[K]>>
    : MaybePartiallyMockedDeep<T[K]>;
};

import { mocked, spyOn } from '@storybook/test';

interface ApiService {
  fetchUser(id: string): Promise<User>;
  updateUser(user: User): Promise<void>;
  nested: {
    helper(data: any): string;
  };
}

export const TypeHelpersStory = {
  play: async () => {
    const apiService: ApiService = {
      fetchUser: async (id) => ({ id, name: 'John' }),
      updateUser: async (user) => {},
      nested: {
        helper: (data) => JSON.stringify(data),
      },
    };
    
    // Spy on methods and use type helper
    spyOn(apiService, 'fetchUser');
    spyOn(apiService, 'updateUser');
    
    // Type helper provides correct typing for mocked methods
    const mockedService = mocked(apiService);
    mockedService.fetchUser.mockResolvedValue({ id: '1', name: 'Mock User' });
    
    const user = await apiService.fetchUser('1');
    expect(user.name).toBe('Mock User');
    
    // Deep mocking for nested objects
    const deepMocked = mocked(apiService, { deep: true });
    // Now deepMocked.nested.helper is also typed as a mock
  },
};