query-orchestration/api/queryAPI/documents_test.go

package queryapi_test

import (
	"bytes"
	"mime/multipart"
	"net/http"
	"path/filepath"
	"testing"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/test"

	queryapi "queryorchestration/api/queryAPI"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestUploadDocument(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	test.CreateBucket(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)

	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)

	fileWriter, err := writer.CreateFormFile("file", filepath.Base("../../assets/original/helloWorld.pdf"))
	require.NoError(t, err)

	_, err = fileWriter.Write([]byte("hi"))
	require.NoError(t, err)

	err = writer.Close()
	require.NoError(t, err)

	ctx, rec := createContextWithBody(t, buf.Bytes())

	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	err = cons.UploadDocument(ctx, "client_id")
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
}

func TestListDocumentsByClientId(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "client_id",
		Hash:     "hash",
		BatchID:  nil,
	})
	require.NoError(t, err)

	ctx, rec := createContext(t)

	err = cons.ListDocumentsByClientId(ctx, "client_id")
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.ListDocuments{
		queryapi.DocumentSummary{
			Id:   docId,
			Hash: "hash",
		},
	})
}

func TestGetDocument(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "client_id",
		Hash:     "hash",
		BatchID:  nil,
	})
	require.NoError(t, err)

	ctx, rec := createContext(t)

	err = cons.GetDocument(ctx, docId)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.Document{
		Id:       docId,
		ClientId: "client_id",
		Hash:     "hash",
		Fields:   map[string]interface{}{},
	})
}

func TestListDocumentBatches(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_list"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch List",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test listing batches with default parameters
	ctx, rec := createContext(t)
	err = cons.ListDocumentBatches(ctx, clientID, queryapi.ListDocumentBatchesParams{})
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)

	response := getBody[queryapi.BatchUploadList](t, rec)
	assert.Len(t, response.Batches, 1)
	assert.Equal(t, batchID.String(), response.Batches[0].BatchId.String())

	// Test listing batches with custom limit and offset
	limit := int32(5)
	offset := int32(0)
	ctx2, rec2 := createContext(t)
	err = cons.ListDocumentBatches(ctx2, clientID, queryapi.ListDocumentBatchesParams{
		Limit:  &limit,
		Offset: &offset,
	})
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec2.Code)

	response2 := getBody[queryapi.BatchUploadList](t, rec2)
	assert.Len(t, response2.Batches, 1)
	assert.Equal(t, batchID.String(), response2.Batches[0].BatchId.String())
}

func TestUploadDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_upload"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Upload",
	})
	require.NoError(t, err)

	// Create multipart form with ZIP file
	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)
	fileWriter, err := writer.CreateFormFile("archive", "test.zip")
	require.NoError(t, err)
	_, err = fileWriter.Write([]byte("PK")) // Minimal ZIP file header
	require.NoError(t, err)
	err = writer.Close()
	require.NoError(t, err)

	ctx, rec := createContextWithBody(t, buf.Bytes())
	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	// Test batch upload
	err = cons.UploadDocumentBatch(ctx, clientID)
	require.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, rec.Code)

	response := getBody[queryapi.BatchUploadResponse](t, rec)
	assert.NotEmpty(t, response.BatchId)
	assert.Equal(t, queryapi.BatchStatusProcessing, response.Status)
}

func TestGetDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_get"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Get",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test getting batch details
	ctx, rec := createContext(t)
	err = cons.GetDocumentBatch(ctx, clientID, queryapi.BatchID(batchID))
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)

	response := getBody[queryapi.BatchUploadDetails](t, rec)
	assert.Equal(t, batchID.String(), response.BatchId.String())
	assert.Equal(t, "test.zip", response.OriginalFilename)
}

func TestCancelDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_cancel"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Cancel",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test canceling batch
	ctx, rec := createContext(t)
	err = cons.CancelDocumentBatch(ctx, clientID, queryapi.BatchID(batchID))
	require.NoError(t, err)
	assert.Equal(t, http.StatusNoContent, rec.Code)

	// Verify batch was canceled
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusCancelled, batch.Status)
}