harb/kraiken-lib/src/tests/snatch.test.ts

import { describe, expect, test } from 'vitest';
import { getSnatchList, minimumTaxRate, selectSnatchPositions, type SnatchablePosition } from '../snatch.js';
import { uint256ToBytesLittleEndian } from '../subgraph.js';
import type { Positions as GraphPosition } from '../__generated__/graphql.js';

describe('minimumTaxRate', () => {
  test('finds the lowest tax from a list', () => {
    const rates = [{ taxRate: 0.12 }, { taxRate: 0.05 }, { taxRate: 0.08 }];
    expect(minimumTaxRate(rates, 1)).toBeCloseTo(0.05);
  });

  test('returns fallback for empty array', () => {
    expect(minimumTaxRate([], 0.42)).toBe(0.42);
  });

  test('returns fallback of 0 by default for empty array', () => {
    expect(minimumTaxRate([])).toBe(0);
  });

  test('works with a single element', () => {
    expect(minimumTaxRate([{ taxRate: 0.07 }])).toBeCloseTo(0.07);
  });
});

describe('selectSnatchPositions', () => {
  test('chooses cheapest positions first', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 30n, taxRate: 0.05, taxRateIndex: 5 },
      { id: 2n, stakeShares: 40n, taxRate: 0.03, taxRateIndex: 3 },
      { id: 3n, stakeShares: 50n, taxRate: 0.07, taxRateIndex: 7 },
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 60n,
      maxTaxRate: 0.08,
    });

    expect(result.selected.map(item => item.id)).toEqual([2n, 1n]);
    expect(result.remainingShortfall).toBe(0n);
    expect(result.coveredShares).toBe(60n);
    expect(result.maxSelectedTaxRateIndex).toBe(5);
  });

  test('keeps track of remaining shortfall when not enough candidates', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 10n, taxRate: 0.01 },
      { id: 2n, stakeShares: 10n, taxRate: 0.02 },
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 50n,
      maxTaxRate: 0.05,
    });

    expect(result.remainingShortfall).toBe(30n);
  });

  test('returns empty result when shortfallShares is zero', () => {
    const candidates: SnatchablePosition[] = [{ id: 1n, stakeShares: 100n, taxRate: 0.05 }];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 0n,
      maxTaxRate: 0.1,
    });

    expect(result.selected).toEqual([]);
    expect(result.coveredShares).toBe(0n);
    expect(result.remainingShortfall).toBe(0n);
  });

  test('filters out candidates at or above maxTaxRate', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 50n, taxRate: 0.1 }, // equal to maxTaxRate -> filtered
      { id: 2n, stakeShares: 50n, taxRate: 0.15 }, // above maxTaxRate -> filtered
      { id: 3n, stakeShares: 50n, taxRate: 0.05 }, // below maxTaxRate -> included
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 100n,
      maxTaxRate: 0.1,
    });

    expect(result.selected.map(p => p.id)).toEqual([3n]);
    expect(result.remainingShortfall).toBe(50n);
  });

  test('filters out positions owned by recipient when includeOwned is false', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 50n, taxRate: 0.05, owner: '0xABC' },
      { id: 2n, stakeShares: 50n, taxRate: 0.04, owner: '0xDEF' },
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 100n,
      maxTaxRate: 0.1,
      includeOwned: false,
      recipientAddress: '0xabc', // same as owner of id=1 (case-insensitive)
    });

    // id=1 is owned by recipient -> filtered out
    expect(result.selected.map(p => p.id)).toEqual([2n]);
  });

  test('includes owned positions when includeOwned is true', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 50n, taxRate: 0.05, owner: '0xABC' },
      { id: 2n, stakeShares: 50n, taxRate: 0.04, owner: '0xDEF' },
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 100n,
      maxTaxRate: 0.1,
      includeOwned: true,
      recipientAddress: '0xabc',
    });

    expect(result.selected.map(p => p.id)).toEqual([2n, 1n]);
    expect(result.coveredShares).toBe(100n);
  });

  test('skips candidates with zero stakeShares', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 0n, taxRate: 0.01 }, // zero shares -> skipped
      { id: 2n, stakeShares: 50n, taxRate: 0.02 },
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 30n,
      maxTaxRate: 0.05,
    });

    expect(result.selected.map(p => p.id)).toEqual([2n]);
  });

  test('tracks maxSelectedTaxRateIndex for tie-breaking same tax rate', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 20n, taxRate: 0.05, taxRateIndex: 3 },
      { id: 2n, stakeShares: 20n, taxRate: 0.05, taxRateIndex: 7 }, // higher index, same rate
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 40n,
      maxTaxRate: 0.1,
    });

    // Both are selected; the one with higher taxRateIndex should win
    expect(result.maxSelectedTaxRateIndex).toBe(7);
  });

  test('handles tie-breaking when first candidate has no taxRateIndex (covers ?? -1 fallback)', () => {
    const candidates: SnatchablePosition[] = [
      { id: 1n, stakeShares: 20n, taxRate: 0.05 }, // no taxRateIndex
      { id: 2n, stakeShares: 20n, taxRate: 0.05, taxRateIndex: 5 }, // same rate, has index
    ];

    const result = selectSnatchPositions(candidates, {
      shortfallShares: 40n,
      maxTaxRate: 0.1,
    });

    expect(result.maxSelectedTaxRateIndex).toBe(5);
  });
});

describe('getSnatchList', () => {
  const makePosition = (id: bigint, share: number, taxRate: number, owner = '0xowner'): GraphPosition => ({
    __typename: 'positions' as const,
    id: uint256ToBytesLittleEndian(id) as unknown as string,
    owner,
    share,
    creationTime: '0',
    lastTaxTime: '0',
    taxRate,
    status: 'Active',
    createdAt: 0,
    kraikenDeposit: '0',
    payout: '0',
    snatched: 0,
    stakeDeposit: '0',
    taxPaid: '0',
    totalSupplyInit: '0',
  }) as unknown as GraphPosition;

  test('converts subgraph positions and returns lowest-tax IDs', () => {
    const stakeTotalSupply = 1_000_000n * 10n ** 18n;
    const positions = [makePosition(1n, 0.0001, 0.02), makePosition(2n, 0.0002, 0.01)];

    const result = getSnatchList(positions, 10n ** 18n, 0.05, stakeTotalSupply);

    expect(result).toEqual([2n]);
  });

  test('throws when stakeTotalSupply is zero', () => {
    const positions = [makePosition(1n, 0.5, 0.02)];
    expect(() => getSnatchList(positions, 10n, 0.05, 0n)).toThrow('stakeTotalSupply must be greater than zero');
  });

  test('throws when there is not enough capacity to cover needed shares', () => {
    // Only one small position, but need is larger than its shares
    const stakeTotalSupply = 100n;
    const positions = [makePosition(1n, 0.001, 0.01)]; // tiny share

    expect(() => getSnatchList(positions, 1000n, 0.05, stakeTotalSupply)).toThrow('Not enough capacity');
  });
});