query-orchestration/api/queryAPI/labels_test.go

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)
}