package queryapi_test import ( "bytes" "encoding/json" "fmt" "net/http" "net/http/httptest" "testing" queryapi "queryorchestration/api/queryAPI" "queryorchestration/internal/database/repository" "queryorchestration/internal/label" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/google/uuid" "github.com/labstack/echo/v4" openapi_types "github.com/oapi-codegen/runtime/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type LabelsTestConfig struct { serviceconfig.BaseConfig } func (c *LabelsTestConfig) GetBackgroundRunner() any { return nil } func setupLabelsTestController(t *testing.T) (*queryapi.Controllers, *LabelsTestConfig, string, uuid.UUID) { t.Helper() cfg := &LabelsTestConfig{} test.CreateDB(t, cfg) // Create services labelSvc := label.New(cfg) services := &queryapi.Services{ Label: labelSvc, } // Create test client with unique ID ctx := t.Context() uniqueID := uuid.New().String()[:8] clientID := fmt.Sprintf("test-client-%s", uniqueID) err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: fmt.Sprintf("Test Client %s", uniqueID), }) require.NoError(t, err) // Create a test document filename := "test.pdf" documentID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("labeltesthash-%s", uniqueID), Filename: &filename, }) require.NoError(t, err) ctrl := queryapi.NewControllers(services, nil) return ctrl, cfg, clientID, documentID } func TestApplyLabel_Success(t *testing.T) { ctrl, _, _, documentID := setupLabelsTestController(t) // Use a label that exists in the labels lookup table (seeded in migration) reqBody := queryapi.LabelApplication{ Label: "Ingested", AppliedBy: openapi_types.Email("user@example.com"), } body, err := json.Marshal(reqBody) require.NoError(t, err) e := echo.New() req := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/documents/%s/labels", documentID), bytes.NewReader(body)) req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err = ctrl.ApplyLabel(c, openapi_types.UUID(documentID)) require.NoError(t, err) assert.Equal(t, http.StatusCreated, rec.Code) var resp queryapi.LabelRecord err = json.Unmarshal(rec.Body.Bytes(), &resp) require.NoError(t, err) assert.Equal(t, documentID, uuid.UUID(resp.DocumentId)) assert.Equal(t, "Ingested", resp.Label) assert.Equal(t, "user@example.com", string(resp.AppliedBy)) assert.NotEqual(t, uuid.Nil, uuid.UUID(resp.Id)) } func TestApplyLabel_MissingLabel(t *testing.T) { ctrl, _, _, documentID := setupLabelsTestController(t) reqBody := queryapi.LabelApplication{ Label: "", AppliedBy: openapi_types.Email("user@example.com"), } body, err := json.Marshal(reqBody) require.NoError(t, err) e := echo.New() req := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/documents/%s/labels", documentID), bytes.NewReader(body)) req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err = ctrl.ApplyLabel(c, openapi_types.UUID(documentID)) require.Error(t, err) httpErr, ok := err.(*echo.HTTPError) require.True(t, ok) assert.Equal(t, http.StatusBadRequest, httpErr.Code) } func TestApplyLabel_MissingAppliedBy(t *testing.T) { ctrl, _, _, documentID := setupLabelsTestController(t) // Send raw JSON with empty appliedBy field body := []byte(`{"label":"Ingested","appliedBy":""}`) e := echo.New() req := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/documents/%s/labels", documentID), bytes.NewReader(body)) req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err := ctrl.ApplyLabel(c, openapi_types.UUID(documentID)) require.Error(t, err) httpErr, ok := err.(*echo.HTTPError) require.True(t, ok) assert.Equal(t, http.StatusBadRequest, httpErr.Code) } func TestApplyLabel_SameLabelMultipleTimes(t *testing.T) { ctrl, _, _, documentID := setupLabelsTestController(t) // Use a label that exists in the labels lookup table (seeded in migration) reqBody := queryapi.LabelApplication{ Label: "Ingested", AppliedBy: openapi_types.Email("user@example.com"), } body, err := json.Marshal(reqBody) require.NoError(t, err) // Apply label first time e := echo.New() req := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/documents/%s/labels", documentID), bytes.NewReader(body)) req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err = ctrl.ApplyLabel(c, openapi_types.UUID(documentID)) require.NoError(t, err) assert.Equal(t, http.StatusCreated, rec.Code) var resp1 queryapi.LabelRecord _ = json.Unmarshal(rec.Body.Bytes(), &resp1) // Apply same label second time - should create another record for history req = httptest.NewRequest(http.MethodPost, fmt.Sprintf("/documents/%s/labels", documentID), bytes.NewReader(body)) req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) rec = httptest.NewRecorder() c = e.NewContext(req, rec) err = ctrl.ApplyLabel(c, openapi_types.UUID(documentID)) require.NoError(t, err) assert.Equal(t, http.StatusCreated, rec.Code) var resp2 queryapi.LabelRecord _ = json.Unmarshal(rec.Body.Bytes(), &resp2) // Both should have the same label but different record IDs assert.Equal(t, "Ingested", resp1.Label) assert.Equal(t, "Ingested", resp2.Label) assert.NotEqual(t, resp1.Id, resp2.Id) } func TestGetDocumentLabels_Success(t *testing.T) { ctrl, cfg, _, documentID := setupLabelsTestController(t) ctx := t.Context() // Apply multiple labels to the document (using labels from the lookup table) _, err := cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID, Label: "Ingested", Appliedby: "user1@example.com", }) require.NoError(t, err) _, err = cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID, Label: "OCR_Processed", Appliedby: "user2@example.com", }) require.NoError(t, err) _, err = cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID, Label: "GenAI_Processed", Appliedby: "user3@example.com", }) require.NoError(t, err) e := echo.New() req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/documents/%s/labels", documentID), nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err = ctrl.GetDocumentLabels(c, openapi_types.UUID(documentID)) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) var resp struct { Labels []queryapi.LabelRecord `json:"labels"` } err = json.Unmarshal(rec.Body.Bytes(), &resp) require.NoError(t, err) assert.Len(t, resp.Labels, 3) // Verify all labels are present labelNames := make(map[string]bool) for _, lbl := range resp.Labels { labelNames[lbl.Label] = true } assert.True(t, labelNames["Ingested"]) assert.True(t, labelNames["OCR_Processed"]) assert.True(t, labelNames["GenAI_Processed"]) } func TestGetDocumentLabels_NoLabels(t *testing.T) { ctrl, _, _, documentID := setupLabelsTestController(t) e := echo.New() req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/documents/%s/labels", documentID), nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) err := ctrl.GetDocumentLabels(c, openapi_types.UUID(documentID)) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) var resp struct { Labels []queryapi.LabelRecord `json:"labels"` } err = json.Unmarshal(rec.Body.Bytes(), &resp) require.NoError(t, err) assert.Len(t, resp.Labels, 0) } func TestGetDocumentsByLabel_Success(t *testing.T) { cfg := &LabelsTestConfig{} test.CreateDB(t, cfg) // Create services labelSvc := label.New(cfg) services := &queryapi.Services{ Label: labelSvc, } // Create test client ctx := t.Context() uniqueID := uuid.New().String()[:8] clientID := fmt.Sprintf("test-client-%s", uniqueID) // Use normal string clientId err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: fmt.Sprintf("Test Client %s", clientID[:8]), }) require.NoError(t, err) // Create test documents filename1 := "doc1.pdf" documentID1, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("labelhash1-%s", uuid.New().String()[:8]), Filename: &filename1, }) require.NoError(t, err) filename2 := "doc2.pdf" documentID2, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("labelhash2-%s", uuid.New().String()[:8]), Filename: &filename2, }) require.NoError(t, err) // Create a third document filename3 := "doc3.pdf" documentID3, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("labelhash3-%s", uuid.New().String()[:8]), Filename: &filename3, }) require.NoError(t, err) // Apply "Ingested" label to doc1 and doc2, but not doc3 (using labels from lookup table) _, err = cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID1, Label: "Ingested", Appliedby: "user@example.com", }) require.NoError(t, err) _, err = cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID2, Label: "Ingested", Appliedby: "user@example.com", }) require.NoError(t, err) // Apply different label to doc3 _, err = cfg.GetDBQueries().ApplyLabel(ctx, &repository.ApplyLabelParams{ Documentid: documentID3, Label: "OCR_Processed", Appliedby: "user@example.com", }) require.NoError(t, err) ctrl := queryapi.NewControllers(services, nil) // Get documents with "Ingested" label e := echo.New() req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/labels/Ingested/documents?clientId=%s", clientID), nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) params := queryapi.GetDocumentsByLabelParams{ ClientId: queryapi.ClientID(clientID), } err = ctrl.GetDocumentsByLabel(c, "Ingested", params) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) var resp struct { Documents []queryapi.DocumentSummary `json:"documents"` } err = json.Unmarshal(rec.Body.Bytes(), &resp) require.NoError(t, err) assert.Len(t, resp.Documents, 2) // Verify the correct documents are returned docIDs := make(map[uuid.UUID]bool) for _, doc := range resp.Documents { docIDs[uuid.UUID(doc.Id)] = true } assert.True(t, docIDs[documentID1]) assert.True(t, docIDs[documentID2]) assert.False(t, docIDs[documentID3]) } func TestGetDocumentsByLabel_NoDocuments(t *testing.T) { cfg := &LabelsTestConfig{} test.CreateDB(t, cfg) // Create services labelSvc := label.New(cfg) services := &queryapi.Services{ Label: labelSvc, } // Create test client ctx := t.Context() uniqueID := uuid.New().String()[:8] clientID := fmt.Sprintf("test-client-%s", uniqueID) err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: fmt.Sprintf("Test Client %s", uniqueID), }) require.NoError(t, err) ctrl := queryapi.NewControllers(services, nil) e := echo.New() req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/labels/nonexistent/documents?clientId=%s", clientID), nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) params := queryapi.GetDocumentsByLabelParams{ ClientId: queryapi.ClientID(clientID), } err = ctrl.GetDocumentsByLabel(c, "nonexistent", params) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) var resp struct { Documents []queryapi.DocumentSummary `json:"documents"` } err = json.Unmarshal(rec.Body.Bytes(), &resp) require.NoError(t, err) assert.Len(t, resp.Documents, 0) } func TestGetDocumentsByLabel_MissingLabelName(t *testing.T) { cfg := &LabelsTestConfig{} test.CreateDB(t, cfg) // Create services labelSvc := label.New(cfg) services := &queryapi.Services{ Label: labelSvc, } // Create test client ctx := t.Context() uniqueID := uuid.New().String()[:8] clientID := fmt.Sprintf("test-client-%s", uniqueID) err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: fmt.Sprintf("Test Client %s", uniqueID), }) require.NoError(t, err) ctrl := queryapi.NewControllers(services, nil) e := echo.New() req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/labels//documents?clientId=%s", clientID), nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) params := queryapi.GetDocumentsByLabelParams{ ClientId: queryapi.ClientID(clientID), } err = ctrl.GetDocumentsByLabel(c, "", params) require.Error(t, err) httpErr, ok := err.(*echo.HTTPError) require.True(t, ok) assert.Equal(t, http.StatusBadRequest, httpErr.Code) }