gadget/gadget-code/tests/code-session.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { Types } from 'mongoose';
import { CodeSession } from '../src/lib/code-session';
import { IChatSession, IProject, IUser, IDroneRegistration, ChatTurnStatus } from '@gadget/api';
import SocketService from '../src/services/socket';
import { ChatTurn } from '../src/models/chat-turn';

// Mock dependencies
vi.mock('../src/services/socket');
vi.mock('../src/models/chat-turn');

describe('CodeSession', () => {
  let mockSocket: any;
  let mockUser: IUser;
  let mockDrone: IDroneRegistration;
  let mockChatSession: IChatSession;
  let mockProject: IProject;
  let codeSession: CodeSession;

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

    mockSocket = {
      id: 'test-socket-id',
      on: vi.fn(),
      emit: vi.fn(),
    };

    mockUser = {
      _id: new Types.ObjectId(),
      email: 'test@example.com',
      displayName: 'Test User',
    } as IUser;

    mockDrone = {
      _id: new Types.ObjectId(),
      hostname: 'test-host',
      workspaceDir: '/test/workspace',
      status: 'available',
    } as IDroneRegistration;

    mockChatSession = {
      _id: new Types.ObjectId(),
      name: 'Test Session',
      mode: 'build',
      provider: new Types.ObjectId(),
      selectedModel: 'llama3.1',
    } as IChatSession;

    mockProject = {
      _id: new Types.ObjectId(),
      slug: 'test-project',
      name: 'Test Project',
    } as IProject;

    codeSession = new CodeSession(mockSocket, mockUser);
  });

  describe('setSelectedDrone', () => {
    it('should set the selected drone', () => {
      codeSession.setSelectedDrone(mockDrone);
      // Can't directly access private property, but we can verify through behavior
      expect(() => codeSession.setSelectedDrone(mockDrone)).not.toThrow();
    });
  });

  describe('setChatSession', () => {
    it('should set the chat session and project', () => {
      codeSession.setChatSession(mockChatSession, mockProject);
      expect(() => codeSession.setChatSession(mockChatSession, mockProject)).not.toThrow();
    });
  });

  describe('onSubmitPrompt', () => {
    it('should throw error if no drone is selected', async () => {
      codeSession.setChatSession(mockChatSession, mockProject);
      
      await expect(codeSession.onSubmitPrompt('test prompt'))
        .rejects.toThrow('No drone selected');
    });

    it('should throw error if no chat session is active', async () => {
      codeSession.setSelectedDrone(mockDrone);
      
      await expect(codeSession.onSubmitPrompt('test prompt'))
        .rejects.toThrow('No chat session active');
    });

    it('should throw error if no project is selected', async () => {
      codeSession.setSelectedDrone(mockDrone);
      codeSession.setChatSession(mockChatSession, undefined as any);
      
      await expect(codeSession.onSubmitPrompt('test prompt'))
        .rejects.toThrow('No project selected');
    });

    it('should create a ChatTurn and emit processWorkOrder to drone', async () => {
      codeSession.setSelectedDrone(mockDrone);
      codeSession.setChatSession(mockChatSession, mockProject);

      const mockTurn = {
        _id: new Types.ObjectId(),
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(ChatTurn).mockImplementation(function() {
        return mockTurn as any;
      });

      const mockDroneSession = {
        socket: {
          emit: vi.fn(),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(mockDroneSession as any);

      await codeSession.onSubmitPrompt('test prompt');

      expect(ChatTurn).toHaveBeenCalledWith(expect.objectContaining({
        user: mockUser._id,
        project: mockProject._id,
        session: mockChatSession._id,
        provider: mockChatSession.provider,
        llm: mockChatSession.selectedModel,
        status: ChatTurnStatus.Processing,
        prompts: {
          user: 'test prompt',
          system: undefined,
        },
      }));

      expect(mockTurn.save).toHaveBeenCalled();
      expect(mockDroneSession.socket.emit).toHaveBeenCalledWith(
        'processWorkOrder',
        mockDrone,
        mockProject,
        mockChatSession,
        mockTurn,
        expect.any(Function)
      );
    });

    it('should update ChatTurn to Error status if drone rejects work order', async () => {
      codeSession.setSelectedDrone(mockDrone);
      codeSession.setChatSession(mockChatSession, mockProject);

      const mockTurn = {
        _id: new Types.ObjectId(),
        status: ChatTurnStatus.Processing,
        response: '',
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(ChatTurn).mockImplementation(function() {
        return mockTurn as any;
      });

      const mockDroneSession = {
        socket: {
          emit: vi.fn((event, ...args) => {
            const callback = args[args.length - 1];
            callback(false, 'Drone is busy');
          }),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(mockDroneSession as any);

      await codeSession.onSubmitPrompt('test prompt');

      expect(mockTurn.status).toBe(ChatTurnStatus.Error);
      expect(mockTurn.response).toBe('Drone is busy');
      expect(mockTurn.save).toHaveBeenCalled();
    });
  });

  describe('onRequestSessionLock', () => {
    it('should set selected drone, chat session, and project on success', () => {
      const mockDroneSession = {
        socket: {
          emit: vi.fn((event, ...args) => {
            const callback = args[args.length - 1];
            callback(true, mockChatSession._id.toHexString());
          }),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(mockDroneSession as any);

      const callback = vi.fn();
      codeSession.onRequestSessionLock(mockDrone, mockProject, mockChatSession, callback);

      expect(callback).toHaveBeenCalledWith(true, mockChatSession._id.toHexString());
    });

    it('should not set session data on failure', () => {
      const mockDroneSession = {
        socket: {
          emit: vi.fn((event, ...args) => {
            const callback = args[args.length - 1];
            callback(false, '');
          }),
        },
        setChatSessionId: vi.fn(),
        setCurrentTurnId: vi.fn(),
      };
      vi.mocked(SocketService.getDroneSession).mockReturnValue(mockDroneSession as any);

      const callback = vi.fn();
      codeSession.onRequestSessionLock(mockDrone, mockProject, mockChatSession, callback);

      expect(callback).toHaveBeenCalledWith(false, '');
    });
  });
});