stack/apps/orchestrator/src/common/guards/throttler.guard.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { ExecutionContext } from "@nestjs/common";
import { ThrottlerException, ThrottlerModuleOptions, ThrottlerStorage } from "@nestjs/throttler";
import { Reflector } from "@nestjs/core";
import { OrchestratorThrottlerGuard } from "./throttler.guard";

describe("OrchestratorThrottlerGuard", () => {
  let guard: OrchestratorThrottlerGuard;

  beforeEach(() => {
    // Create guard with minimal mocks for testing protected methods
    const options: ThrottlerModuleOptions = {
      throttlers: [{ name: "default", ttl: 60000, limit: 100 }],
    };
    const storageService = {} as ThrottlerStorage;
    const reflector = {} as Reflector;

    guard = new OrchestratorThrottlerGuard(options, storageService, reflector);
  });

  describe("getTracker", () => {
    it("should extract IP from X-Forwarded-For header", async () => {
      const req = {
        headers: {
          "x-forwarded-for": "192.168.1.1, 10.0.0.1",
        },
        ip: "127.0.0.1",
      };

      // Access protected method for testing
      const tracker = await (
        guard as unknown as { getTracker: (req: unknown) => Promise<string> }
      ).getTracker(req);

      expect(tracker).toBe("192.168.1.1");
    });

    it("should handle X-Forwarded-For as array", async () => {
      const req = {
        headers: {
          "x-forwarded-for": ["192.168.1.1, 10.0.0.1"],
        },
        ip: "127.0.0.1",
      };

      const tracker = await (
        guard as unknown as { getTracker: (req: unknown) => Promise<string> }
      ).getTracker(req);

      expect(tracker).toBe("192.168.1.1");
    });

    it("should fallback to request IP when no X-Forwarded-For", async () => {
      const req = {
        headers: {},
        ip: "192.168.2.2",
      };

      const tracker = await (
        guard as unknown as { getTracker: (req: unknown) => Promise<string> }
      ).getTracker(req);

      expect(tracker).toBe("192.168.2.2");
    });

    it("should fallback to connection remoteAddress when no IP", async () => {
      const req = {
        headers: {},
        connection: {
          remoteAddress: "192.168.3.3",
        },
      };

      const tracker = await (
        guard as unknown as { getTracker: (req: unknown) => Promise<string> }
      ).getTracker(req);

      expect(tracker).toBe("192.168.3.3");
    });

    it("should return 'unknown' when no IP available", async () => {
      const req = {
        headers: {},
      };

      const tracker = await (
        guard as unknown as { getTracker: (req: unknown) => Promise<string> }
      ).getTracker(req);

      expect(tracker).toBe("unknown");
    });
  });

  describe("throwThrottlingException", () => {
    it("should throw ThrottlerException with endpoint info", () => {
      const mockRequest = {
        url: "/agents/spawn",
      };

      const mockContext = {
        switchToHttp: vi.fn().mockReturnValue({
          getRequest: vi.fn().mockReturnValue(mockRequest),
        }),
      } as unknown as ExecutionContext;

      expect(() => {
        (
          guard as unknown as { throwThrottlingException: (context: ExecutionContext) => void }
        ).throwThrottlingException(mockContext);
      }).toThrow(ThrottlerException);

      try {
        (
          guard as unknown as { throwThrottlingException: (context: ExecutionContext) => void }
        ).throwThrottlingException(mockContext);
      } catch (error) {
        expect(error).toBeInstanceOf(ThrottlerException);
        expect((error as ThrottlerException).message).toContain("/agents/spawn");
      }
    });
  });
});