src/test/java/com/dalab/autoarchival/service/impl/ArchivalTaskServiceImplTest.java

package com.dalab.autoarchival.service.impl;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.UUID;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;

import com.dalab.autoarchival.dto.ArchivalConfigDTO;
import com.dalab.autoarchival.dto.ArchivalTaskListResponse;
import com.dalab.autoarchival.dto.ArchivalTaskRequest;
import com.dalab.autoarchival.dto.ArchivalTaskResponse;
import com.dalab.autoarchival.dto.ArchivalTaskStatusDTO;
import com.dalab.autoarchival.dto.RestoreRequestDTO;
import com.dalab.autoarchival.dto.RestoreResponseDTO;
import com.dalab.autoarchival.dto.TaskApprovalRequest;
import com.dalab.autoarchival.mapper.ArchivalTaskMapper;
import com.dalab.autoarchival.model.ArchivalStatus;
import com.dalab.autoarchival.model.ArchivalTaskEntity;
import com.dalab.autoarchival.repository.ArchivalTaskRepository;
import com.dalab.autoarchival.service.IArchivalConfigService;

@ExtendWith(MockitoExtension.class)
class ArchivalTaskServiceImplTest {

    @Mock
    private ArchivalTaskRepository taskRepository;

    @Mock
    private ArchivalTaskMapper taskMapper;

    @Mock
    private IArchivalConfigService archivalConfigService; // JPA version

    @InjectMocks
    private ArchivalTaskServiceImpl taskService;

    private ArchivalTaskRequest sampleTaskRequest;
    private ArchivalTaskEntity sampleTaskEntity;
    private ArchivalTaskStatusDTO sampleTaskStatusDTO;
    private ArchivalTaskResponse sampleTaskResponse;
    private String sampleTaskId;
    private ArchivalConfigDTO globalConfigNoApproval;
    private ArchivalConfigDTO globalConfigWithApproval;

    @BeforeEach
    void setUp() {
        sampleTaskId = UUID.randomUUID().toString();
        ArchivalTaskRequest.ArchivalScope scope = ArchivalTaskRequest.ArchivalScope.builder()
                .assetIds(List.of("asset1")).build();

        sampleTaskRequest = ArchivalTaskRequest.builder()
                .taskName("Test Task")
                .scope(scope)
                .triggeredBy("user@test.com")
                .build();

        sampleTaskEntity = new ArchivalTaskEntity();
        sampleTaskEntity.setTaskId(sampleTaskId);
        sampleTaskEntity.setTaskName("Test Task");
        sampleTaskEntity.setStatus(ArchivalStatus.SUBMITTED);
        sampleTaskEntity.setSubmittedAt(LocalDateTime.now());
        sampleTaskEntity.setScope(new ArchivalTaskEntity.ArchivalScopeData(List.of("asset1")));

        sampleTaskStatusDTO = ArchivalTaskStatusDTO.builder().taskId(sampleTaskId).status("SUBMITTED").build();
        sampleTaskResponse = ArchivalTaskResponse.builder().taskId(sampleTaskId).status("SUBMITTED").build();
        
        globalConfigNoApproval = ArchivalConfigDTO.builder().requireApprovalForArchival(false).build();
        globalConfigWithApproval = ArchivalConfigDTO.builder().requireApprovalForArchival(true).build();
    }

    @Test
    void submitArchivalTask_ValidRequest_NoApprovalNeeded_ShouldSaveAndReturnResponse() {
        when(archivalConfigService.getArchivalConfig()).thenReturn(globalConfigNoApproval);
        when(taskMapper.requestToEntity(sampleTaskRequest)).thenReturn(sampleTaskEntity);
        when(taskRepository.save(any(ArchivalTaskEntity.class))).thenReturn(sampleTaskEntity);
        when(taskMapper.entityToTaskResponse(sampleTaskEntity)).thenReturn(sampleTaskResponse);

        ArchivalTaskResponse response = taskService.submitArchivalTask(sampleTaskRequest);

        assertNotNull(response);
        assertEquals("SUBMITTED", response.getStatus());
        ArgumentCaptor<ArchivalTaskEntity> entityCaptor = ArgumentCaptor.forClass(ArchivalTaskEntity.class);
        verify(taskRepository).save(entityCaptor.capture());
        assertEquals(ArchivalStatus.SUBMITTED, entityCaptor.getValue().getStatus());
        assertNotNull(entityCaptor.getValue().getAssetStatuses());
        assertEquals(1, entityCaptor.getValue().getAssetsPendingArchival());
    }
    
    @Test
    void submitArchivalTask_ValidRequest_ApprovalNeeded_ShouldSetPendingApproval() {
        when(archivalConfigService.getArchivalConfig()).thenReturn(globalConfigWithApproval);
        when(taskMapper.requestToEntity(sampleTaskRequest)).thenReturn(sampleTaskEntity);
        when(taskRepository.save(any(ArchivalTaskEntity.class))).thenReturn(sampleTaskEntity);
        when(taskMapper.entityToTaskResponse(sampleTaskEntity)).thenReturn(sampleTaskResponse);

        taskService.submitArchivalTask(sampleTaskRequest);

        ArgumentCaptor<ArchivalTaskEntity> entityCaptor = ArgumentCaptor.forClass(ArchivalTaskEntity.class);
        verify(taskRepository).save(entityCaptor.capture());
        assertEquals(ArchivalStatus.PENDING_APPROVAL, entityCaptor.getValue().getStatus());
    }

    @Test
    void submitArchivalTask_InvalidScope_ShouldReturnFailedValidation() {
        ArchivalTaskRequest invalidRequest = ArchivalTaskRequest.builder().scope(null).build();
        ArchivalTaskResponse response = taskService.submitArchivalTask(invalidRequest);
        assertEquals("FAILED_VALIDATION", response.getStatus());
        verify(taskRepository, never()).save(any());
    }

    @Test
    void getTaskStatus_TaskExists_ShouldReturnDTO() {
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));
        when(taskMapper.entityToStatusDTO(sampleTaskEntity)).thenReturn(sampleTaskStatusDTO);

        ArchivalTaskStatusDTO result = taskService.getTaskStatus(sampleTaskId);

        assertNotNull(result);
        assertEquals(sampleTaskId, result.getTaskId());
    }

    @Test
    void getTaskStatus_TaskNotExists_ShouldReturnNull() {
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.empty());
        ArchivalTaskStatusDTO result = taskService.getTaskStatus(sampleTaskId);
        assertNull(result);
    }

    @Test
    void listTasks_ShouldReturnPaginatedDTOs() {
        Pageable pageable = PageRequest.of(0, 10);
        Page<ArchivalTaskEntity> page = new PageImpl<>(Collections.singletonList(sampleTaskEntity), pageable, 1);
        when(taskRepository.findAll(pageable)).thenReturn(page);
        when(taskMapper.entityToStatusDTO(sampleTaskEntity)).thenReturn(sampleTaskStatusDTO);

        ArchivalTaskListResponse response = taskService.listTasks(pageable);

        assertNotNull(response);
        assertEquals(1, response.getTasks().size());
        assertEquals(sampleTaskId, response.getTasks().get(0).getTaskId());
    }

    @Test
    void approveTask_TaskExistsAndPending_ShouldApprove() {
        sampleTaskEntity.setStatus(ArchivalStatus.PENDING_APPROVAL);
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));
        when(taskRepository.save(any(ArchivalTaskEntity.class))).thenReturn(sampleTaskEntity);
        when(taskMapper.entityToStatusDTO(sampleTaskEntity)).thenAnswer(inv -> {
            ArchivalTaskEntity e = inv.getArgument(0);
            return ArchivalTaskStatusDTO.builder().taskId(e.getTaskId()).status(e.getStatus().name()).build();
        });

        TaskApprovalRequest approvalReq = TaskApprovalRequest.builder().comments("Approved").build();
        ArchivalTaskStatusDTO result = taskService.approveTask(sampleTaskId, approvalReq);

        assertEquals("APPROVED", result.getStatus());
        assertEquals("Approved", sampleTaskEntity.getApprovalComments());
        verify(taskRepository).save(sampleTaskEntity);
    }

    @Test
    void approveTask_TaskNotPending_ShouldReturnCurrentStatusWithErrors() {
        sampleTaskEntity.setStatus(ArchivalStatus.COMPLETED);
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));
        when(taskMapper.entityToStatusDTO(sampleTaskEntity)).thenReturn(ArchivalTaskStatusDTO.builder().status("COMPLETED").build());

        ArchivalTaskStatusDTO result = taskService.approveTask(sampleTaskId, new TaskApprovalRequest());

        assertEquals("COMPLETED", result.getStatus());
        assertNotNull(result.getErrorMessages());
        assertFalse(result.getErrorMessages().isEmpty());
        verify(taskRepository, never()).save(any());
    }
    
    @Test
    void approveTask_TaskNotFound_ShouldThrowException() {
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.empty());
        assertThrows(ArchivalTaskNotFoundException.class, () -> {
            taskService.approveTask(sampleTaskId, new TaskApprovalRequest());
        });
    }

    // Similar tests for rejectTask and requestRestore
    @Test
    void rejectTask_TaskExistsAndPending_ShouldReject() {
        sampleTaskEntity.setStatus(ArchivalStatus.PENDING_APPROVAL);
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));
        when(taskRepository.save(any(ArchivalTaskEntity.class))).thenReturn(sampleTaskEntity);
         when(taskMapper.entityToStatusDTO(sampleTaskEntity)).thenAnswer(inv -> {
            ArchivalTaskEntity e = inv.getArgument(0);
            return ArchivalTaskStatusDTO.builder().taskId(e.getTaskId()).status(e.getStatus().name()).build();
        });

        TaskApprovalRequest rejectionReq = TaskApprovalRequest.builder().comments("Rejected").build();
        ArchivalTaskStatusDTO result = taskService.rejectTask(sampleTaskId, rejectionReq);

        assertEquals("REJECTED", result.getStatus());
        assertEquals("Rejected", sampleTaskEntity.getRejectionComments());
        verify(taskRepository).save(sampleTaskEntity);
    }

    @Test
    void requestRestore_TaskCompleted_ShouldInitiateRestore() {
        sampleTaskEntity.setStatus(ArchivalStatus.COMPLETED);
        RestoreRequestDTO restoreDto = RestoreRequestDTO.builder().restoreTier("STANDARD").retentionDays(1).build();
        ArchivalTaskEntity.RestoreRequestData restoreData = new ArchivalTaskEntity.RestoreRequestData("STANDARD", 1, null, null, null);
        
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));
        when(taskMapper.restoreRequestDtoToData(restoreDto)).thenReturn(restoreData);
        when(taskRepository.save(any(ArchivalTaskEntity.class))).thenReturn(sampleTaskEntity);

        RestoreResponseDTO response = taskService.requestRestore(sampleTaskId, restoreDto);

        assertEquals(ArchivalStatus.RESTORATION_REQUESTED.name(), response.getStatus());
        assertNotNull(sampleTaskEntity.getRestoreRequestDetails());
        verify(taskRepository).save(sampleTaskEntity);
    }
    
    @Test
    void requestRestore_TaskNotRestorable_ShouldReturnNotRestorable() {
        sampleTaskEntity.setStatus(ArchivalStatus.IN_PROGRESS);
        RestoreRequestDTO restoreDto = RestoreRequestDTO.builder().build();
        when(taskRepository.findById(sampleTaskId)).thenReturn(Optional.of(sampleTaskEntity));

        RestoreResponseDTO response = taskService.requestRestore(sampleTaskId, restoreDto);
        assertEquals("NOT_RESTORABLE", response.getStatus());
        verify(taskRepository, never()).save(any());
    }
}