query-orchestration/internal/document/upload/get_test.go

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
}