stack/apps/orchestrator/src/valkey/valkey.service.spec.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { ConfigService } from '@nestjs/config';
import { ValkeyService } from './valkey.service';
import type { TaskState, AgentState, OrchestratorEvent } from './types';

// Create mock client methods that will be shared
const mockClient = {
  getTaskState: vi.fn(),
  setTaskState: vi.fn(),
  deleteTaskState: vi.fn(),
  updateTaskStatus: vi.fn(),
  listTasks: vi.fn(),
  getAgentState: vi.fn(),
  setAgentState: vi.fn(),
  deleteAgentState: vi.fn(),
  updateAgentStatus: vi.fn(),
  listAgents: vi.fn(),
  publishEvent: vi.fn(),
  subscribeToEvents: vi.fn(),
  disconnect: vi.fn(),
};

// Mock ValkeyClient before importing
vi.mock('./valkey.client', () => {
  return {
    ValkeyClient: class {
      constructor() {
        return mockClient;
      }
    },
  };
});

describe('ValkeyService', () => {
  let service: ValkeyService;
  let mockConfigService: ConfigService;

  beforeEach(() => {
    // Clear all mock calls
    vi.clearAllMocks();

    // Create mock config service
    mockConfigService = {
      get: vi.fn((key: string, defaultValue?: unknown) => {
        const config: Record<string, unknown> = {
          'orchestrator.valkey.host': 'localhost',
          'orchestrator.valkey.port': 6379,
        };
        return config[key] ?? defaultValue;
      }),
    } as any;

    // Create service directly
    service = new ValkeyService(mockConfigService);
  });

  describe('Initialization', () => {
    it('should be defined', () => {
      expect(service).toBeDefined();
    });

    it('should create ValkeyClient with config from ConfigService', () => {
      expect(mockConfigService.get).toHaveBeenCalledWith('orchestrator.valkey.host', 'localhost');
      expect(mockConfigService.get).toHaveBeenCalledWith('orchestrator.valkey.port', 6379);
    });

    it('should use password from config if provided', () => {
      const configWithPassword = {
        get: vi.fn((key: string, defaultValue?: unknown) => {
          const config: Record<string, unknown> = {
            'orchestrator.valkey.host': 'localhost',
            'orchestrator.valkey.port': 6379,
            'orchestrator.valkey.password': 'secret',
          };
          return config[key] ?? defaultValue;
        }),
      } as any;

      const serviceWithPassword = new ValkeyService(configWithPassword);

      expect(configWithPassword.get).toHaveBeenCalledWith('orchestrator.valkey.password');
    });
  });

  describe('Lifecycle', () => {
    it('should disconnect on module destroy', async () => {
      mockClient.disconnect.mockResolvedValue(undefined);

      await service.onModuleDestroy();

      expect(mockClient.disconnect).toHaveBeenCalled();
    });
  });

  describe('Task State Management', () => {
    const mockTaskState: TaskState = {
      taskId: 'task-123',
      status: 'pending',
      context: {
        repository: 'https://github.com/example/repo',
        branch: 'main',
        workItems: ['item-1'],
      },
      createdAt: '2026-02-02T10:00:00Z',
      updatedAt: '2026-02-02T10:00:00Z',
    };

    it('should get task state', async () => {
      mockClient.getTaskState.mockResolvedValue(mockTaskState);

      const result = await service.getTaskState('task-123');

      expect(mockClient.getTaskState).toHaveBeenCalledWith('task-123');
      expect(result).toEqual(mockTaskState);
    });

    it('should set task state', async () => {
      mockClient.setTaskState.mockResolvedValue(undefined);

      await service.setTaskState(mockTaskState);

      expect(mockClient.setTaskState).toHaveBeenCalledWith(mockTaskState);
    });

    it('should delete task state', async () => {
      mockClient.deleteTaskState.mockResolvedValue(undefined);

      await service.deleteTaskState('task-123');

      expect(mockClient.deleteTaskState).toHaveBeenCalledWith('task-123');
    });

    it('should update task status', async () => {
      const updatedTask = { ...mockTaskState, status: 'assigned' as const };
      mockClient.updateTaskStatus.mockResolvedValue(updatedTask);

      const result = await service.updateTaskStatus('task-123', 'assigned', 'agent-456');

      expect(mockClient.updateTaskStatus).toHaveBeenCalledWith(
        'task-123',
        'assigned',
        'agent-456',
        undefined
      );
      expect(result).toEqual(updatedTask);
    });

    it('should list all tasks', async () => {
      const tasks = [mockTaskState];
      mockClient.listTasks.mockResolvedValue(tasks);

      const result = await service.listTasks();

      expect(mockClient.listTasks).toHaveBeenCalled();
      expect(result).toEqual(tasks);
    });
  });

  describe('Agent State Management', () => {
    const mockAgentState: AgentState = {
      agentId: 'agent-456',
      status: 'spawning',
      taskId: 'task-123',
    };

    it('should get agent state', async () => {
      mockClient.getAgentState.mockResolvedValue(mockAgentState);

      const result = await service.getAgentState('agent-456');

      expect(mockClient.getAgentState).toHaveBeenCalledWith('agent-456');
      expect(result).toEqual(mockAgentState);
    });

    it('should set agent state', async () => {
      mockClient.setAgentState.mockResolvedValue(undefined);

      await service.setAgentState(mockAgentState);

      expect(mockClient.setAgentState).toHaveBeenCalledWith(mockAgentState);
    });

    it('should delete agent state', async () => {
      mockClient.deleteAgentState.mockResolvedValue(undefined);

      await service.deleteAgentState('agent-456');

      expect(mockClient.deleteAgentState).toHaveBeenCalledWith('agent-456');
    });

    it('should update agent status', async () => {
      const updatedAgent = { ...mockAgentState, status: 'running' as const };
      mockClient.updateAgentStatus.mockResolvedValue(updatedAgent);

      const result = await service.updateAgentStatus('agent-456', 'running');

      expect(mockClient.updateAgentStatus).toHaveBeenCalledWith(
        'agent-456',
        'running',
        undefined
      );
      expect(result).toEqual(updatedAgent);
    });

    it('should list all agents', async () => {
      const agents = [mockAgentState];
      mockClient.listAgents.mockResolvedValue(agents);

      const result = await service.listAgents();

      expect(mockClient.listAgents).toHaveBeenCalled();
      expect(result).toEqual(agents);
    });
  });

  describe('Event Pub/Sub', () => {
    const mockEvent: OrchestratorEvent = {
      type: 'agent.spawned',
      agentId: 'agent-456',
      taskId: 'task-123',
      timestamp: '2026-02-02T10:00:00Z',
    };

    it('should publish events', async () => {
      mockClient.publishEvent.mockResolvedValue(undefined);

      await service.publishEvent(mockEvent);

      expect(mockClient.publishEvent).toHaveBeenCalledWith(mockEvent);
    });

    it('should subscribe to events', async () => {
      mockClient.subscribeToEvents.mockResolvedValue(undefined);

      const handler = vi.fn();
      await service.subscribeToEvents(handler);

      expect(mockClient.subscribeToEvents).toHaveBeenCalledWith(handler);
    });
  });

  describe('Convenience Methods', () => {
    it('should create task state with timestamps', async () => {
      mockClient.setTaskState.mockResolvedValue(undefined);

      const context = {
        repository: 'https://github.com/example/repo',
        branch: 'main',
        workItems: ['item-1'],
      };

      await service.createTask('task-123', context);

      expect(mockClient.setTaskState).toHaveBeenCalledWith({
        taskId: 'task-123',
        status: 'pending',
        context,
        createdAt: expect.any(String),
        updatedAt: expect.any(String),
      });
    });

    it('should create agent state', async () => {
      mockClient.setAgentState.mockResolvedValue(undefined);

      await service.createAgent('agent-456', 'task-123');

      expect(mockClient.setAgentState).toHaveBeenCalledWith({
        agentId: 'agent-456',
        status: 'spawning',
        taskId: 'task-123',
      });
    });
  });
});