package documentupload_test import ( "io" "log" "strings" "testing" "queryorchestration/internal/client" "queryorchestration/internal/database/repository" documentupload "queryorchestration/internal/document/upload" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/test" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type Config struct { serviceconfig.BaseConfig objectstore.ObjectStoreConfig } // createTestClient creates a client using the client service, which automatically creates the root folder. func createTestClient(t *testing.T, cfg *Config, clientID, clientName string) { t.Helper() clientSvc := client.New(cfg) _, err := clientSvc.Create(t.Context(), client.CreateParams{ ID: clientID, Name: clientName, }) require.NoError(t, err) } func TestUpload(t *testing.T) { t.Parallel() cfg := &Config{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) // Use client service to create client (which auto-creates root folder) createTestClient(t, cfg, "client_id", "client") svc := documentupload.New(cfg) file := documentupload.File{ ClientID: "client_id", Content: strings.NewReader("abc"), Filename: "test.txt", } log.Print("test") err := svc.Upload(t.Context(), file) require.NoError(t, err) os, err := cfg.StoreClient.ListObjects(t.Context(), &s3.ListObjectsInput{ Bucket: &cfg.Bucket, }) require.NoError(t, err) assert.Len(t, os.Contents, 1) key := *os.Contents[0].Key o, err := cfg.StoreClient.GetObject(t.Context(), &s3.GetObjectInput{ Bucket: &cfg.Bucket, Key: &key, }) require.NoError(t, err) body, err := io.ReadAll(o.Body) require.NoError(t, err) assert.Equal(t, []byte("abc"), body) } func TestUploadWithPath(t *testing.T) { t.Parallel() cfg := &Config{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) // Use client service to create client (which auto-creates root folder) createTestClient(t, cfg, "client_id", "client") svc := documentupload.New(cfg) // Upload with a full path including folders file := documentupload.File{ ClientID: "client_id", Content: strings.NewReader("content with path"), Filename: "agreement.pdf", Path: "contracts/2025/Q1/agreement.pdf", } result, err := svc.UploadWithResult(t.Context(), file) require.NoError(t, err) // Verify original path is set correctly assert.Equal(t, "contracts/2025/Q1/agreement.pdf", result.OriginalPath) // Verify folder was assigned (always non-nil now since we always have at least root folder) assert.NotEqual(t, result.FolderID.String(), "00000000-0000-0000-0000-000000000000") // Verify S3 key is set assert.NotEmpty(t, result.Key.String()) // Verify folder hierarchy was created in database folder, err := cfg.GetDBQueries().GetFolderByID(t.Context(), result.FolderID) require.NoError(t, err) assert.Equal(t, "/contracts/2025/Q1", folder.Path) // Verify parent folders exist (root "/" + /contracts, /contracts/2025, /contracts/2025/Q1) folders, err := cfg.GetDBQueries().GetFoldersByClientID(t.Context(), "client_id") require.NoError(t, err) assert.Len(t, folders, 4) // /, /contracts, /contracts/2025, /contracts/2025/Q1 // Verify the root folder exists and is the parent of /contracts rootFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{ Clientid: "client_id", Path: "/", }) require.NoError(t, err) assert.Nil(t, rootFolder.Parentid, "Root folder should have no parent") // Verify /contracts has root as parent contractsFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{ Clientid: "client_id", Path: "/contracts", }) require.NoError(t, err) assert.Equal(t, rootFolder.ID, *contractsFolder.Parentid, "/contracts should have root as parent") // Verify S3 object exists os, err := cfg.StoreClient.ListObjects(t.Context(), &s3.ListObjectsInput{ Bucket: &cfg.Bucket, }) require.NoError(t, err) assert.Len(t, os.Contents, 1) } func TestUploadWithPathNoFolder(t *testing.T) { t.Parallel() cfg := &Config{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) // Use client service to create client (which auto-creates root folder) createTestClient(t, cfg, "client_id", "client") svc := documentupload.New(cfg) // Upload with just a filename (no folder path) file := documentupload.File{ ClientID: "client_id", Content: strings.NewReader("root content"), Filename: "root_file.pdf", } result, err := svc.UploadWithResult(t.Context(), file) require.NoError(t, err) // Verify original path is just the filename assert.Equal(t, "root_file.pdf", result.OriginalPath) // Verify document is assigned to root folder rootFolder, err := cfg.GetDBQueries().GetFolderByPath(t.Context(), &repository.GetFolderByPathParams{ Clientid: "client_id", Path: "/", }) require.NoError(t, err) assert.Equal(t, rootFolder.ID, result.FolderID, "Document without path should be in root folder") // Verify only root folder exists in database folders, err := cfg.GetDBQueries().GetFoldersByClientID(t.Context(), "client_id") require.NoError(t, err) assert.Len(t, folders, 1) // Only root folder "/" assert.Equal(t, "/", folders[0].Path) } func TestUploadAndReturnKey(t *testing.T) { t.Parallel() cfg := &Config{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) // Use client service to create client (which auto-creates root folder) createTestClient(t, cfg, "client_id", "client") svc := documentupload.New(cfg) file := documentupload.File{ ClientID: "client_id", Content: strings.NewReader("key test content"), Filename: "keytest.pdf", Path: "folder/keytest.pdf", } key, err := svc.UploadAndReturnKey(t.Context(), file) require.NoError(t, err) // Verify the key is valid assert.Equal(t, "client_id", key.ClientID) assert.NotEmpty(t, key.String()) } // TestUpload_SetsFileTypeOnKey verifies that UploadWithResult sets the FileType // field on the BucketKey based on the filename extension. func TestUpload_SetsFileTypeOnKey(t *testing.T) { t.Parallel() cfg := &Config{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) createTestClient(t, cfg, "client_id", "client") svc := documentupload.New(cfg) tests := []struct { name string filename string wantFileType *string }{ { name: "pdf extension", filename: "test.pdf", wantFileType: strPtr("pdf"), }, { name: "txt extension", filename: "notes.txt", wantFileType: strPtr("txt"), }, { name: "docx extension", filename: "report.docx", wantFileType: strPtr("docx"), }, { name: "no extension", filename: "noext", wantFileType: nil, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { file := documentupload.File{ ClientID: "client_id", Content: strings.NewReader("file content"), Filename: tt.filename, } result, err := svc.UploadWithResult(t.Context(), file) require.NoError(t, err) if tt.wantFileType == nil { assert.Nil(t, result.Key.FileType, "FileType should be nil for filename without extension") } else { require.NotNil(t, result.Key.FileType, "FileType should not be nil for filename with extension") assert.Equal(t, *tt.wantFileType, *result.Key.FileType) } }) } } func strPtr(s string) *string { return &s }