docfast/src/__tests__/billing.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";
import express from "express";
import request from "supertest";

// We need to mock Stripe before importing billing router
vi.mock("stripe", () => {
  const mockStripe = {
    checkout: {
      sessions: {
        create: vi.fn(),
        retrieve: vi.fn(),
      },
    },
    webhooks: {
      constructEvent: vi.fn(),
    },
    products: {
      search: vi.fn(),
      create: vi.fn(),
    },
    prices: {
      list: vi.fn(),
      create: vi.fn(),
    },
    subscriptions: {
      retrieve: vi.fn(),
    },
  };
  return { default: vi.fn(function() { return mockStripe; }), __mockStripe: mockStripe };
});

let app: express.Express;
let mockStripe: any;

beforeEach(async () => {
  vi.clearAllMocks();
  vi.resetModules();

  process.env.STRIPE_SECRET_KEY = "sk_test_fake";
  process.env.STRIPE_WEBHOOK_SECRET = "whsec_test_fake";

  // Re-import to get fresh mocks
  const stripeMod = await import("stripe");
  mockStripe = (stripeMod as any).__mockStripe;

  // Default: product search returns existing product+price
  mockStripe.products.search.mockResolvedValue({ data: [{ id: "prod_TygeG8tQPtEAdE" }] });
  mockStripe.prices.list.mockResolvedValue({ data: [{ id: "price_123" }] });

  const { createProKey, findKeyByCustomerId, downgradeByCustomer, updateEmailByCustomer } = await import("../services/keys.js");
  vi.mocked(createProKey).mockResolvedValue({ key: "pro-key-123", tier: "pro", email: "test@test.com", createdAt: new Date().toISOString() } as any);
  vi.mocked(findKeyByCustomerId).mockResolvedValue(null);
  vi.mocked(downgradeByCustomer).mockResolvedValue(undefined as any);
  vi.mocked(updateEmailByCustomer).mockResolvedValue(true as any);

  const { billingRouter } = await import("../routes/billing.js");
  app = express();
  // Webhook needs raw body
  app.use("/v1/billing/webhook", express.raw({ type: "application/json" }));
  app.use(express.json());
  app.use("/v1/billing", billingRouter);
});

describe("POST /v1/billing/checkout", () => {
  it("returns url on success", async () => {
    mockStripe.checkout.sessions.create.mockResolvedValue({ id: "cs_123", url: "https://checkout.stripe.com/pay/cs_123" });
    const res = await request(app).post("/v1/billing/checkout").send({});
    expect(res.status).toBe(200);
    expect(res.body.url).toBe("https://checkout.stripe.com/pay/cs_123");
  });

  it("returns 413 for body too large", async () => {
    // The route checks content-length header; send a large body to trigger it
    const largeBody = JSON.stringify({ padding: "x".repeat(2000) });
    const res = await request(app)
      .post("/v1/billing/checkout")
      .set("content-type", "application/json")
      .send(largeBody);
    expect(res.status).toBe(413);
  });

  it("returns 500 on Stripe error", async () => {
    mockStripe.checkout.sessions.create.mockRejectedValue(new Error("Stripe down"));
    const res = await request(app).post("/v1/billing/checkout").send({});
    expect(res.status).toBe(500);
    expect(res.body.error).toMatch(/Failed to create checkout session/);
  });
});

describe("GET /v1/billing/success", () => {
  it("returns 400 for missing session_id", async () => {
    const res = await request(app).get("/v1/billing/success");
    expect(res.status).toBe(400);
  });

  it("returns 409 for duplicate session", async () => {
    // First call succeeds
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_dup",
      customer: "cus_123",
      customer_details: { email: "test@test.com" },
    });
    await request(app).get("/v1/billing/success?session_id=cs_dup");
    // Second call with same session
    const res = await request(app).get("/v1/billing/success?session_id=cs_dup");
    expect(res.status).toBe(409);
  });

  it("returns existing key page when key already in DB", async () => {
    const { findKeyByCustomerId } = await import("../services/keys.js");
    vi.mocked(findKeyByCustomerId).mockResolvedValue({ key: "existing-key", tier: "pro" } as any);
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_existing",
      customer: "cus_existing",
      customer_details: { email: "test@test.com" },
    });
    const res = await request(app).get("/v1/billing/success?session_id=cs_existing");
    expect(res.status).toBe(200);
    expect(res.text).toContain("Key Already Provisioned");
  });

  it("provisions new key on success", async () => {
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_new",
      customer: "cus_new",
      customer_details: { email: "new@test.com" },
    });
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app).get("/v1/billing/success?session_id=cs_new");
    expect(res.status).toBe(200);
    expect(res.text).toContain("Welcome to Pro");
    expect(createProKey).toHaveBeenCalledWith("new@test.com", "cus_new");
  });

  it("returns 500 on Stripe error", async () => {
    mockStripe.checkout.sessions.retrieve.mockRejectedValue(new Error("Stripe error"));
    const res = await request(app).get("/v1/billing/success?session_id=cs_err");
    expect(res.status).toBe(500);
  });

  it("returns 400 when session has no customer", async () => {
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_no_cust",
      customer: null,
      customer_details: { email: "test@test.com" },
    });
    const res = await request(app).get("/v1/billing/success?session_id=cs_no_cust");
    expect(res.status).toBe(400);
    expect(res.body.error).toMatch(/No customer found/);
  });

  it("escapes HTML in displayed key to prevent XSS", async () => {
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_xss",
      customer: "cus_xss",
      customer_details: { email: "xss@test.com" },
    });
    const { createProKey } = await import("../services/keys.js");
    vi.mocked(createProKey).mockResolvedValue({
      key: '<script>alert("xss")</script>',
      tier: "pro",
      email: "xss@test.com",
      createdAt: new Date().toISOString(),
    } as any);
    const res = await request(app).get("/v1/billing/success?session_id=cs_xss");
    expect(res.status).toBe(200);
    expect(res.text).not.toContain('<script>alert("xss")</script>');
    expect(res.text).toContain("&lt;script&gt;");
  });
});

describe("POST /v1/billing/webhook", () => {
  it("returns 500 when webhook secret missing", async () => {
    delete process.env.STRIPE_WEBHOOK_SECRET;
    // Need to re-import to pick up env change - but the router is already loaded
    // The router reads env at request time, so this should work
    const savedSecret = process.env.STRIPE_WEBHOOK_SECRET;
    process.env.STRIPE_WEBHOOK_SECRET = "";
    delete process.env.STRIPE_WEBHOOK_SECRET;

    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "sig_test")
      .send(JSON.stringify({ type: "test" }));
    expect(res.status).toBe(500);

    process.env.STRIPE_WEBHOOK_SECRET = "whsec_test_fake";
  });

  it("returns 400 for missing signature", async () => {
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .send(JSON.stringify({ type: "test" }));
    expect(res.status).toBe(400);
    expect(res.body.error).toMatch(/Missing stripe-signature/);
  });

  it("returns 400 for invalid signature", async () => {
    mockStripe.webhooks.constructEvent.mockImplementation(() => {
      throw new Error("Invalid signature");
    });
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "bad_sig")
      .send(JSON.stringify({ type: "test" }));
    expect(res.status).toBe(400);
    expect(res.body.error).toMatch(/Invalid signature/);
  });

  it("provisions key on checkout.session.completed for DocFast product", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "checkout.session.completed",
      data: {
        object: {
          id: "cs_wh",
          customer: "cus_wh",
          customer_details: { email: "wh@test.com" },
        },
      },
    });
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_wh",
      line_items: {
        data: [{ price: { product: "prod_TygeG8tQPtEAdE" } }],
      },
    });
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "checkout.session.completed" }));
    expect(res.status).toBe(200);
    expect(res.body.received).toBe(true);
    expect(createProKey).toHaveBeenCalledWith("wh@test.com", "cus_wh");
  });

  it("ignores checkout.session.completed for non-DocFast product", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "checkout.session.completed",
      data: {
        object: {
          id: "cs_other",
          customer: "cus_other",
          customer_details: { email: "other@test.com" },
        },
      },
    });
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_other",
      line_items: {
        data: [{ price: { product: "prod_OTHER" } }],
      },
    });
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "checkout.session.completed" }));
    expect(res.status).toBe(200);
    expect(createProKey).not.toHaveBeenCalled();
  });

  it("downgrades on customer.subscription.deleted", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.deleted",
      data: {
        object: { id: "sub_del", customer: "cus_del" },
      },
    });
    mockStripe.subscriptions.retrieve.mockResolvedValue({
      items: { data: [{ price: { product: { id: "prod_TygeG8tQPtEAdE" } } }] },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.deleted" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).toHaveBeenCalledWith("cus_del");
  });

  it("downgrades on customer.subscription.updated with cancel_at_period_end", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.updated",
      data: {
        object: { id: "sub_cancel", customer: "cus_cancel", status: "active", cancel_at_period_end: true },
      },
    });
    mockStripe.subscriptions.retrieve.mockResolvedValue({
      items: { data: [{ price: { product: { id: "prod_TygeG8tQPtEAdE" } } }] },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.updated" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).toHaveBeenCalledWith("cus_cancel");
  });

  it("does not provision key when checkout.session.completed has missing customer", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "checkout.session.completed",
      data: {
        object: {
          id: "cs_no_cust",
          customer: null,
          customer_details: { email: "nocust@test.com" },
        },
      },
    });
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_no_cust",
      line_items: { data: [{ price: { product: "prod_TygeG8tQPtEAdE" } }] },
    });
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "checkout.session.completed" }));
    expect(res.status).toBe(200);
    expect(createProKey).not.toHaveBeenCalled();
  });

  it("does not provision key when checkout.session.completed has missing email", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "checkout.session.completed",
      data: {
        object: {
          id: "cs_no_email",
          customer: "cus_no_email",
          customer_details: {},
        },
      },
    });
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_no_email",
      line_items: { data: [{ price: { product: "prod_TygeG8tQPtEAdE" } }] },
    });
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "checkout.session.completed" }));
    expect(res.status).toBe(200);
    expect(createProKey).not.toHaveBeenCalled();
  });

  it("does not downgrade on customer.subscription.updated with non-DocFast product", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.updated",
      data: {
        object: { id: "sub_other", customer: "cus_other", status: "canceled", cancel_at_period_end: false },
      },
    });
    mockStripe.subscriptions.retrieve.mockResolvedValue({
      items: { data: [{ price: { product: { id: "prod_OTHER" } } }] },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.updated" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).not.toHaveBeenCalled();
  });

  it("downgrades on customer.subscription.updated with past_due status", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.updated",
      data: {
        object: { id: "sub_past", customer: "cus_past", status: "past_due", cancel_at_period_end: false },
      },
    });
    mockStripe.subscriptions.retrieve.mockResolvedValue({
      items: { data: [{ price: { product: { id: "prod_TygeG8tQPtEAdE" } } }] },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.updated" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).toHaveBeenCalledWith("cus_past");
  });

  it("does not downgrade on customer.subscription.updated with active status and no cancel", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.updated",
      data: {
        object: { id: "sub_ok", customer: "cus_ok", status: "active", cancel_at_period_end: false },
      },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.updated" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).not.toHaveBeenCalled();
  });

  it("does not downgrade on customer.subscription.deleted with non-DocFast product", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.subscription.deleted",
      data: {
        object: { id: "sub_del_other", customer: "cus_del_other" },
      },
    });
    mockStripe.subscriptions.retrieve.mockResolvedValue({
      items: { data: [{ price: { product: { id: "prod_OTHER" } } }] },
    });
    const { downgradeByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.subscription.deleted" }));
    expect(res.status).toBe(200);
    expect(downgradeByCustomer).not.toHaveBeenCalled();
  });

  it("returns 200 for unknown event type", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "invoice.payment_failed",
      data: { object: {} },
    });
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "invoice.payment_failed" }));
    expect(res.status).toBe(200);
    expect(res.body.received).toBe(true);
  });

  it("returns 200 when session retrieve fails on checkout.session.completed", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "checkout.session.completed",
      data: {
        object: {
          id: "cs_fail_retrieve",
          customer: "cus_fail",
          customer_details: { email: "fail@test.com" },
        },
      },
    });
    mockStripe.checkout.sessions.retrieve.mockRejectedValue(new Error("Stripe retrieve failed"));
    const { createProKey } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "checkout.session.completed" }));
    expect(res.status).toBe(200);
    expect(res.body.received).toBe(true);
    expect(createProKey).not.toHaveBeenCalled();
  });

  it("syncs email on customer.updated", async () => {
    mockStripe.webhooks.constructEvent.mockReturnValue({
      type: "customer.updated",
      data: {
        object: { id: "cus_email", email: "newemail@test.com" },
      },
    });
    const { updateEmailByCustomer } = await import("../services/keys.js");
    const res = await request(app)
      .post("/v1/billing/webhook")
      .set("content-type", "application/json")
      .set("stripe-signature", "valid_sig")
      .send(JSON.stringify({ type: "customer.updated" }));
    expect(res.status).toBe(200);
    expect(updateEmailByCustomer).toHaveBeenCalledWith("cus_email", "newemail@test.com");
  });
});

describe("Provisioned Sessions TTL (Memory Leak Fix)", () => {
  it("should allow fresh entries that haven't expired", async () => {
    mockStripe.checkout.sessions.retrieve.mockResolvedValue({
      id: "cs_fresh",
      customer: "cus_fresh",
      customer_details: { email: "fresh@test.com" },
    });

    // First call - should provision
    const res1 = await request(app).get("/v1/billing/success?session_id=cs_fresh");
    expect(res1.status).toBe(200);
    expect(res1.text).toContain("Welcome to Pro");

    // Second call immediately - should be duplicate (409)  
    const res2 = await request(app).get("/v1/billing/success?session_id=cs_fresh");
    expect(res2.status).toBe(409);
    expect(res2.body.error).toContain("already been used");
  });

  it("should remove stale entries older than 24 hours from provisionedSessions", async () => {
    // This test will verify that the cleanup mechanism removes old entries
    // For now, this will fail because the current implementation doesn't have TTL
    
    // Mock Date.now to control time
    const originalDateNow = Date.now;
    let currentTime = 1640995200000; // Jan 1, 2022 00:00:00 GMT
    vi.spyOn(Date, 'now').mockImplementation(() => currentTime);

    try {
      mockStripe.checkout.sessions.retrieve.mockResolvedValue({
        id: "cs_old",
        customer: "cus_old", 
        customer_details: { email: "old@test.com" },
      });

      // Add an entry at time T
      const res1 = await request(app).get("/v1/billing/success?session_id=cs_old");
      expect(res1.status).toBe(200);

      // Advance time by 25 hours (more than 24h TTL)
      currentTime += 25 * 60 * 60 * 1000;

      // The old entry should be cleaned up and session should work again
      const { findKeyByCustomerId } = await import("../services/keys.js");
      vi.mocked(findKeyByCustomerId).mockResolvedValueOnce(null); // No existing key in DB
      
      const res2 = await request(app).get("/v1/billing/success?session_id=cs_old");
      expect(res2.status).toBe(200); // Should provision again, not 409
      expect(res2.text).toContain("Welcome to Pro");

    } finally {
      vi.restoreAllMocks();
      Date.now = originalDateNow;
    }
  });

  it("should preserve fresh entries during cleanup", async () => {
    // This test verifies that cleanup doesn't remove fresh entries
    const originalDateNow = Date.now;
    let currentTime = 1640995200000;
    vi.spyOn(Date, 'now').mockImplementation(() => currentTime);

    try {
      // Add an old entry
      mockStripe.checkout.sessions.retrieve.mockResolvedValue({
        id: "cs_stale",
        customer: "cus_stale",
        customer_details: { email: "stale@test.com" },
      });
      await request(app).get("/v1/billing/success?session_id=cs_stale");

      // Advance time by 1 hour
      currentTime += 60 * 60 * 1000;

      // Add a fresh entry  
      mockStripe.checkout.sessions.retrieve.mockResolvedValue({
        id: "cs_recent",
        customer: "cus_recent",
        customer_details: { email: "recent@test.com" },
      });
      await request(app).get("/v1/billing/success?session_id=cs_recent");

      // Advance time by 24 more hours (stale entry is now 25h old, recent is 24h old)
      currentTime += 24 * 60 * 60 * 1000;

      // Recent entry should still be treated as duplicate (preserved), stale should be cleaned
      const res = await request(app).get("/v1/billing/success?session_id=cs_recent");
      expect(res.status).toBe(409); // Still duplicate - not cleaned up
      expect(res.body.error).toContain("already been used");

    } finally {
      vi.restoreAllMocks(); 
      Date.now = originalDateNow;
    }
  });

  it("should have bounded size even with many entries", async () => {
    // This test verifies that the Set/Map doesn't grow unbounded
    // We'll check that it doesn't exceed a reasonable size
    const originalDateNow = Date.now;
    let currentTime = 1640995200000;
    vi.spyOn(Date, 'now').mockImplementation(() => currentTime);

    try {
      // Create many entries over time
      for (let i = 0; i < 50; i++) {
        mockStripe.checkout.sessions.retrieve.mockResolvedValue({
          id: `cs_bulk_${i}`,
          customer: `cus_bulk_${i}`,
          customer_details: { email: `bulk${i}@test.com` },
        });
        
        await request(app).get(`/v1/billing/success?session_id=cs_bulk_${i}`);
        
        // Advance time by 1 hour each iteration
        currentTime += 60 * 60 * 1000;
      }

      // After processing 50 entries over 50 hours, old ones should be cleaned up
      // The first ~25 entries should be expired (older than 24h)
      
      // Try to use a very old session - should work again (cleaned up)
      const { findKeyByCustomerId } = await import("../services/keys.js");
      vi.mocked(findKeyByCustomerId).mockResolvedValueOnce(null);
      
      const res = await request(app).get("/v1/billing/success?session_id=cs_bulk_0");
      expect(res.status).toBe(200); // Should provision again, indicating it was cleaned up

    } finally {
      vi.restoreAllMocks();
      Date.now = originalDateNow; 
    }
  });
});