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";
import { nanoid } from "nanoid";

// 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: nanoid(),
      email: "test@example.com",
      displayName: "Test User",
    } as IUser;

    mockDrone = {
      _id: nanoid(),
      hostname: "test-host",
      workspaceDir: "/test/workspace",
      status: "available",
    } as IDroneRegistration;

    mockProject = {
      _id: nanoid(),
      slug: "test-project",
      name: "Test Project",
    } as IProject;

    mockChatSession = {
      _id: nanoid(),
      name: "Test Session",
      mode: "build",
      provider: nanoid(),
      selectedModel: "llama3.1",
      user: mockUser,
      project: mockProject,
    } as IChatSession;

    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 return error if no drone is selected", async () => {
      codeSession.setChatSession(mockChatSession, mockProject);
      const cb = vi.fn();

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, { message: "No drone selected" });
    });

    it("should return error if no chat session is active", async () => {
      codeSession.setSelectedDrone(mockDrone);
      const cb = vi.fn();

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, { message: "No chat session active" });
    });

    it("should return error if no project is selected", async () => {
      codeSession.setSelectedDrone(mockDrone);
      codeSession.setChatSession(mockChatSession, undefined as any);
      const cb = vi.fn();

      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(cb).toHaveBeenCalledWith(false, { message: "No project selected" });
    });

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

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

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

      const cb = vi.fn();
      await codeSession.onSubmitPrompt("test prompt", cb);

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

      expect(mockTurn.save).toHaveBeenCalled();
      expect(mockDroneSession.socket.emit).toHaveBeenCalledWith(
        "processWorkOrder",
        mockDrone,
        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: nanoid(),
        status: ChatTurnStatus.Processing,
        errorMessage: "",
        save: vi.fn().mockResolvedValue(undefined),
      };
      vi.mocked(ChatTurn).mockImplementation(function () {
        return mockTurn as any;
      });
      (vi.mocked(ChatTurn) as any).populate = vi
        .fn()
        .mockResolvedValue(mockTurn);

      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,
      );

      const cb = vi.fn();
      await codeSession.onSubmitPrompt("test prompt", cb);

      expect(mockTurn.status).toBe(ChatTurnStatus.Error);
      expect(mockTurn.errorMessage).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);
          }),
        },
        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);
    });

    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, "");
    });
  });
});