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';

// 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.toHexString(),
        mockDroneSession
      );

      // Verify both lookups work
      expect(SocketService.droneSessions.get(mockSocket.id)).toBe(mockDroneSession);
      expect(SocketService.droneRegistrationIndex.get(mockDrone._id.toHexString())).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.toHexString(),
        mockDroneSession
      );

      // Mock getDroneSession to use the registration index
      SocketService.getDroneSession.mockImplementation((registration: any) => {
        const session = SocketService.droneRegistrationIndex.get(registration._id.toHexString());
        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.toHexString(),
        mockDroneSession
      );

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

      // Verify removal from all indexes
      expect(SocketService.droneSessions.get(mockSocket.id)).toBeUndefined();
      expect(SocketService.droneRegistrationIndex.get(mockDrone._id.toHexString())).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.toHexString(),
        mockCodeSession
      );

      // Verify both lookups work
      expect(SocketService.codeSessions.get(mockSocket.id)).toBe(mockCodeSession);
      expect(SocketService.codeSessionUserIndex.get(mockUser._id.toHexString())).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.toHexString(),
        mockCodeSession
      );

      // Mock getCodeSession to use the user index
      SocketService.getCodeSession.mockImplementation((user: any) => {
        const session = SocketService.codeSessionUserIndex.get(user._id.toHexString());
        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.toHexString(),
        mockCodeSession
      );

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

      // Verify removal from all indexes
      expect(SocketService.codeSessions.get(mockSocket.id)).toBeUndefined();
      expect(SocketService.codeSessionUserIndex.get(mockUser._id.toHexString())).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.toHexString();
      
      // 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 or ObjectId', () => {
      const mockSocket = createMockSocket('code-socket-chat2');
      const mockCodeSession = {
        socket: mockSocket,
        user: mockUser,
        type: 'code',
      };

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

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

      // Test with ObjectId
      const found2 = SocketService.getCodeSessionByChatSessionId(mockChatSession._id);
      expect(found2).toBe(mockCodeSession);
    });

    it('should throw 404 when chat session not found', () => {
      const nonExistentChatSessionId = new Types.ObjectId();
      
      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.toHexString();
      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);
    });
  });
});