gadget/gadget-code/tests/socket-service.test.ts

// tests/socket-service.test.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Types } from "mongoose";
import { DroneStatus } from "@gadget/api";
import {
  createMockSocket,
  createMockUser,
  createMockDroneRegistration,
  createMockChatSession,
  createMockProject,
} from "./fixtures";
import { nanoid } from "nanoid";

// Mock the entire socket service module
vi.mock("../src/services/socket", async () => {
  const chatSessionIndex = new Map();
  return {
    default: {
      codeSessions: new Map(),
      droneSessions: new Map(),
      chatSessionIndex: chatSessionIndex,
      droneRegistrationIndex: new Map(),
      codeSessionUserIndex: new Map(),
      getDroneSession: vi.fn(),
      getCodeSession: vi.fn(),
      getCodeSessionByChatSessionId: vi.fn(),
      registerChatSession: vi.fn((chatSessionId: string, codeSession: any) => {
        chatSessionIndex.set(chatSessionId, codeSession);
      }),
      unregisterChatSession: vi.fn((chatSessionId: string) => {
        chatSessionIndex.delete(chatSessionId);
      }),
    },
  };
});

describe("SocketService Session Indexing", () => {
  let SocketService: any;
  let mockUser: any;
  let mockDrone: any;
  let mockChatSession: any;
  let mockProject: any;

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

    const mod = await import("../src/services/socket");
    SocketService = mod.default;

    // Clear all session maps
    SocketService.codeSessions.clear();
    SocketService.droneSessions.clear();
    SocketService.chatSessionIndex.clear();
    SocketService.droneRegistrationIndex.clear();
    SocketService.codeSessionUserIndex.clear();

    mockUser = createMockUser();
    mockDrone = createMockDroneRegistration(mockUser);
    mockChatSession = createMockChatSession(mockUser);
    mockProject = createMockProject(mockUser);
  });

  describe("Drone Session Indexing", () => {
    it("should store drone session by socket.id and registration._id", () => {
      const mockSocket = createMockSocket("drone-socket-123");
      const mockDroneSession = {
        socket: mockSocket,
        registration: mockDrone,
        type: "drone",
      };

      // Store in both indexes
      SocketService.droneSessions.set(mockSocket.id, mockDroneSession);
      SocketService.droneRegistrationIndex.set(mockDrone._id, mockDroneSession);

      // Verify both lookups work
      expect(SocketService.droneSessions.get(mockSocket.id)).toBe(
        mockDroneSession,
      );
      expect(SocketService.droneRegistrationIndex.get(mockDrone._id)).toBe(
        mockDroneSession,
      );
    });

    it("should find drone session by registration._id", () => {
      const mockSocket = createMockSocket("drone-socket-456");
      const mockDroneSession = {
        socket: mockSocket,
        registration: mockDrone,
        type: "drone",
      };

      // Store in both indexes
      SocketService.droneSessions.set(mockSocket.id, mockDroneSession);
      SocketService.droneRegistrationIndex.set(mockDrone._id, mockDroneSession);

      // Mock getDroneSession to use the registration index
      SocketService.getDroneSession.mockImplementation((registration: any) => {
        const session = SocketService.droneRegistrationIndex.get(
          registration._id,
        );
        if (!session) {
          const error = new Error("drone session not found");
          (error as any).statusCode = 404;
          throw error;
        }
        return session;
      });

      const found = SocketService.getDroneSession(mockDrone);
      expect(found).toBe(mockDroneSession);
    });

    it("should throw 404 when drone session not found", () => {
      const nonExistentDrone = createMockDroneRegistration(mockUser);

      SocketService.getDroneSession.mockImplementation(() => {
        const error = new Error("drone session not found");
        (error as any).statusCode = 404;
        throw error;
      });

      expect(() => SocketService.getDroneSession(nonExistentDrone)).toThrow(
        "drone session not found",
      );
      expect(() =>
        SocketService.getDroneSession(nonExistentDrone),
      ).toThrowError(
        expect.objectContaining({
          statusCode: 404,
        }),
      );
    });

    it("should remove drone session from all indexes on disconnect", () => {
      const mockSocket = createMockSocket("drone-socket-789");
      const mockDroneSession = {
        socket: mockSocket,
        registration: mockDrone,
        type: "drone",
      };

      // Store in indexes
      SocketService.droneSessions.set(mockSocket.id, mockDroneSession);
      SocketService.droneRegistrationIndex.set(mockDrone._id, mockDroneSession);

      // Simulate disconnect
      SocketService.droneSessions.delete(mockSocket.id);
      SocketService.droneRegistrationIndex.delete(mockDrone._id);

      // Verify removal from all indexes
      expect(SocketService.droneSessions.get(mockSocket.id)).toBeUndefined();
      expect(
        SocketService.droneRegistrationIndex.get(mockDrone._id),
      ).toBeUndefined();
    });
  });

  describe("Code Session Indexing", () => {
    it("should store code session by socket.id and user._id", () => {
      const mockSocket = createMockSocket("code-socket-123");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      // Store in both indexes
      SocketService.codeSessions.set(mockSocket.id, mockCodeSession);
      SocketService.codeSessionUserIndex.set(mockUser._id, mockCodeSession);

      // Verify both lookups work
      expect(SocketService.codeSessions.get(mockSocket.id)).toBe(
        mockCodeSession,
      );
      expect(SocketService.codeSessionUserIndex.get(mockUser._id)).toBe(
        mockCodeSession,
      );
    });

    it("should find code session by user._id", () => {
      const mockSocket = createMockSocket("code-socket-456");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      // Store in both indexes
      SocketService.codeSessions.set(mockSocket.id, mockCodeSession);
      SocketService.codeSessionUserIndex.set(mockUser._id, mockCodeSession);

      // Mock getCodeSession to use the user index
      SocketService.getCodeSession.mockImplementation((user: any) => {
        const session = SocketService.codeSessionUserIndex.get(user._id);
        if (!session) {
          const error = new Error("code session not found");
          (error as any).statusCode = 404;
          throw error;
        }
        return session;
      });

      const found = SocketService.getCodeSession(mockUser);
      expect(found).toBe(mockCodeSession);
    });

    it("should throw 404 when code session not found", () => {
      const nonExistentUser = createMockUser();

      SocketService.getCodeSession.mockImplementation(() => {
        const error = new Error("code session not found");
        (error as any).statusCode = 404;
        throw error;
      });

      expect(() => SocketService.getCodeSession(nonExistentUser)).toThrow(
        "code session not found",
      );
      expect(() => SocketService.getCodeSession(nonExistentUser)).toThrowError(
        expect.objectContaining({
          statusCode: 404,
        }),
      );
    });

    it("should remove code session from all indexes on disconnect", () => {
      const mockSocket = createMockSocket("code-socket-789");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      // Store in indexes
      SocketService.codeSessions.set(mockSocket.id, mockCodeSession);
      SocketService.codeSessionUserIndex.set(mockUser._id, mockCodeSession);

      // Simulate disconnect
      SocketService.codeSessions.delete(mockSocket.id);
      SocketService.codeSessionUserIndex.delete(mockUser._id);

      // Verify removal from all indexes
      expect(SocketService.codeSessions.get(mockSocket.id)).toBeUndefined();
      expect(
        SocketService.codeSessionUserIndex.get(mockUser._id),
      ).toBeUndefined();
    });
  });

  describe("Chat Session Index", () => {
    it("should register and retrieve code session by chatSessionId", () => {
      const mockSocket = createMockSocket("code-socket-chat");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      const chatSessionId = mockChatSession._id;

      // Mock the retrieval to return our session after registration
      SocketService.registerChatSession(chatSessionId, mockCodeSession);
      SocketService.getCodeSessionByChatSessionId.mockReturnValue(
        mockCodeSession,
      );

      const found = SocketService.getCodeSessionByChatSessionId(
        mockChatSession._id,
      );
      expect(found).toBe(mockCodeSession);
      expect(SocketService.registerChatSession).toHaveBeenCalledWith(
        chatSessionId,
        mockCodeSession,
      );
      expect(SocketService.getCodeSessionByChatSessionId).toHaveBeenCalledWith(
        mockChatSession._id,
      );
    });

    it("should handle chatSessionId as string", () => {
      const mockSocket = createMockSocket("code-socket-chat2");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      const chatSessionId = mockChatSession._id;
      SocketService.registerChatSession(chatSessionId, mockCodeSession);
      SocketService.getCodeSessionByChatSessionId.mockReturnValue(
        mockCodeSession,
      );

      // Test with string
      const found1 = SocketService.getCodeSessionByChatSessionId(chatSessionId);
      expect(found1).toBe(mockCodeSession);

      // Test with same ID again
      const found2 = SocketService.getCodeSessionByChatSessionId(
        mockChatSession._id,
      );
      expect(found2).toBe(mockCodeSession);
    });

    it("should throw 404 when chat session not found", () => {
      const nonExistentChatSessionId = nanoid();

      SocketService.getCodeSessionByChatSessionId.mockImplementation(() => {
        const error = new Error("code session not found for chat session");
        (error as any).statusCode = 404;
        throw error;
      });

      expect(() =>
        SocketService.getCodeSessionByChatSessionId(nonExistentChatSessionId),
      ).toThrow("code session not found for chat session");
      expect(() =>
        SocketService.getCodeSessionByChatSessionId(nonExistentChatSessionId),
      ).toThrowError(
        expect.objectContaining({
          statusCode: 404,
        }),
      );
    });

    it("should unregister chat session", () => {
      const mockSocket = createMockSocket("code-socket-chat3");
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: "code",
      };

      const chatSessionId = mockChatSession._id;
      SocketService.registerChatSession(chatSessionId, mockCodeSession);
      SocketService.getCodeSessionByChatSessionId.mockReturnValue(
        mockCodeSession,
      );

      // Verify it's registered
      const found = SocketService.getCodeSessionByChatSessionId(chatSessionId);
      expect(found).toBe(mockCodeSession);

      // Unregister
      SocketService.unregisterChatSession(chatSessionId);
      expect(SocketService.unregisterChatSession).toHaveBeenCalledWith(
        chatSessionId,
      );
    });
  });
});