stack/apps/api/src/job-steps/job-steps.service.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { JobStepsService } from "./job-steps.service";
import { PrismaService } from "../prisma/prisma.service";
import { JobStepPhase, JobStepType, JobStepStatus } from "@prisma/client";
import { NotFoundException } from "@nestjs/common";
import { CreateStepDto, UpdateStepDto } from "./dto";

describe("JobStepsService", () => {
  let service: JobStepsService;
  let prisma: PrismaService;

  const mockPrismaService = {
    jobStep: {
      create: vi.fn(),
      findMany: vi.fn(),
      findUnique: vi.fn(),
      update: vi.fn(),
    },
    runnerJob: {
      findUnique: vi.fn(),
    },
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        JobStepsService,
        {
          provide: PrismaService,
          useValue: mockPrismaService,
        },
      ],
    }).compile();

    service = module.get<JobStepsService>(JobStepsService);
    prisma = module.get<PrismaService>(PrismaService);

    // Clear all mocks before each test
    vi.clearAllMocks();
  });

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

  describe("create", () => {
    it("should create a job step", async () => {
      const jobId = "job-123";
      const createDto: CreateStepDto = {
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repository",
        type: JobStepType.COMMAND,
      };

      const mockStep = {
        id: "step-123",
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repository",
        type: JobStepType.COMMAND,
        status: JobStepStatus.PENDING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: null,
        completedAt: null,
        durationMs: null,
      };

      mockPrismaService.jobStep.create.mockResolvedValue(mockStep);

      const result = await service.create(jobId, createDto);

      expect(result).toEqual(mockStep);
      expect(prisma.jobStep.create).toHaveBeenCalledWith({
        data: {
          job: { connect: { id: jobId } },
          ordinal: 1,
          phase: JobStepPhase.SETUP,
          name: "Clone repository",
          type: JobStepType.COMMAND,
          status: JobStepStatus.PENDING,
        },
      });
    });

    it("should use provided status when creating step", async () => {
      const jobId = "job-123";
      const createDto: CreateStepDto = {
        ordinal: 2,
        phase: JobStepPhase.EXECUTION,
        name: "Run tests",
        type: JobStepType.COMMAND,
        status: JobStepStatus.RUNNING,
      };

      const mockStep = {
        id: "step-124",
        jobId,
        ordinal: 2,
        phase: JobStepPhase.EXECUTION,
        name: "Run tests",
        type: JobStepType.COMMAND,
        status: JobStepStatus.RUNNING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: new Date(),
        completedAt: null,
        durationMs: null,
      };

      mockPrismaService.jobStep.create.mockResolvedValue(mockStep);

      const result = await service.create(jobId, createDto);

      expect(result).toEqual(mockStep);
      expect(prisma.jobStep.create).toHaveBeenCalledWith({
        data: {
          job: { connect: { id: jobId } },
          ordinal: 2,
          phase: JobStepPhase.EXECUTION,
          name: "Run tests",
          type: JobStepType.COMMAND,
          status: JobStepStatus.RUNNING,
        },
      });
    });
  });

  describe("findAllByJob", () => {
    it("should return all steps for a job ordered by ordinal", async () => {
      const jobId = "job-123";

      const mockSteps = [
        {
          id: "step-1",
          jobId,
          ordinal: 1,
          phase: JobStepPhase.SETUP,
          name: "Clone repo",
          type: JobStepType.COMMAND,
          status: JobStepStatus.COMPLETED,
          output: "Cloned successfully",
          tokensInput: null,
          tokensOutput: null,
          startedAt: new Date("2024-01-01T10:00:00Z"),
          completedAt: new Date("2024-01-01T10:00:05Z"),
          durationMs: 5000,
        },
        {
          id: "step-2",
          jobId,
          ordinal: 2,
          phase: JobStepPhase.EXECUTION,
          name: "Run tests",
          type: JobStepType.COMMAND,
          status: JobStepStatus.RUNNING,
          output: null,
          tokensInput: null,
          tokensOutput: null,
          startedAt: new Date("2024-01-01T10:00:05Z"),
          completedAt: null,
          durationMs: null,
        },
      ];

      mockPrismaService.jobStep.findMany.mockResolvedValue(mockSteps);

      const result = await service.findAllByJob(jobId);

      expect(result).toEqual(mockSteps);
      expect(prisma.jobStep.findMany).toHaveBeenCalledWith({
        where: { jobId },
        orderBy: { ordinal: "asc" },
      });
    });
  });

  describe("findOne", () => {
    it("should return a single step by ID", async () => {
      const stepId = "step-123";
      const jobId = "job-123";

      const mockStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repo",
        type: JobStepType.COMMAND,
        status: JobStepStatus.COMPLETED,
        output: "Cloned successfully",
        tokensInput: null,
        tokensOutput: null,
        startedAt: new Date("2024-01-01T10:00:00Z"),
        completedAt: new Date("2024-01-01T10:00:05Z"),
        durationMs: 5000,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(mockStep);

      const result = await service.findOne(stepId, jobId);

      expect(result).toEqual(mockStep);
      expect(prisma.jobStep.findUnique).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
      });
    });

    it("should throw NotFoundException when step not found", async () => {
      const stepId = "step-999";
      const jobId = "job-123";

      mockPrismaService.jobStep.findUnique.mockResolvedValue(null);

      await expect(service.findOne(stepId, jobId)).rejects.toThrow(NotFoundException);
      await expect(service.findOne(stepId, jobId)).rejects.toThrow(
        `JobStep with ID ${stepId} not found`
      );
    });
  });

  describe("update", () => {
    it("should update step status", async () => {
      const stepId = "step-123";
      const jobId = "job-123";
      const updateDto: UpdateStepDto = {
        status: JobStepStatus.COMPLETED,
      };

      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repo",
        type: JobStepType.COMMAND,
        status: JobStepStatus.RUNNING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: new Date("2024-01-01T10:00:00Z"),
        completedAt: null,
        durationMs: null,
      };

      const updatedStep = {
        ...existingStep,
        status: JobStepStatus.COMPLETED,
        completedAt: new Date("2024-01-01T10:00:05Z"),
        durationMs: 5000,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(updatedStep);

      const result = await service.update(stepId, jobId, updateDto);

      expect(result).toEqual(updatedStep);
      expect(prisma.jobStep.update).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
        data: { status: JobStepStatus.COMPLETED },
      });
    });

    it("should update step with output and token usage", async () => {
      const stepId = "step-123";
      const jobId = "job-123";
      const updateDto: UpdateStepDto = {
        status: JobStepStatus.COMPLETED,
        output: "Analysis complete",
        tokensInput: 1000,
        tokensOutput: 500,
      };

      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 2,
        phase: JobStepPhase.EXECUTION,
        name: "AI Analysis",
        type: JobStepType.AI_ACTION,
        status: JobStepStatus.RUNNING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: new Date("2024-01-01T10:00:00Z"),
        completedAt: null,
        durationMs: null,
      };

      const updatedStep = {
        ...existingStep,
        status: JobStepStatus.COMPLETED,
        output: "Analysis complete",
        tokensInput: 1000,
        tokensOutput: 500,
        completedAt: new Date("2024-01-01T10:00:10Z"),
        durationMs: 10000,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(updatedStep);

      const result = await service.update(stepId, jobId, updateDto);

      expect(result).toEqual(updatedStep);
      expect(prisma.jobStep.update).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
        data: {
          status: JobStepStatus.COMPLETED,
          output: "Analysis complete",
          tokensInput: 1000,
          tokensOutput: 500,
        },
      });
    });

    it("should throw NotFoundException when step not found", async () => {
      const stepId = "step-999";
      const jobId = "job-123";
      const updateDto: UpdateStepDto = {
        status: JobStepStatus.COMPLETED,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(null);

      await expect(service.update(stepId, jobId, updateDto)).rejects.toThrow(NotFoundException);
    });
  });

  describe("startStep", () => {
    it("should mark step as running and set startedAt", async () => {
      const stepId = "step-123";
      const jobId = "job-123";

      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repo",
        type: JobStepType.COMMAND,
        status: JobStepStatus.PENDING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: null,
        completedAt: null,
        durationMs: null,
      };

      const startedStep = {
        ...existingStep,
        status: JobStepStatus.RUNNING,
        startedAt: new Date("2024-01-01T10:00:00Z"),
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(startedStep);

      const result = await service.startStep(stepId, jobId);

      expect(result).toEqual(startedStep);
      expect(prisma.jobStep.update).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
        data: {
          status: JobStepStatus.RUNNING,
          startedAt: expect.any(Date),
        },
      });
    });
  });

  describe("completeStep", () => {
    it("should mark step as completed and calculate duration", async () => {
      const stepId = "step-123";
      const jobId = "job-123";

      const startTime = new Date("2024-01-01T10:00:00Z");
      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repo",
        type: JobStepType.COMMAND,
        status: JobStepStatus.RUNNING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: startTime,
        completedAt: null,
        durationMs: null,
      };

      const completedStep = {
        ...existingStep,
        status: JobStepStatus.COMPLETED,
        output: "Success",
        completedAt: new Date("2024-01-01T10:00:05Z"),
        durationMs: 5000,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(completedStep);

      const result = await service.completeStep(stepId, jobId, "Success");

      expect(result).toEqual(completedStep);
      expect(prisma.jobStep.update).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
        data: {
          status: JobStepStatus.COMPLETED,
          output: "Success",
          completedAt: expect.any(Date),
          durationMs: expect.any(Number),
        },
      });
    });

    it("should handle step without startedAt by setting durationMs to null", async () => {
      const stepId = "step-123";
      const jobId = "job-123";

      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.SETUP,
        name: "Clone repo",
        type: JobStepType.COMMAND,
        status: JobStepStatus.PENDING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: null,
        completedAt: null,
        durationMs: null,
      };

      const completedStep = {
        ...existingStep,
        status: JobStepStatus.COMPLETED,
        output: "Success",
        completedAt: new Date("2024-01-01T10:00:05Z"),
        durationMs: null,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(completedStep);

      const result = await service.completeStep(stepId, jobId, "Success");

      expect(result.durationMs).toBeNull();
    });
  });

  describe("failStep", () => {
    it("should mark step as failed with error output", async () => {
      const stepId = "step-123";
      const jobId = "job-123";
      const error = "Command failed with exit code 1";

      const startTime = new Date("2024-01-01T10:00:00Z");
      const existingStep = {
        id: stepId,
        jobId,
        ordinal: 1,
        phase: JobStepPhase.VALIDATION,
        name: "Run tests",
        type: JobStepType.GATE,
        status: JobStepStatus.RUNNING,
        output: null,
        tokensInput: null,
        tokensOutput: null,
        startedAt: startTime,
        completedAt: null,
        durationMs: null,
      };

      const failedStep = {
        ...existingStep,
        status: JobStepStatus.FAILED,
        output: error,
        completedAt: new Date("2024-01-01T10:00:03Z"),
        durationMs: 3000,
      };

      mockPrismaService.jobStep.findUnique.mockResolvedValue(existingStep);
      mockPrismaService.jobStep.update.mockResolvedValue(failedStep);

      const result = await service.failStep(stepId, jobId, error);

      expect(result).toEqual(failedStep);
      expect(prisma.jobStep.update).toHaveBeenCalledWith({
        where: { id: stepId, jobId },
        data: {
          status: JobStepStatus.FAILED,
          output: error,
          completedAt: expect.any(Date),
          durationMs: expect.any(Number),
        },
      });
    });
  });
});