yao/neo/attachment/manager_test.go

package attachment

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"mime/multipart"
	"strings"
	"testing"
)

func TestManagerUpload(t *testing.T) {
	// Create a local storage manager
	manager, err := New(ManagerOption{
		Driver:       "local",
		MaxSize:      "10M",
		ChunkSize:    "2M",
		AllowedTypes: []string{"text/*", "image/*", ".txt", ".jpg", ".png"},
		Options: map[string]interface{}{
			"path": "/tmp/test_attachments",
		},
	})
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	// Test simple text file upload
	t.Run("SimpleTextUpload", func(t *testing.T) {
		content := "Hello, World!"
		reader := strings.NewReader(content)

		// Create a mock file header
		fileHeader := &FileHeader{
			FileHeader: &multipart.FileHeader{
				Filename: "test.txt",
				Size:     int64(len(content)),
				Header:   make(map[string][]string),
			},
		}
		fileHeader.Header.Set("Content-Type", "text/plain")

		option := UploadOption{
			UserID:      "user123",
			ChatID:      "chat456",
			AssistantID: "assistant789",
		}

		file, err := manager.Upload(context.Background(), fileHeader, reader, option)
		if err != nil {
			t.Fatalf("Failed to upload file: %v", err)
		}

		if file.Filename != "test.txt" {
			t.Errorf("Expected filename 'test.txt', got '%s'", file.Filename)
		}

		if file.ContentType != "text/plain" {
			t.Errorf("Expected content type 'text/plain', got '%s'", file.ContentType)
		}

		// Test download
		response, err := manager.Download(context.Background(), file.ID)
		if err != nil {
			t.Fatalf("Failed to download file: %v", err)
		}
		defer response.Reader.Close()

		downloadedContent, err := manager.Read(context.Background(), file.ID)
		if err != nil {
			t.Fatalf("Failed to read file: %v", err)
		}

		if string(downloadedContent) != content {
			t.Errorf("Expected content '%s', got '%s'", content, string(downloadedContent))
		}

		// Test ReadBase64
		base64Content, err := manager.ReadBase64(context.Background(), file.ID)
		if err != nil {
			t.Fatalf("Failed to read file as base64: %v", err)
		}

		expectedBase64 := base64.StdEncoding.EncodeToString([]byte(content))
		if base64Content != expectedBase64 {
			t.Errorf("Expected base64 '%s', got '%s'", expectedBase64, base64Content)
		}
	})

	// Test gzip compression
	t.Run("GzipUpload", func(t *testing.T) {
		content := "This is a test file that will be compressed with gzip."
		reader := strings.NewReader(content)

		fileHeader := &FileHeader{
			FileHeader: &multipart.FileHeader{
				Filename: "test_gzip.txt",
				Size:     int64(len(content)),
				Header:   make(map[string][]string),
			},
		}
		fileHeader.Header.Set("Content-Type", "text/plain")

		option := UploadOption{
			Gzip:   true,
			UserID: "user123",
		}

		file, err := manager.Upload(context.Background(), fileHeader, reader, option)
		if err != nil {
			t.Fatalf("Failed to upload gzipped file: %v", err)
		}

		// The stored file should be compressed, but when we read it back,
		// we should get the original content (if the storage handles decompression)
		downloadedContent, err := manager.Read(context.Background(), file.ID)
		if err != nil {
			t.Fatalf("Failed to read gzipped file: %v", err)
		}

		// Since we're storing compressed data, we need to decompress it
		decompressed, err := Gunzip(downloadedContent)
		if err != nil {
			t.Fatalf("Failed to decompress file: %v", err)
		}

		if string(decompressed) != content {
			t.Errorf("Expected content '%s', got '%s'", content, string(decompressed))
		}
	})

	// Test chunked upload
	t.Run("ChunkedUpload", func(t *testing.T) {
		content := "This is a large file that will be uploaded in chunks. " +
			strings.Repeat("Lorem ipsum dolor sit amet, consectetur adipiscing elit. ", 100)

		chunkSize := 1024
		totalSize := len(content)

		var lastFile *File
		for start := 0; start < totalSize; start += chunkSize {
			end := start + chunkSize - 1
			if end >= totalSize {
				end = totalSize - 1
			}

			chunk := []byte(content[start : end+1])

			fileHeader := &FileHeader{
				FileHeader: &multipart.FileHeader{
					Filename: "large_file.txt",
					Size:     int64(len(chunk)),
					Header:   make(map[string][]string),
				},
			}
			fileHeader.Header.Set("Content-Type", "text/plain")
			fileHeader.Header.Set("Content-Range",
				fmt.Sprintf("bytes %d-%d/%d", start, end, totalSize))
			fileHeader.Header.Set("Content-Uid", "unique-file-id-123")

			option := UploadOption{
				UserID: "user123",
			}

			file, err := manager.Upload(context.Background(), fileHeader, bytes.NewReader(chunk), option)
			if err != nil {
				t.Fatalf("Failed to upload chunk starting at %d: %v", start, err)
			}

			lastFile = file
		}

		// After uploading all chunks, read the complete file
		if lastFile != nil {
			downloadedContent, err := manager.Read(context.Background(), lastFile.ID)
			if err != nil {
				t.Fatalf("Failed to read chunked file: %v", err)
			}

			if string(downloadedContent) != content {
				t.Errorf("Chunked upload content mismatch. Expected length %d, got %d",
					len(content), len(downloadedContent))
			}
		}
	})
}

func TestManagerValidation(t *testing.T) {
	manager, err := New(ManagerOption{
		Driver:       "local",
		MaxSize:      "1K", // Very small max size for testing
		AllowedTypes: []string{"text/plain"},
		Options: map[string]interface{}{
			"path": "/tmp/test_attachments",
		},
	})
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	// Test file size validation
	t.Run("FileSizeValidation", func(t *testing.T) {
		content := strings.Repeat("a", 2048) // 2KB, exceeds 1KB limit
		reader := strings.NewReader(content)

		fileHeader := &FileHeader{
			FileHeader: &multipart.FileHeader{
				Filename: "large.txt",
				Size:     int64(len(content)),
				Header:   make(map[string][]string),
			},
		}
		fileHeader.Header.Set("Content-Type", "text/plain")

		option := UploadOption{}

		_, err := manager.Upload(context.Background(), fileHeader, reader, option)
		if err == nil {
			t.Error("Expected error for file size exceeding limit")
		}
	})

	// Test file type validation
	t.Run("FileTypeValidation", func(t *testing.T) {
		content := "test"
		reader := strings.NewReader(content)

		fileHeader := &FileHeader{
			FileHeader: &multipart.FileHeader{
				Filename: "test.jpg",
				Size:     int64(len(content)),
				Header:   make(map[string][]string),
			},
		}
		fileHeader.Header.Set("Content-Type", "image/jpeg") // Not allowed

		option := UploadOption{}

		_, err := manager.Upload(context.Background(), fileHeader, reader, option)
		if err == nil {
			t.Error("Expected error for disallowed file type")
		}
	})
}