fit/src/trading.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { printAsset, accountBalance, waitForNextOpen, wait, setWaitFn, resetWaitFn } from './trading';
import type { Alpaca } from './alpaca';

function mockAlpaca(overrides: Partial<Alpaca> = {}): Alpaca {
    return {
        getAccount: vi.fn(),
        getAssets: vi.fn(),
        getAsset: vi.fn(),
        getClock: vi.fn(),
        getLatestAsk: vi.fn(),
        getLatestBid: vi.fn(),
        getLatestSpread: vi.fn(),
        getLatestTrades: vi.fn(),
        buy: vi.fn(),
        sell: vi.fn(),
        getOrder: vi.fn(),
        ...overrides,
    } as unknown as Alpaca;
}

beforeEach(() => {
    vi.useFakeTimers();
});

afterEach(() => {
    vi.useRealTimers();
    vi.restoreAllMocks();
    resetWaitFn();
});

describe('printAsset', () => {
    beforeEach(() => {
        vi.spyOn(console, 'debug').mockImplementation(() => {});
    });

    it('logs fractional when asset is fractionable', async () => {
        const alpaca = mockAlpaca({
            getAsset: vi.fn().mockResolvedValue({ symbol: 'TQQQ', fractionable: true, status: 'active', asset_class: 'us_equity' }),
        });

        await printAsset(alpaca, 'TQQQ');
        expect(console.debug).toHaveBeenCalledWith(expect.stringContaining('[DEBUG]'), 'TQQQ is fractional');
    });

    it('logs not fractional when asset is not fractionable', async () => {
        const alpaca = mockAlpaca({
            getAsset: vi.fn().mockResolvedValue({ symbol: 'BRK.A', fractionable: false, status: 'active', asset_class: 'us_equity' }),
        });

        await printAsset(alpaca, 'BRK.A');
        expect(console.debug).toHaveBeenCalledWith(expect.stringContaining('[DEBUG]'), 'BRK.A is not fractional');
    });
});

describe('accountBalance', () => {
    it('returns the cash value from the account', async () => {
        const alpaca = mockAlpaca({
            getAccount: vi.fn().mockResolvedValue({ cash: '10000.00' }),
        });

        const result = await accountBalance(alpaca);
        expect(result).toBe('10000.00');
    });
});

describe('waitForNextOpen', () => {
    beforeEach(() => {
        vi.spyOn(console, 'log').mockImplementation(() => {});
        vi.spyOn(console, 'debug').mockImplementation(() => {});
    });

    it('polls until market opens', async () => {
        const futureDate = new Date(Date.now() + 60000).toISOString();
        const alpaca = mockAlpaca({
            getClock: vi.fn()
                .mockResolvedValueOnce({ is_open: false, next_open: futureDate, next_close: futureDate })
                .mockResolvedValueOnce({ is_open: false, next_open: futureDate, next_close: futureDate })
                .mockResolvedValueOnce({ is_open: true, next_open: futureDate, next_close: futureDate }),
        });

        const promise = waitForNextOpen(alpaca, 100);
        await vi.advanceTimersByTimeAsync(300);
        await promise;

        expect(alpaca.getClock).toHaveBeenCalledTimes(3);
    });

    it('sleeps at most pollInterval even when open is far away', async () => {
        const farFuture = new Date(Date.now() + 86_400_000).toISOString();
        const alpaca = mockAlpaca({
            getClock: vi.fn()
                .mockResolvedValueOnce({ is_open: false, next_open: farFuture, next_close: farFuture })
                .mockResolvedValueOnce({ is_open: true, next_open: farFuture, next_close: farFuture }),
        });

        const promise = waitForNextOpen(alpaca, 200);
        await vi.advanceTimersByTimeAsync(300);
        await promise;

        // Should have polled twice: first sleep capped to 200ms, then saw is_open
        expect(alpaca.getClock).toHaveBeenCalledTimes(2);
    });
});

describe('setWaitFn / resetWaitFn', () => {
    it('uses custom wait function when set', async () => {
        const calls: number[] = [];
        setWaitFn(async (ms) => { calls.push(ms); });

        await wait(5000);
        await wait(3000);

        expect(calls).toEqual([5000, 3000]);
    });

    it('restores default wait after resetWaitFn', async () => {
        setWaitFn(async () => {});
        resetWaitFn();

        // After reset, wait should use real setTimeout again
        const promise = wait(1000);
        await vi.advanceTimersByTimeAsync(1000);
        await promise;
        // If we got here without hanging, default wait is restored
    });
});