jay brown
476 lines
14 KiB
Go
476 lines
14 KiB
Go
package repository_test
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
|
|
"queryorchestration/internal/database/repository"
|
|
"queryorchestration/internal/serviceconfig"
|
|
"queryorchestration/internal/test"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestBatchUpload(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
t.Run("CreateBatchUpload", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client first
|
|
clientID := fmt.Sprintf("TEST_CLIENT_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
params := &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test_batch.zip",
|
|
TotalDocuments: 10,
|
|
}
|
|
|
|
batchID, err := queries.CreateBatchUpload(ctx, params)
|
|
require.NoError(t, err)
|
|
assert.NotEqual(t, uuid.Nil, batchID)
|
|
})
|
|
|
|
t.Run("GetBatchUpload", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_GET_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Get %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test_get.zip",
|
|
TotalDocuments: 5,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Get batch upload
|
|
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, batchID, batch.ID)
|
|
assert.Equal(t, clientID, batch.ClientID)
|
|
assert.Equal(t, "test_get.zip", batch.OriginalFilename)
|
|
assert.Equal(t, int32(5), batch.TotalDocuments)
|
|
assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
|
|
assert.Equal(t, int32(0), batch.ProcessedDocuments)
|
|
assert.Equal(t, int32(0), batch.FailedDocuments)
|
|
assert.NotNil(t, batch.CreatedAt)
|
|
assert.False(t, batch.CompletedAt.Valid) // pgtype.Timestamp has Valid field
|
|
})
|
|
|
|
t.Run("ListBatchUploads", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_LIST_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client List %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create multiple batch uploads
|
|
for i := 0; i < 3; i++ {
|
|
totalDocs := int32(1) + int32(i) //nolint:gosec // safe conversion for test data
|
|
_, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: fmt.Sprintf("batch%d.zip", i),
|
|
TotalDocuments: totalDocs,
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// List batch uploads
|
|
batches, err := queries.ListBatchUploads(ctx, &repository.ListBatchUploadsParams{
|
|
ClientID: clientID,
|
|
Limit: 10,
|
|
Offset: 0,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Len(t, batches, 3)
|
|
})
|
|
|
|
t.Run("UpdateBatchProgress", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_PROGRESS_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Progress %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "progress.zip",
|
|
TotalDocuments: 100,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Update progress
|
|
err = queries.UpdateBatchProgress(ctx, &repository.UpdateBatchProgressParams{
|
|
ID: batchID,
|
|
ProcessedDocuments: 42,
|
|
FailedDocuments: 2,
|
|
InvalidTypeDocuments: 1,
|
|
ProgressPercent: 45,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Verify update
|
|
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int32(42), batch.ProcessedDocuments)
|
|
assert.Equal(t, int32(2), batch.FailedDocuments)
|
|
assert.Equal(t, int32(1), batch.InvalidTypeDocuments)
|
|
assert.Equal(t, int32(45), batch.ProgressPercent)
|
|
})
|
|
|
|
t.Run("UpdateBatchStatus", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_STATUS_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Status %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "status.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Update status to completed
|
|
err = queries.UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
|
|
ID: batchID,
|
|
Column2: repository.BatchStatusCompleted,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Verify update
|
|
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusCompleted, batch.Status)
|
|
assert.True(t, batch.CompletedAt.Valid)
|
|
assert.True(t, batch.CompletedAt.Time.After(batch.CreatedAt.Time))
|
|
})
|
|
|
|
t.Run("AddFailedFilename", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_FAILED_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Failed %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "failed.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Add failed filenames
|
|
err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{
|
|
ID: batchID,
|
|
Column2: "document1.pdf",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{
|
|
ID: batchID,
|
|
Column2: "document2.pdf",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Verify failed filenames were added
|
|
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Verify failed filenames were added as JSON array
|
|
assert.NotNil(t, batch.FailedFilenames)
|
|
// The failed filenames should contain our test filenames
|
|
failedFilenamesStr := string(batch.FailedFilenames)
|
|
assert.Contains(t, failedFilenamesStr, "document1.pdf")
|
|
assert.Contains(t, failedFilenamesStr, "document2.pdf")
|
|
})
|
|
|
|
t.Run("CancelBatchUpload", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_CANCEL_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Cancel %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "cancel.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Cancel batch upload
|
|
err = queries.CancelBatchUpload(ctx, &repository.CancelBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Verify cancellation
|
|
batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusCancelled, batch.Status)
|
|
assert.True(t, batch.CompletedAt.Valid)
|
|
})
|
|
|
|
t.Run("GetDocumentsByBatchId", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_DOCS_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Docs %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "docs.zip",
|
|
TotalDocuments: 3,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create documents with batch ID
|
|
var docIDs []uuid.UUID
|
|
for i := 0; i < 3; i++ {
|
|
docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: fmt.Sprintf("hash%d", i),
|
|
BatchID: &batchID,
|
|
})
|
|
require.NoError(t, err)
|
|
docIDs = append(docIDs, docID)
|
|
}
|
|
|
|
// Get documents by batch ID
|
|
docs, err := queries.GetDocumentsByBatchId(ctx, &batchID)
|
|
require.NoError(t, err)
|
|
assert.Len(t, docs, 3)
|
|
|
|
// Verify all documents are returned
|
|
for _, doc := range docs {
|
|
assert.Contains(t, docIDs, doc.ID)
|
|
assert.Equal(t, clientID, doc.Clientid)
|
|
}
|
|
})
|
|
|
|
t.Run("CountDocumentsByBatchId", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := t.Context()
|
|
|
|
cfg := &serviceconfig.BaseConfig{}
|
|
test.CreateDB(t, cfg)
|
|
queries := cfg.GetDBQueries()
|
|
|
|
// Create a client
|
|
clientID := fmt.Sprintf("TEST_CLIENT_COUNT_%s", uuid.New().String()[:8])
|
|
err := queries.CreateClient(ctx, &repository.CreateClientParams{
|
|
Name: fmt.Sprintf("Test Client Count %s", uuid.New().String()[:8]),
|
|
Clientid: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create batch upload
|
|
batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "count.zip",
|
|
TotalDocuments: 5,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create documents with batch ID
|
|
for i := 0; i < 4; i++ {
|
|
_, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: fmt.Sprintf("hash%d", i),
|
|
BatchID: &batchID,
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Count documents by batch ID
|
|
count, err := queries.CountDocumentsByBatchId(ctx, &batchID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(4), count)
|
|
})
|
|
}
|
|
|
|
func TestBatchStatusEnum(t *testing.T) {
|
|
t.Run("Scan method", func(t *testing.T) {
|
|
// Test valid scan with string
|
|
var status repository.BatchStatus
|
|
err := status.Scan("processing")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusProcessing, status)
|
|
|
|
// Test valid scan with []byte
|
|
var statusBytes repository.BatchStatus
|
|
err = statusBytes.Scan([]byte("completed"))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusCompleted, statusBytes)
|
|
|
|
// Test all enum values
|
|
testCases := []struct {
|
|
str string
|
|
expected repository.BatchStatus
|
|
}{
|
|
{"processing", repository.BatchStatusProcessing},
|
|
{"completed", repository.BatchStatusCompleted},
|
|
{"failed", repository.BatchStatusFailed},
|
|
{"cancelled", repository.BatchStatusCancelled}, //nolint:misspell // matches database enum value
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
var s repository.BatchStatus
|
|
err := s.Scan(tc.str)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, tc.expected, s)
|
|
}
|
|
|
|
// Test scan with unsupported type
|
|
var invalidStatus repository.BatchStatus
|
|
err = invalidStatus.Scan(123) // int is not supported
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "unsupported scan type")
|
|
})
|
|
|
|
t.Run("NullBatchStatus", func(t *testing.T) {
|
|
// Test scan with nil
|
|
var nullStatus repository.NullBatchStatus
|
|
err := nullStatus.Scan(nil)
|
|
require.NoError(t, err)
|
|
assert.False(t, nullStatus.Valid)
|
|
assert.Equal(t, repository.BatchStatus(""), nullStatus.BatchStatus)
|
|
|
|
// Test scan with valid value
|
|
err = nullStatus.Scan("processing")
|
|
require.NoError(t, err)
|
|
assert.True(t, nullStatus.Valid)
|
|
assert.Equal(t, repository.BatchStatusProcessing, nullStatus.BatchStatus)
|
|
|
|
// Test Value method when Valid is false
|
|
nullStatus.Valid = false // reset to false
|
|
val, err := nullStatus.Value()
|
|
require.NoError(t, err)
|
|
assert.Nil(t, val)
|
|
|
|
// Test Value method when Valid is true
|
|
nullStatus.Valid = true
|
|
nullStatus.BatchStatus = repository.BatchStatusCompleted
|
|
val, err = nullStatus.Value()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "completed", val)
|
|
})
|
|
|
|
t.Run("BatchStatus Valid method", func(t *testing.T) {
|
|
// Test valid enum values
|
|
assert.True(t, repository.BatchStatusProcessing.Valid())
|
|
assert.True(t, repository.BatchStatusCompleted.Valid())
|
|
assert.True(t, repository.BatchStatusFailed.Valid())
|
|
assert.True(t, repository.BatchStatusCancelled.Valid()) //nolint:misspell // matches database enum value
|
|
|
|
// Test invalid enum value
|
|
invalidStatus := repository.BatchStatus("invalid")
|
|
assert.False(t, invalidStatus.Valid())
|
|
})
|
|
}
|