edr-platform/apps/edr-freight-api/src/modules/train-scheduling/train-scheduling.service.spec.ts

import { ConflictException } from '@nestjs/common';

import { TrainSchedulingService } from './train-scheduling.service';

const nw5 = {
  id: 'wagon-type-1',
  code: 'NW5',
  name: 'Flat Wagon',
  capacityTons: 70,
  lengthMeters: 14,
  maxWagonsPerTrain: 53,
  supportedLoadTypes: ['CONTAINER'],
  isActive: true,
};

const locomotive = {
  id: 'loc-1',
  code: 'LOC-001',
  maxPullWeightTons: 3500,
  status: 'AVAILABLE',
};

const makeBooking = (
  id: string,
  reference: string,
  weight: number,
  quantity: number,
  containerCode: string,
  scheduledDate = '2026-06-20T08:00:00.000Z',
  originYardId = 'yard-origin',
  destinationYardId = 'yard-destination',
) => ({
  id,
  reference,
  freightType: 'CONTAINER',
  cargoTotalWeightVgm: weight,
  scheduledDate: new Date(scheduledDate),
  originYardId,
  destinationYardId,
  status: 'PAID',
  customer: { companyName: 'Demo Customer' },
  originYard: { label: 'Djibouti', code: 'DJIBOUTI' },
  destinationYard: { label: 'Addis Ababa', code: 'ADDIS_ABABA' },
  bookingContainers: [
    {
      quantity,
      containerType: { code: containerCode, label: containerCode },
    },
  ],
});

describe('TrainSchedulingService', () => {
  let service: TrainSchedulingService;
  let dataSource: {
    getRepository: jest.Mock;
    transaction: jest.Mock;
  };
  let locomotivesRepository: {
    findById: jest.Mock;
  };
  let wagonTypesRepository: {
    findAll: jest.Mock;
  };

  beforeEach(() => {
    dataSource = {
      getRepository: jest.fn(),
      transaction: jest.fn(),
    };
    locomotivesRepository = {
      findById: jest.fn(),
    };
    wagonTypesRepository = {
      findAll: jest.fn(),
    };

    service = new TrainSchedulingService(
      dataSource as never,
      locomotivesRepository as never,
      wagonTypesRepository as never,
    );
  });

  it('computes the expected valid preview for Group A', async () => {
    const bookings = [
      makeBooking('b1', 'BKG-CONT-001', 500, 20, '40FT'),
      makeBooking('b2', 'BKG-CONT-002', 300, 10, '20FT'),
      makeBooking('b3', 'BKG-CONT-003', 450, 15, '40FT'),
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    dataSource.getRepository.mockImplementation((entity: { name?: string }) => {
      if (entity?.name === 'Booking') {
        return { find: jest.fn().mockResolvedValue(bookings) };
      }
      if (entity?.name === 'TrainScheduleBooking') {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity?.name === 'Locomotive') {
        return {
          count: jest.fn().mockResolvedValue(2),
          find: jest.fn().mockResolvedValue([locomotive]),
        };
      }
      throw new Error(`Unexpected repository ${entity?.name}`);
    });

    const result = await service.previewContainerTrainSchedule({
      bookingIds: bookings.map((booking) => booking.id),
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(true);
    expect(result.violations).toEqual([]);
    expect(result.summary).toEqual({
      totalBookings: 3,
      totalWeightTons: 1250,
      wagonType: 'NW5',
      wagonsNeeded: 18,
      totalLengthMeters: 252,
    });
    expect(result.wagonPlan).toHaveLength(18);
    expect(result.wagonPlan[0]?.allocations[0]).toEqual({
      bookingId: 'b1',
      bookingReference: 'BKG-CONT-001',
      allocatedWeightTons: 70,
    });
  });

  it('flags the overweight booking as invalid', async () => {
    const bookings = [makeBooking('b6', 'BKG-CONT-006', 3600, 80, '40FT')];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    dataSource.getRepository.mockImplementation((entity: { name?: string }) => {
      if (entity?.name === 'Booking') {
        return { find: jest.fn().mockResolvedValue(bookings) };
      }
      if (entity?.name === 'TrainScheduleBooking') {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity?.name === 'Locomotive') {
        return {
          count: jest.fn().mockResolvedValue(1),
          find: jest.fn().mockResolvedValue([locomotive]),
        };
      }
      throw new Error(`Unexpected repository ${entity?.name}`);
    });

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b6'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(false);
    expect(result.summary.totalWeightTons).toBe(3600);
    expect(result.violations).toContain(
      'Total booking weight 3600T exceeds max train weight 3500T',
    );
  });

  it('rejects bookings that are not in schedulable status', async () => {
    const bookings = [
      {
        ...makeBooking('b7', 'BKG-CONT-007', 120, 2, '40FT'),
        status: 'APPROVED',
      },
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    dataSource.getRepository.mockImplementation((entity: { name?: string }) => {
      if (entity?.name === 'Booking') {
        return { find: jest.fn().mockResolvedValue(bookings) };
      }
      if (entity?.name === 'TrainScheduleBooking') {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity?.name === 'Locomotive') {
        return {
          count: jest.fn().mockResolvedValue(1),
          find: jest.fn().mockResolvedValue([locomotive]),
        };
      }
      throw new Error(`Unexpected repository ${entity?.name}`);
    });

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b7'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(false);
    expect(result.violations).toContain(
      'Only PAID bookings can be scheduled; received: APPROVED',
    );
  });

  it('creates a schedule transactionally when validation passes', async () => {
    const bookings = [makeBooking('b1', 'BKG-CONT-001', 140, 2, '40FT')];
    const validation = {
      valid: true,
      violations: [],
      bookings,
      wagonType: nw5,
      summary: {
        totalBookings: 1,
        totalWeightTons: 140,
        wagonType: 'NW5',
        wagonsNeeded: 2,
        totalLengthMeters: 28,
      },
      wagonPlan: [
        {
          sequenceNo: 1,
          capacityTons: 70,
          lengthMeters: 14,
          assignedWeightTons: 70,
          allocations: [
            {
              bookingId: 'b1',
              bookingReference: 'BKG-CONT-001',
              allocatedWeightTons: 70,
            },
          ],
        },
        {
          sequenceNo: 2,
          capacityTons: 70,
          lengthMeters: 14,
          assignedWeightTons: 70,
          allocations: [
            {
              bookingId: 'b1',
              bookingReference: 'BKG-CONT-001',
              allocatedWeightTons: 70,
            },
          ],
        },
      ],
    };

    const lockedLocomotiveRepo = {
      findOne: jest.fn().mockResolvedValue(locomotive),
      update: jest.fn().mockResolvedValue(undefined),
    };
    const trainScheduleRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue({ id: 'schedule-1' }),
    };
    const trainScheduleBookingRepo = {
      count: jest.fn().mockResolvedValue(0),
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue(undefined),
    };
    const trainSetWagonRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue(undefined),
      find: jest.fn().mockResolvedValue([
        { id: 'wagon-1', sequenceNo: 1 },
        { id: 'wagon-2', sequenceNo: 2 },
      ]),
    };
    const wagonAllocRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue(undefined),
    };
    const trainSetRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue({ id: 'train-set-1' }),
    };
    const manager = {
      getRepository: jest.fn((entity: { name?: string }) => {
        switch (entity?.name) {
          case 'Locomotive':
            return lockedLocomotiveRepo;
          case 'TrainSchedule':
            return trainScheduleRepo;
          case 'TrainScheduleBooking':
            return trainScheduleBookingRepo;
          case 'TrainSetWagon':
            return trainSetWagonRepo;
          case 'WagonBookingAllocation':
            return wagonAllocRepo;
          case 'TrainSet':
            return trainSetRepo;
          default:
            throw new Error(`Unexpected transaction repository ${entity?.name}`);
        }
      }),
    };

    jest.spyOn(service, 'validateContainerBookingsForScheduling').mockResolvedValue(validation as never);
    jest.spyOn(service, 'selectOrValidateLocomotive').mockResolvedValue(locomotive as never);
    jest.spyOn(service, 'getContainerTrainScheduleById').mockResolvedValue({ id: 'schedule-1' } as never);
    dataSource.transaction.mockImplementation(async (callback: (tx: typeof manager) => Promise<string>) =>
      callback(manager),
    );

    const result = await service.createContainerTrainSchedule({
      bookingIds: ['b1'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
      locomotiveId: 'loc-1',
    });

    expect(trainSetRepo.save).toHaveBeenCalled();
    expect(trainScheduleRepo.save).toHaveBeenCalled();
    expect(trainSetWagonRepo.save).toHaveBeenCalled();
    expect(wagonAllocRepo.save).toHaveBeenCalled();
    expect(lockedLocomotiveRepo.update).toHaveBeenCalledWith('loc-1', { status: 'ASSIGNED' });
    expect(result).toEqual({ id: 'schedule-1' });
  });

  it('rejects create when the locked locomotive is no longer available', async () => {
    const validation = {
      valid: true,
      violations: [],
      bookings: [makeBooking('b1', 'BKG-CONT-001', 70, 1, '40FT')],
      wagonType: nw5,
      summary: {
        totalBookings: 1,
        totalWeightTons: 70,
        wagonType: 'NW5',
        wagonsNeeded: 1,
        totalLengthMeters: 14,
      },
      wagonPlan: [
        {
          sequenceNo: 1,
          capacityTons: 70,
          lengthMeters: 14,
          assignedWeightTons: 70,
          allocations: [],
        },
      ],
    };
    const manager = {
      getRepository: jest.fn(() => ({
        findOne: jest.fn().mockResolvedValue({ ...locomotive, status: 'ASSIGNED' }),
      })),
    };

    jest.spyOn(service, 'validateContainerBookingsForScheduling').mockResolvedValue(validation as never);
    jest.spyOn(service, 'selectOrValidateLocomotive').mockResolvedValue(locomotive as never);
    dataSource.transaction.mockImplementation(async (callback: (tx: typeof manager) => Promise<string>) =>
      callback(manager),
    );

    await expect(
      service.createContainerTrainSchedule({
        bookingIds: ['b1'],
        scheduleDate: '2026-06-20T08:00:00.000Z',
        originStationId: 'yard-origin',
        destinationStationId: 'yard-destination',
        locomotiveId: 'loc-1',
      }),
    ).rejects.toBeInstanceOf(ConflictException);
  });
});