directus-task-management/tests/services/ai/ai-task.service.test.ts

import 'reflect-metadata';
import { AITaskService } from '../../../src/services/ai/ai-task.service';
import { OpenAIService } from '../../../src/services/ai/openai.service';
import { LangChainService } from '../../../src/services/ai/langchain.service';
import { AIConfig } from '../../../src/config/ai.config';
import { Repository } from 'typeorm';
import Redis from 'ioredis';

// Mock dependencies
jest.mock('../../../src/services/ai/openai.service');
jest.mock('../../../src/services/ai/langchain.service');
jest.mock('typeorm', () => ({
  getRepository: jest.fn(),
  Entity: jest.fn(),
  PrimaryGeneratedColumn: jest.fn(),
  Column: jest.fn(),
  CreateDateColumn: jest.fn(),
  UpdateDateColumn: jest.fn(),
  Index: jest.fn(),
}));
jest.mock('ioredis');
jest.mock('rate-limiter-flexible');

describe('AITaskService', () => {
  let service: AITaskService;
  let mockOpenAIService: jest.Mocked<OpenAIService>;
  let mockLangChainService: jest.Mocked<LangChainService>;
  let mockRedis: jest.Mocked<Redis>;
  let mockRepository: jest.Mocked<Repository<any>>;

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

    // Mock repository
    mockRepository = {
      create: jest.fn(),
      save: jest.fn(),
      count: jest.fn().mockResolvedValue(10),
    } as any;

    // Mock getRepository
    (require('typeorm').getRepository as jest.Mock).mockReturnValue(mockRepository);

    // Create mock services
    mockOpenAIService = {
      complete: jest.fn(),
      testConnection: jest.fn().mockResolvedValue(true),
      getUsageStats: jest.fn().mockReturnValue({
        requestCount: 5,
        totalTokensUsed: 1000,
        estimatedCost: 0.02,
      }),
    } as any;

    mockLangChainService = {
      createTaskFromNaturalLanguage: jest.fn(),
      breakdownTask: jest.fn(),
      getTaskSuggestions: jest.fn(),
      testConnection: jest.fn().mockResolvedValue(true),
    } as any;

    mockRedis = {
      get: jest.fn(),
      set: jest.fn(),
      del: jest.fn(),
      keys: jest.fn().mockResolvedValue([]),
      ping: jest.fn().mockResolvedValue('PONG'),
    } as any;

    // Create service instance
    service = new AITaskService(
      mockOpenAIService,
      mockLangChainService,
      mockRedis,
    );
  });

  describe('createTaskFromNaturalLanguage', () => {
    it('should create a task from natural language', async () => {
      const mockTask = {
        title: 'Fix login bug',
        description: 'Users cannot login with email',
        priority: 'high',
        complexity: 'minor',
        estimatedHours: 2,
        executionTimeMs: 150,
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);
      mockRedis.get.mockResolvedValue(null); // No cache hit

      const result = await service.createTaskFromNaturalLanguage({
        prompt: 'Fix the bug where users cant login with email',
        projectId: 'project-123',
      });

      expect(result).toMatchObject({
        title: 'Fix login bug',
        description: 'Users cannot login with email',
        priority: 'high',
        complexity: 'minor',
      });

      expect(mockLangChainService.createTaskFromNaturalLanguage).toHaveBeenCalledWith(
        'Fix the bug where users cant login with email',
        '',
        { temperature: 0.5 },
      );

      // Should cache the result
      expect(mockRedis.set).toHaveBeenCalled();
    });

    it('should enhance task with acceptance criteria if missing', async () => {
      const mockTask = {
        title: 'Add feature',
        description: 'New feature implementation',
        priority: 'medium',
        complexity: 'major',
        acceptanceCriteria: [],
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);
      mockOpenAIService.complete.mockResolvedValue({
        content: '1. Feature works correctly\n2. Tests pass\n3. Documentation updated',
        promptTokens: 50,
        completionTokens: 30,
        totalTokens: 80,
        model: 'gpt-4-turbo',
        finishReason: 'stop',
        executionTimeMs: 200,
      });
      mockRedis.get.mockResolvedValue(null);

      const result = await service.createTaskFromNaturalLanguage({
        prompt: 'Add new feature',
      });

      expect(result.acceptanceCriteria).toEqual([
        'Feature works correctly',
        'Tests pass',
        'Documentation updated',
      ]);

      expect(mockOpenAIService.complete).toHaveBeenCalled();
    });

    it('should use cached result if available', async () => {
      const cachedTask = {
        title: 'Cached task',
        description: 'From cache',
        priority: 'low',
        complexity: 'trivial',
      };

      mockRedis.get.mockResolvedValue(JSON.stringify(cachedTask));

      const result = await service.createTaskFromNaturalLanguage({
        prompt: 'Create a task',
      });

      expect(result).toEqual(cachedTask);
      expect(mockLangChainService.createTaskFromNaturalLanguage).not.toHaveBeenCalled();
    });

    it('should save AI context on success', async () => {
      const mockTask = {
        title: 'New task',
        description: 'Task description',
        priority: 'medium',
        complexity: 'minor',
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);
      mockRedis.get.mockResolvedValue(null);

      await service.createTaskFromNaturalLanguage({
        prompt: 'Create new task',
      });

      expect(mockRepository.create).toHaveBeenCalled();
      expect(mockRepository.save).toHaveBeenCalled();
    });

    it('should save error context on failure', async () => {
      mockLangChainService.createTaskFromNaturalLanguage.mockRejectedValue(
        new Error('AI service error'),
      );
      mockRedis.get.mockResolvedValue(null);

      await expect(service.createTaskFromNaturalLanguage({
        prompt: 'Create task',
      })).rejects.toThrow('AI service error');

      expect(mockRepository.create).toHaveBeenCalledWith(
        expect.objectContaining({
          context_type: 'error',
          success: false,
        }),
      );
    });
  });

  describe('breakdownTask', () => {
    it('should breakdown a task into subtasks', async () => {
      const mockBreakdown = {
        title: 'Main task',
        description: 'Task to break down',
        subtasks: [
          { id: '1', title: 'Subtask 1', description: 'First subtask' },
          { id: '2', title: 'Subtask 2', description: 'Second subtask' },
        ],
        complexity: 'major' as const,
        estimatedTotalHours: 8,
      };

      mockLangChainService.breakdownTask.mockResolvedValue(mockBreakdown);
      mockRedis.get.mockResolvedValue(null);

      const result = await service.breakdownTask(
        'task-123',
        'Implement authentication system',
        'Web application',
      );

      expect(result).toEqual(mockBreakdown);
      expect(mockLangChainService.breakdownTask).toHaveBeenCalledWith(
        'Implement authentication system',
        'Web application',
        { temperature: 0.3, maxTokens: 2000 },
      );

      // Should cache the result
      expect(mockRedis.set).toHaveBeenCalledWith(
        expect.stringContaining('ai:breakdown:task-123'),
        JSON.stringify(mockBreakdown),
        'EX',
        7200,
      );
    });

    it('should use cached breakdown if available', async () => {
      const cachedBreakdown = {
        title: 'Cached breakdown',
        subtasks: [],
        complexity: 'minor' as const,
      };

      mockRedis.get.mockResolvedValue(JSON.stringify(cachedBreakdown));

      const result = await service.breakdownTask('task-456', 'Task description');

      expect(result).toEqual(cachedBreakdown);
      expect(mockLangChainService.breakdownTask).not.toHaveBeenCalled();
    });
  });

  describe('getTaskSuggestions', () => {
    it('should generate task suggestions for a project', async () => {
      const mockSuggestionsText = `
        1. Implement user profile page - Users need to view and edit their profiles
        2. Add password reset functionality - Security requirement for production
        3. Setup email notifications - Keep users informed about important events
      `;

      mockLangChainService.getTaskSuggestions.mockResolvedValue(mockSuggestionsText);
      mockRedis.get.mockResolvedValue(null);

      const result = await service.getTaskSuggestions('project-123', {
        projectDescription: 'User management system',
        completedTasks: ['User registration', 'Login'],
        currentTasks: ['Dashboard'],
        goals: ['Complete MVP'],
      });

      expect(result).toHaveLength(3);
      expect(result[0]).toMatchObject({
        id: 'suggestion-1',
        title: 'Implement user profile page',
        priority: 'medium',
        confidence: 0.7,
      });

      expect(mockLangChainService.getTaskSuggestions).toHaveBeenCalledWith(
        {
          projectDescription: 'User management system',
          completedTasks: ['User registration', 'Login'],
          currentTasks: ['Dashboard'],
          goals: ['Complete MVP'],
        },
        { temperature: 0.7, maxTokens: 1500 },
      );
    });

    it('should parse suggestions with different formats', async () => {
      const mockSuggestionsText = `
        1) First task - Description here
        2. Second task - Another description
        3) Third task
      `;

      mockLangChainService.getTaskSuggestions.mockResolvedValue(mockSuggestionsText);
      mockRedis.get.mockResolvedValue(null);

      const result = await service.getTaskSuggestions('project-456', {
        projectDescription: 'Test project',
        completedTasks: [],
        currentTasks: [],
        goals: [],
      });

      expect(result).toHaveLength(3);
      expect(result[0].title).toBe('First task');
      expect(result[1].title).toBe('Second task');
      expect(result[2].title).toBe('Third task');
    });
  });

  describe('updateTaskContext', () => {
    it('should update task context with feedback', async () => {
      await service.updateTaskContext('task-789', 'Task needs more detail');

      expect(mockRepository.create).toHaveBeenCalledWith(
        expect.objectContaining({
          task_id: 'task-789',
          context_type: 'feedback',
          context_data: expect.objectContaining({
            feedback: 'Task needs more detail',
          }),
        }),
      );

      expect(mockRepository.save).toHaveBeenCalled();

      // Should invalidate cache
      expect(mockRedis.keys).toHaveBeenCalledWith('ai:breakdown:task-789');
    });
  });

  describe('warmCache', () => {
    it('should pre-populate cache with common queries', async () => {
      const mockTask = {
        title: 'Generic task',
        description: 'Task description',
        priority: 'medium',
        complexity: 'minor',
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);
      mockRedis.get.mockResolvedValue(null);

      await service.warmCache('project-123');

      // Should attempt to cache common prompts
      expect(mockLangChainService.createTaskFromNaturalLanguage).toHaveBeenCalledTimes(4);
      expect(mockLangChainService.createTaskFromNaturalLanguage).toHaveBeenCalledWith(
        expect.stringContaining('bug fix'),
        expect.any(String),
        expect.any(Object),
      );
    });

    it('should handle errors during cache warming', async () => {
      mockLangChainService.createTaskFromNaturalLanguage.mockRejectedValue(
        new Error('Service unavailable'),
      );
      mockRedis.get.mockResolvedValue(null);

      // Should not throw
      await expect(service.warmCache('project-456')).resolves.not.toThrow();
    });
  });

  describe('getHealthStatus', () => {
    it('should return health status of all services', async () => {
      mockOpenAIService.testConnection.mockResolvedValue(true);
      mockLangChainService.testConnection.mockResolvedValue(true);

      const health = await service.getHealthStatus();

      expect(health).toEqual({
        openai: true,
        langchain: true,
        redis: true,
        rateLimiter: true,
      });
    });

    it('should handle service failures', async () => {
      mockOpenAIService.testConnection.mockResolvedValue(false);
      mockLangChainService.testConnection.mockResolvedValue(true);
      mockRedis.ping.mockRejectedValue(new Error('Connection failed'));

      const health = await service.getHealthStatus();

      expect(health).toEqual({
        openai: false,
        langchain: true,
        redis: false,
        rateLimiter: true,
      });
    });
  });

  describe('getUsageStats', () => {
    it('should return combined usage statistics', async () => {
      const stats = await service.getUsageStats();

      expect(stats).toEqual({
        openai: {
          requestCount: 5,
          totalTokensUsed: 1000,
          estimatedCost: 0.02,
        },
        totalContextsSaved: 10,
        cacheHitRate: 0.0,
      });

      expect(mockOpenAIService.getUsageStats).toHaveBeenCalled();
      expect(mockRepository.count).toHaveBeenCalled();
    });
  });

  describe('rate limiting', () => {
    it('should enforce rate limits', async () => {
      // Mock rate limiter to throw on second call
      const mockRateLimiter = {
        consume: jest.fn()
          .mockResolvedValueOnce(undefined)
          .mockRejectedValueOnce(new Error('Rate limit exceeded')),
      };

      (service as any).rateLimiter = mockRateLimiter;

      const mockTask = {
        title: 'Task',
        description: 'Description',
        priority: 'medium',
        complexity: 'minor',
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);
      mockRedis.get.mockResolvedValue(null);

      // First call should succeed
      await service.createTaskFromNaturalLanguage({ prompt: 'First task' });

      // Second call should fail with rate limit
      await expect(service.createTaskFromNaturalLanguage({ 
        prompt: 'Second task',
      })).rejects.toThrow('Rate limit exceeded');
    });
  });

  describe('caching', () => {
    it('should handle cache errors gracefully', async () => {
      mockRedis.get.mockRejectedValue(new Error('Redis error'));
      mockRedis.set.mockRejectedValue(new Error('Redis error'));

      const mockTask = {
        title: 'Task',
        description: 'Description',
        priority: 'medium',
        complexity: 'minor',
      };

      mockLangChainService.createTaskFromNaturalLanguage.mockResolvedValue(mockTask);

      // Should still work despite cache errors
      const result = await service.createTaskFromNaturalLanguage({
        prompt: 'Create task',
      });

      expect(result).toEqual(mockTask);
    });

    it('should invalidate cache patterns', async () => {
      mockRedis.keys.mockResolvedValue([
        'ai:breakdown:task-123',
        'ai:breakdown:task-123:v2',
      ]);

      await service.updateTaskContext('task-123', 'Feedback');

      expect(mockRedis.del).toHaveBeenCalledWith(
        'ai:breakdown:task-123',
        'ai:breakdown:task-123:v2',
      );
    });
  });
});