hikari-desktop/src/lib/stores/costTracking.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";
import { get } from "svelte/store";
import { invoke } from "@tauri-apps/api/core";
import { setMockInvokeResult } from "../../../vitest.setup";
import {
  formatCost,
  formatAlertType,
  getAlertMessage,
  costTrackingStore,
  formattedCosts,
  type AlertType,
  type CostAlert,
} from "./costTracking";

vi.mock("$lib/notifications/notificationManager", () => ({
  notificationManager: {
    notifyCostAlert: vi.fn(),
  },
}));

describe("formatCost", () => {
  it("formats amounts below $0.01 to 4 decimal places", () => {
    expect(formatCost(0)).toBe("$0.0000");
    expect(formatCost(0.001)).toBe("$0.0010");
    expect(formatCost(0.0099)).toBe("$0.0099");
  });

  it("formats amounts between $0.01 and $1 to 3 decimal places", () => {
    expect(formatCost(0.01)).toBe("$0.010");
    expect(formatCost(0.123)).toBe("$0.123");
    expect(formatCost(0.999)).toBe("$0.999");
  });

  it("formats amounts $1 and above to 2 decimal places", () => {
    expect(formatCost(1)).toBe("$1.00");
    expect(formatCost(1.5)).toBe("$1.50");
    expect(formatCost(100.99)).toBe("$100.99");
  });
});

describe("formatAlertType", () => {
  it("formats Daily as Today", () => {
    expect(formatAlertType("Daily")).toBe("Today");
  });

  it("formats Weekly as This Week", () => {
    expect(formatAlertType("Weekly")).toBe("This Week");
  });

  it("formats Monthly as This Month", () => {
    expect(formatAlertType("Monthly")).toBe("This Month");
  });

  it("handles all AlertType values", () => {
    const types: AlertType[] = ["Daily", "Weekly", "Monthly"];
    const results = types.map(formatAlertType);
    expect(results).toEqual(["Today", "This Week", "This Month"]);
  });
});

describe("getAlertMessage", () => {
  it("generates a message for a Daily alert", () => {
    const alert: CostAlert = {
      alert_type: "Daily",
      threshold: 1.0,
      current_cost: 1.5,
    };
    const message = getAlertMessage(alert);
    expect(message).toContain("Today");
    expect(message).toContain("$1.50");
    expect(message).toContain("$1.00");
  });

  it("generates a message for a Weekly alert", () => {
    const alert: CostAlert = {
      alert_type: "Weekly",
      threshold: 5.0,
      current_cost: 6.0,
    };
    const message = getAlertMessage(alert);
    expect(message).toContain("This Week");
    expect(message).toContain("$6.00");
    expect(message).toContain("$5.00");
  });

  it("generates a message for a Monthly alert", () => {
    const alert: CostAlert = {
      alert_type: "Monthly",
      threshold: 20.0,
      current_cost: 25.0,
    };
    const message = getAlertMessage(alert);
    expect(message).toContain("This Month");
    expect(message).toContain("$25.00");
    expect(message).toContain("$20.00");
  });

  it("includes threshold and current cost in the message", () => {
    const alert: CostAlert = {
      alert_type: "Daily",
      threshold: 0.005,
      current_cost: 0.007,
    };
    const message = getAlertMessage(alert);
    expect(message).toContain("$0.0070");
    expect(message).toContain("$0.0050");
  });
});

describe("costTrackingStore", () => {
  beforeEach(() => {
    vi.clearAllMocks();
    costTrackingStore.reset();
  });

  describe("refresh", () => {
    it("loads cost data from the backend", async () => {
      setMockInvokeResult("get_today_cost", 1.5);
      setMockInvokeResult("get_week_cost", 5.25);
      setMockInvokeResult("get_month_cost", 20.0);
      setMockInvokeResult("get_cost_alerts", []);

      await costTrackingStore.refresh();

      const state = get(costTrackingStore);
      expect(state.todayCost).toBe(1.5);
      expect(state.weekCost).toBe(5.25);
      expect(state.monthCost).toBe(20.0);
      expect(state.isLoading).toBe(false);
    });

    it("triggers notifications for any returned alerts", async () => {
      const { notificationManager } = await import("$lib/notifications/notificationManager");
      const alert: CostAlert = { alert_type: "Daily", threshold: 1.0, current_cost: 1.5 };
      setMockInvokeResult("get_today_cost", 1.5);
      setMockInvokeResult("get_week_cost", 0);
      setMockInvokeResult("get_month_cost", 0);
      setMockInvokeResult("get_cost_alerts", [alert]);

      await costTrackingStore.refresh();

      expect(notificationManager.notifyCostAlert).toHaveBeenCalledWith(
        expect.stringContaining("Today")
      );
    });

    it("returns alerts from refresh", async () => {
      const alert: CostAlert = { alert_type: "Weekly", threshold: 5.0, current_cost: 6.0 };
      setMockInvokeResult("get_today_cost", 0);
      setMockInvokeResult("get_week_cost", 6.0);
      setMockInvokeResult("get_month_cost", 0);
      setMockInvokeResult("get_cost_alerts", [alert]);

      const result = await costTrackingStore.refresh();

      expect(result).toEqual([alert]);
    });

    it("handles refresh errors gracefully and returns empty array", async () => {
      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
      setMockInvokeResult("get_today_cost", new Error("Backend error"));

      const result = await costTrackingStore.refresh();

      expect(result).toEqual([]);
      expect(errorSpy).toHaveBeenCalledWith("Failed to refresh cost tracking:", expect.any(Error));
      errorSpy.mockRestore();
    });
  });

  describe("getSummary", () => {
    it("fetches and stores a cost summary", async () => {
      const mockSummary = {
        period_days: 7,
        total_input_tokens: 1000,
        total_output_tokens: 500,
        total_cost: 0.05,
        total_messages: 20,
        total_sessions: 3,
        average_daily_cost: 0.007,
        daily_breakdown: [],
      };
      setMockInvokeResult("get_cost_summary", mockSummary);

      const result = await costTrackingStore.getSummary(7);

      expect(result).toEqual(mockSummary);
      expect(invoke).toHaveBeenCalledWith("get_cost_summary", { days: 7 });
    });

    it("returns null and logs error on failure", async () => {
      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
      setMockInvokeResult("get_cost_summary", new Error("Summary unavailable"));

      const result = await costTrackingStore.getSummary(30);

      expect(result).toBeNull();
      expect(errorSpy).toHaveBeenCalledWith("Failed to get cost summary:", expect.any(Error));
      errorSpy.mockRestore();
    });
  });

  describe("setAlertThresholds", () => {
    it("persists new thresholds to the backend", async () => {
      const thresholds = { daily: 2.0, weekly: 10.0, monthly: 40.0 };

      await costTrackingStore.setAlertThresholds(thresholds);

      expect(invoke).toHaveBeenCalledWith("set_cost_alert_thresholds", {
        daily: 2.0,
        weekly: 10.0,
        monthly: 40.0,
      });
    });

    it("logs an error when the backend call fails", async () => {
      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
      setMockInvokeResult("set_cost_alert_thresholds", new Error("Failed"));

      await costTrackingStore.setAlertThresholds({ daily: 1.0, weekly: null, monthly: null });

      expect(errorSpy).toHaveBeenCalledWith("Failed to set alert thresholds:", expect.any(Error));
      errorSpy.mockRestore();
    });
  });

  describe("exportCsv", () => {
    it("returns the CSV string from the backend", async () => {
      setMockInvokeResult("export_cost_csv", "date,cost\n2026-01-01,1.50");

      const result = await costTrackingStore.exportCsv(7);

      expect(result).toBe("date,cost\n2026-01-01,1.50");
      expect(invoke).toHaveBeenCalledWith("export_cost_csv", { days: 7 });
    });

    it("returns null and logs error on failure", async () => {
      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
      setMockInvokeResult("export_cost_csv", new Error("Export failed"));

      const result = await costTrackingStore.exportCsv(30);

      expect(result).toBeNull();
      expect(errorSpy).toHaveBeenCalledWith("Failed to export CSV:", expect.any(Error));
      errorSpy.mockRestore();
    });
  });

  describe("reset", () => {
    it("resets costs back to zero", async () => {
      setMockInvokeResult("get_today_cost", 5.0);
      setMockInvokeResult("get_week_cost", 15.0);
      setMockInvokeResult("get_month_cost", 50.0);
      setMockInvokeResult("get_cost_alerts", []);
      await costTrackingStore.refresh();

      costTrackingStore.reset();

      const state = get(costTrackingStore);
      expect(state.todayCost).toBe(0);
      expect(state.weekCost).toBe(0);
      expect(state.monthCost).toBe(0);
    });
  });
});

describe("formattedCosts", () => {
  beforeEach(() => {
    costTrackingStore.reset();
  });

  it("formats the initial zero costs correctly", () => {
    const costs = get(formattedCosts);
    expect(costs.today).toBe("$0.0000");
    expect(costs.week).toBe("$0.0000");
    expect(costs.month).toBe("$0.0000");
  });

  it("reflects updated costs after a refresh", async () => {
    setMockInvokeResult("get_today_cost", 1.5);
    setMockInvokeResult("get_week_cost", 5.0);
    setMockInvokeResult("get_month_cost", 20.0);
    setMockInvokeResult("get_cost_alerts", []);
    await costTrackingStore.refresh();

    const costs = get(formattedCosts);
    expect(costs.today).toBe("$1.50");
    expect(costs.week).toBe("$5.00");
    expect(costs.month).toBe("$20.00");
    expect(costs.todayRaw).toBe(1.5);
  });
});