src/test/java/com/dalab/autolabel/controller/LabelingJobControllerTest.java

package com.dalab.autolabel.controller;

import com.dalab.autolabel.client.rest.dto.*;
import com.dalab.autolabel.exception.AutoLabelingException;
import com.dalab.autolabel.service.ILabelingJobService;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.autoconfigure.web.servlet.WebMvcTest;
import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.http.MediaType;
import org.springframework.security.test.context.support.WithMockUser;
import org.springframework.test.web.servlet.MockMvc;

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

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.when;
import static org.springframework.security.test.web.servlet.request.SecurityMockMvcRequestPostProcessors.csrf;
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post;
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;

@WebMvcTest(LabelingJobController.class)
class LabelingJobControllerTest {

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private ILabelingJobService labelingJobService;

    @Autowired
    private ObjectMapper objectMapper;

    private LabelingJobRequest labelingJobRequest;
    private LabelingJobResponse labelingJobResponse;
    private LabelingJobStatusResponse labelingJobStatusResponse;
    private LabelingFeedbackRequest labelingFeedbackRequest;
    private LabelingFeedbackResponse labelingFeedbackResponse;

    @BeforeEach
    void setUp() {
        String jobId = UUID.randomUUID().toString();
        labelingJobRequest = LabelingJobRequest.builder().jobName("Test Job").build();
        labelingJobResponse = LabelingJobResponse.builder()
                .jobId(jobId)
                .status("SUBMITTED")
                .submittedAt(LocalDateTime.now())
                .message("Job submitted")
                .build();
        labelingJobStatusResponse = LabelingJobStatusResponse.builder()
                .jobId(jobId).jobName("Test Job").status("RUNNING").build();

        labelingFeedbackRequest = LabelingFeedbackRequest.builder()
                .assetId("asset-123")
                .labelingJobId(jobId)
                .feedbackItems(List.of(LabelingFeedbackRequest.FeedbackItem.builder()
                        .suggestedLabel("PII")
                        .type(LabelingFeedbackRequest.FeedbackType.CONFIRMED)
                        .build()))
                .build();
        labelingFeedbackResponse = LabelingFeedbackResponse.builder()
                .feedbackId(UUID.randomUUID().toString())
                .assetId("asset-123")
                .status("RECEIVED")
                .build();
    }

    @Test
    @WithMockUser(authorities = "ROLE_DATA_STEWARD")
    void submitLabelingJob_ValidRequest_ShouldSucceed() throws Exception {
        when(labelingJobService.submitLabelingJob(any(LabelingJobRequest.class))).thenReturn(labelingJobResponse);

        mockMvc.perform(post("/api/v1/labeling/jobs")
                        .with(csrf())
                        .contentType(MediaType.APPLICATION_JSON)
                        .content(objectMapper.writeValueAsString(labelingJobRequest)))
                .andExpect(status().isAccepted())
                .andExpect(jsonPath("$.jobId").value(labelingJobResponse.getJobId()));
    }
    
    @Test
    @WithMockUser(authorities = "ROLE_USER") // Insufficient role
    void submitLabelingJob_UserRole_ShouldBeForbidden() throws Exception {
        mockMvc.perform(post("/api/v1/labeling/jobs")
                        .with(csrf())
                        .contentType(MediaType.APPLICATION_JSON)
                        .content(objectMapper.writeValueAsString(labelingJobRequest)))
                .andExpect(status().isForbidden());
    }

    @Test
    @WithMockUser(authorities = "ROLE_DATA_STEWARD")
    void getJobStatus_ExistingJob_ShouldReturnStatus() throws Exception {
        when(labelingJobService.getJobStatus(eq(labelingJobStatusResponse.getJobId()))).thenReturn(labelingJobStatusResponse);

        mockMvc.perform(get("/api/v1/labeling/jobs/{jobId}", labelingJobStatusResponse.getJobId()))
                .andExpect(status().isOk())
                .andExpect(jsonPath("$.jobId").value(labelingJobStatusResponse.getJobId()))
                .andExpect(jsonPath("$.status").value("RUNNING"));
    }

    @Test
    @WithMockUser(authorities = "ROLE_DATA_STEWARD")
    void getJobStatus_NonExistingJob_ShouldReturnNotFound() throws Exception {
        String nonExistentJobId = UUID.randomUUID().toString();
        when(labelingJobService.getJobStatus(eq(nonExistentJobId))).thenReturn(null);

        mockMvc.perform(get("/api/v1/labeling/jobs/{jobId}", nonExistentJobId))
                .andExpect(status().isNotFound());
    }

    @Test
    @WithMockUser(authorities = "ROLE_USER")
    void listJobs_ShouldReturnJobList() throws Exception {
        LabelingJobListResponse listResponse = LabelingJobListResponse.builder()
                .jobs(Collections.singletonList(labelingJobStatusResponse))
                .pageNumber(0).pageSize(20).totalElements(1L).totalPages(1).last(true)
                .build();
        when(labelingJobService.listJobs(any(Pageable.class))).thenReturn(listResponse);

        mockMvc.perform(get("/api/v1/labeling/jobs").param("page", "0").param("size", "20"))
                .andExpect(status().isOk())
                .andExpect(jsonPath("$.jobs[0].jobId").value(labelingJobStatusResponse.getJobId()));
    }

    @Test
    @WithMockUser(authorities = "ROLE_DATA_STEWARD")
    void submitLabelingFeedback_ValidRequest_ShouldSucceed() throws Exception {
        when(labelingJobService.processLabelingFeedback(any(LabelingFeedbackRequest.class))).thenReturn(labelingFeedbackResponse);

        mockMvc.perform(post("/api/v1/labeling/feedback")
                        .with(csrf())
                        .contentType(MediaType.APPLICATION_JSON)
                        .content(objectMapper.writeValueAsString(labelingFeedbackRequest)))
                .andExpect(status().isOk())
                .andExpect(jsonPath("$.feedbackId").value(labelingFeedbackResponse.getFeedbackId()));
    }
    
    @Test
    @WithMockUser(authorities = "ROLE_USER") // Insufficient role
    void submitLabelingFeedback_UserRole_ShouldBeForbidden() throws Exception {
         mockMvc.perform(post("/api/v1/labeling/feedback")
                        .with(csrf())
                        .contentType(MediaType.APPLICATION_JSON)
                        .content(objectMapper.writeValueAsString(labelingFeedbackRequest)))
                .andExpect(status().isForbidden());
    }

    @Test
    @WithMockUser(authorities = "ROLE_DATA_STEWARD")
    void submitLabelingJob_ServiceThrowsAutoLabelingException_ShouldReturnBadRequest() throws Exception {
        when(labelingJobService.submitLabelingJob(any(LabelingJobRequest.class)))
            .thenThrow(new AutoLabelingException("Test ML config error"));

        mockMvc.perform(post("/api/v1/labeling/jobs")
                        .with(csrf())
                        .contentType(MediaType.APPLICATION_JSON)
                        .content(objectMapper.writeValueAsString(labelingJobRequest)))
                .andExpect(status().isBadRequest())
                .andExpect(jsonPath("$.message").value("Test ML config error"));
    }
}