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

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

import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { Types } from 'mongoose';
import { DroneStatus, IDroneRegistration } from '@gadget/api';
import { createMockUser, createMockDroneRegistration } from './fixtures';

// Mock SocketService
vi.mock('../src/services/socket');

import DroneService from '../src/services/drone';
import SocketService from '../src/services/socket';

describe('DroneService.requestTermination', () => {
  let mockUser: any;
  let mockDrone: any;
  let mockDroneSession: any;
  let mockSocket: any;

  beforeEach(() => {
    vi.clearAllMocks();
    mockUser = createMockUser();
    mockDrone = createMockDroneRegistration(mockUser);
    mockSocket = {
      emit: vi.fn(),
    };
    mockDroneSession = {
      socket: mockSocket,
      registration: mockDrone,
    };
    
    // Setup mocks for DroneService methods
    vi.spyOn(DroneService, 'getById').mockResolvedValue(mockDrone);
    vi.spyOn(DroneService, 'setStatus').mockResolvedValue(mockDrone);
    vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession);
  });

  it('should return error if drone is already offline', async () => {
    const offlineDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Offline,
    });
    vi.spyOn(DroneService, 'getById').mockResolvedValue(offlineDrone);

    const result = await DroneService.requestTermination(offlineDrone._id);

    expect(result.success).toBe(false);
    expect(result.message).toBe('Drone is already offline');
    expect(SocketService.getDroneSession).not.toHaveBeenCalled();
  });

  it('should mark drone offline if not connected and return success', async () => {
    const connectedDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Available,
    });
    vi.spyOn(DroneService, 'getById').mockResolvedValue(connectedDrone);
    vi.spyOn(SocketService, 'getDroneSession').mockImplementation(() => {
      throw new Error('drone session not found');
    });

    const result = await DroneService.requestTermination(connectedDrone._id);

    expect(result.success).toBe(true);
    expect(result.message).toContain('not connected');
    expect(DroneService.setStatus).toHaveBeenCalledWith(connectedDrone, DroneStatus.Offline);
  });

  it('should emit requestTermination to drone socket', async () => {
    const availableDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Available,
    });
    vi.spyOn(DroneService, 'getById')
      .mockResolvedValueOnce(availableDrone)
      .mockResolvedValueOnce(availableDrone)
      .mockResolvedValueOnce(createMockDroneRegistration(mockUser, { status: DroneStatus.Offline }));
    vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession);

    // Mock socket emit to call the callback with success
    mockSocket.emit.mockImplementation((event: string, cb: Function) => {
      if (event === 'requestTermination') {
        cb(true);
      }
    });

    const result = await DroneService.requestTermination(availableDrone._id);

    expect(SocketService.getDroneSession).toHaveBeenCalledWith(availableDrone);
    expect(mockSocket.emit).toHaveBeenCalledWith('requestTermination', expect.any(Function));
    expect(result.success).toBe(true);
  });

  it('should handle drone rejection of termination', async () => {
    const availableDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Available,
    });
    vi.spyOn(DroneService, 'getById').mockResolvedValue(availableDrone);
    vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession);

    // Mock socket emit to call the callback with failure
    mockSocket.emit.mockImplementation((event: string, cb: Function) => {
      if (event === 'requestTermination') {
        cb(false);
      }
    });

    const result = await DroneService.requestTermination(availableDrone._id);

    expect(result.success).toBe(false);
    expect(result.message).toBe('Drone rejected termination request');
  });

  it('should timeout and force offline after 10 seconds', async () => {
    vi.useFakeTimers();
    
    const availableDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Available,
    });
    vi.spyOn(DroneService, 'getById').mockResolvedValue(availableDrone);
    vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession);

    // Mock socket emit but never call the callback (simulating no response)
    mockSocket.emit.mockImplementation(() => {
      // Never call callback
    });

    const terminationPromise = DroneService.requestTermination(availableDrone._id);

    // Advance time past 10 second timeout
    await vi.advanceTimersByTimeAsync(10000);

    const result = await terminationPromise;

    expect(result.success).toBe(true);
    expect(result.message).toContain('timed out');
    expect(DroneService.setStatus).toHaveBeenCalledWith(availableDrone, DroneStatus.Offline);

    vi.useRealTimers();
  });

  it('should poll for offline status after drone accepts termination', async () => {
    vi.useFakeTimers();
    
    const availableDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Available,
    });
    const offlineDrone = createMockDroneRegistration(mockUser, {
      status: DroneStatus.Offline,
    });
    
    vi.spyOn(DroneService, 'getById')
      .mockResolvedValueOnce(availableDrone)  // Initial check
      .mockResolvedValueOnce(offlineDrone);   // Poll after accept
    vi.spyOn(SocketService, 'getDroneSession').mockReturnValue(mockDroneSession);

    // Mock socket emit to call callback with success immediately
    mockSocket.emit.mockImplementation((event: string, cb: Function) => {
      if (event === 'requestTermination') {
        cb(true);
      }
    });

    const terminationPromise = DroneService.requestTermination(availableDrone._id);

    // Advance time to allow callback and one poll cycle (500ms interval)
    await vi.advanceTimersByTimeAsync(600);

    const result = await terminationPromise;

    expect(result.success).toBe(true);
    expect(result.message).toBeUndefined();

    vi.useRealTimers();
  });
});