query-orchestration/internal/document/download/download_test.go

package documentdownload_test

import (
	"bytes"
	"io"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"queryorchestration/internal/client"
	"queryorchestration/internal/database/repository"
	documentdownload "queryorchestration/internal/document/download"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/test"
)

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

// setupTestDocument creates a document in DB, uploads content to S3, and creates the document entry.
// Returns the document ID.
func setupTestDocument(t *testing.T, cfg *Config, clientID, hash string, filename *string, content []byte) uuid.UUID {
	t.Helper()

	docID, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     hash,
		Filename: filename,
	})
	require.NoError(t, err)

	bucket := cfg.GetBucket()
	s3Key := "test/download/" + docID.String()

	_, err = cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{
		Bucket: &bucket,
		Key:    &s3Key,
		Body:   bytes.NewReader(content),
	})
	require.NoError(t, err)

	err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
		Documentid: docID,
		Bucket:     bucket,
		Key:        s3Key,
	})
	require.NoError(t, err)

	return docID
}

func TestDownloadDocument_Success(t *testing.T) {
	t.Parallel()
	cfg := &Config{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	createTestClient(t, cfg, "dl_success_client", "Download Success Client")

	filename := "test-report.pdf"
	fileContent := []byte("test file content for download verification")
	docID := setupTestDocument(t, cfg, "dl_success_client", "dl_success_hash", &filename, fileContent)

	svc := documentdownload.New(cfg)
	result, err := svc.GetDownload(t.Context(), docID)
	require.NoError(t, err)
	defer result.Body.Close()

	body, err := io.ReadAll(result.Body)
	require.NoError(t, err)

	assert.Equal(t, fileContent, body)
	assert.Equal(t, "test-report.pdf", result.Filename)
	assert.Equal(t, "application/pdf", result.ContentType)
}

func TestDownloadDocument_NotFound(t *testing.T) {
	t.Parallel()
	cfg := &Config{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	svc := documentdownload.New(cfg)
	_, err := svc.GetDownload(t.Context(), uuid.New())
	require.Error(t, err)
	assert.Contains(t, err.Error(), "document not found")
}

func TestDownloadDocument_NoS3Entry(t *testing.T) {
	t.Parallel()
	cfg := &Config{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	createTestClient(t, cfg, "dl_noentry_client", "Download NoEntry Client")

	// Create document but do NOT create a document entry or upload to S3
	docID, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "dl_noentry_client",
		Hash:     "dl_noentry_hash",
	})
	require.NoError(t, err)

	svc := documentdownload.New(cfg)
	_, err = svc.GetDownload(t.Context(), docID)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "no S3 entry")
}

func TestDownloadDocument_MimeTypeDetection(t *testing.T) {
	t.Parallel()
	cfg := &Config{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	createTestClient(t, cfg, "dl_mime_client", "Download Mime Client")

	tests := []struct {
		name         string
		filename     *string
		expectedMime string
	}{
		{
			name:         "PDF file",
			filename:     strPtr("report.pdf"),
			expectedMime: "application/pdf",
		},
		{
			name:         "PNG image",
			filename:     strPtr("photo.png"),
			expectedMime: "image/png",
		},
		{
			name:         "CSV file",
			filename:     strPtr("data.csv"),
			expectedMime: "text/csv; charset=utf-8",
		},
		{
			name:         "nil filename defaults to octet-stream",
			filename:     nil,
			expectedMime: "application/octet-stream",
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			hash := "dl_mime_hash_" + uuid.New().String()[:8]
			content := []byte("file content " + tc.name)

			docID := setupTestDocument(t, cfg, "dl_mime_client", hash, tc.filename, content)

			svc := documentdownload.New(cfg)
			result, err := svc.GetDownload(t.Context(), docID)
			require.NoError(t, err)
			defer result.Body.Close()

			body, err := io.ReadAll(result.Body)
			require.NoError(t, err)

			assert.Equal(t, tc.expectedMime, result.ContentType)
			assert.Equal(t, content, body)
		})
	}
}

func strPtr(s string) *string {
	return &s
}