CHORUS/pkg/ucxi/storage_test.go

package ucxi

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func createTempStorageDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "ucxi-storage-test-*")
	if err != nil {
		t.Fatalf("Failed to create temp directory: %v", err)
	}
	return dir
}

func TestNewBasicContentStorage(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Errorf("NewBasicContentStorage failed: %v", err)
	}
	
	if storage == nil {
		t.Error("NewBasicContentStorage should not return nil")
	}
	
	if storage.basePath != tempDir {
		t.Errorf("Base path = %s, want %s", storage.basePath, tempDir)
	}
	
	// Verify directory was created
	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
		t.Error("Storage directory should be created")
	}
}

func TestNewBasicContentStorageWithInvalidPath(t *testing.T) {
	// Try to create storage with invalid path (e.g., a file instead of directory)
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	// Create a file at the path
	invalidPath := filepath.Join(tempDir, "file-not-dir")
	if err := ioutil.WriteFile(invalidPath, []byte("test"), 0644); err != nil {
		t.Fatalf("Failed to create test file: %v", err)
	}
	
	// This should fail because the path exists as a file, not a directory
	_, err := NewBasicContentStorage(invalidPath)
	if err == nil {
		t.Error("NewBasicContentStorage should fail with invalid path")
	}
}

func TestStorageStoreAndRetrieve(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	key := "test-key"
	content := &Content{
		Data:        []byte("test content data"),
		ContentType: "text/plain",
		Metadata: map[string]string{
			"author":  "test-author",
			"version": "1.0",
		},
		Version:   1,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
		Author:    "test-user",
	}
	
	// Test store
	err = storage.Store(ctx, key, content)
	if err != nil {
		t.Errorf("Store failed: %v", err)
	}
	
	// Test retrieve
	retrieved, err := storage.Retrieve(ctx, key)
	if err != nil {
		t.Errorf("Retrieve failed: %v", err)
	}
	
	if retrieved == nil {
		t.Error("Retrieved content should not be nil")
	}
	
	// Verify content matches
	if string(retrieved.Data) != string(content.Data) {
		t.Errorf("Data mismatch: got %s, want %s", string(retrieved.Data), string(content.Data))
	}
	
	if retrieved.ContentType != content.ContentType {
		t.Errorf("ContentType mismatch: got %s, want %s", retrieved.ContentType, content.ContentType)
	}
	
	if retrieved.Author != content.Author {
		t.Errorf("Author mismatch: got %s, want %s", retrieved.Author, content.Author)
	}
	
	if retrieved.Version != content.Version {
		t.Errorf("Version mismatch: got %d, want %d", retrieved.Version, content.Version)
	}
	
	// Verify metadata
	if len(retrieved.Metadata) != len(content.Metadata) {
		t.Errorf("Metadata length mismatch: got %d, want %d", len(retrieved.Metadata), len(content.Metadata))
	}
	
	for key, value := range content.Metadata {
		if retrieved.Metadata[key] != value {
			t.Errorf("Metadata[%s] mismatch: got %s, want %s", key, retrieved.Metadata[key], value)
		}
	}
	
	// Verify checksum is calculated
	if retrieved.Checksum == "" {
		t.Error("Checksum should be calculated and stored")
	}
}

func TestStorageChecksumValidation(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	key := "checksum-test"
	content := &Content{
		Data:        []byte("test content for checksum"),
		ContentType: "text/plain",
	}
	
	// Store content (checksum will be calculated automatically)
	err = storage.Store(ctx, key, content)
	if err != nil {
		t.Errorf("Store failed: %v", err)
	}
	
	// Retrieve and verify checksum validation works
	retrieved, err := storage.Retrieve(ctx, key)
	if err != nil {
		t.Errorf("Retrieve failed: %v", err)
	}
	
	if retrieved.Checksum == "" {
		t.Error("Checksum should be set after storing")
	}
	
	// Manually corrupt the file to test checksum validation
	filePath := storage.getFilePath(key)
	originalData, err := ioutil.ReadFile(filePath)
	if err != nil {
		t.Fatalf("Failed to read file: %v", err)
	}
	
	// Corrupt the data in the JSON by changing base64 encoded data
	// The content is base64 encoded in JSON, so we'll replace some characters
	corruptedData := strings.Replace(string(originalData), "dGVzdCBjb250ZW50IGZvciBjaGVja3N1bQ==", "Y29ycnVwdGVkIGNvbnRlbnQ=", 1)
	if corruptedData == string(originalData) {
		// If the base64 replacement didn't work, try a simpler corruption
		corruptedData = strings.Replace(string(originalData), "\"", "'", 1)
		if corruptedData == string(originalData) {
			t.Fatalf("Failed to corrupt data - no changes made")
		}
	}
	
	err = ioutil.WriteFile(filePath, []byte(corruptedData), 0644)
	if err != nil {
		t.Fatalf("Failed to write corrupted file: %v", err)
	}
	
	// Retrieve should fail due to checksum mismatch
	_, err = storage.Retrieve(ctx, key)
	if err == nil {
		t.Error("Retrieve should fail with corrupted content")
	}
	
	if !strings.Contains(err.Error(), "checksum mismatch") {
		t.Errorf("Error should mention checksum mismatch, got: %v", err)
	}
}

func TestStorageDelete(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	key := "delete-test"
	content := &Content{Data: []byte("content to delete")}
	
	// Store content
	err = storage.Store(ctx, key, content)
	if err != nil {
		t.Errorf("Store failed: %v", err)
	}
	
	// Verify it exists
	exists, err := storage.Exists(ctx, key)
	if err != nil {
		t.Errorf("Exists check failed: %v", err)
	}
	if !exists {
		t.Error("Content should exist after storing")
	}
	
	// Delete content
	err = storage.Delete(ctx, key)
	if err != nil {
		t.Errorf("Delete failed: %v", err)
	}
	
	// Verify it no longer exists
	exists, err = storage.Exists(ctx, key)
	if err != nil {
		t.Errorf("Exists check after delete failed: %v", err)
	}
	if exists {
		t.Error("Content should not exist after deletion")
	}
	
	// Verify retrieve fails
	_, err = storage.Retrieve(ctx, key)
	if err == nil {
		t.Error("Retrieve should fail for deleted content")
	}
	
	// Delete non-existent key should fail
	err = storage.Delete(ctx, "non-existent-key")
	if err == nil {
		t.Error("Delete should fail for non-existent key")
	}
}

func TestStorageList(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	
	// Store multiple pieces of content
	testKeys := []string{
		"prefix1/key1",
		"prefix1/key2", 
		"prefix2/key1",
		"prefix2/key2",
		"different-prefix/key1",
	}
	
	for i, key := range testKeys {
		content := &Content{Data: []byte(fmt.Sprintf("content-%d", i))}
		err = storage.Store(ctx, key, content)
		if err != nil {
			t.Errorf("Store failed for key %s: %v", key, err)
		}
	}
	
	// Test list all
	allKeys, err := storage.List(ctx, "")
	if err != nil {
		t.Errorf("List all failed: %v", err)
	}
	
	if len(allKeys) != len(testKeys) {
		t.Errorf("List all returned %d keys, want %d", len(allKeys), len(testKeys))
	}
	
	// Test list with prefix
	prefix1Keys, err := storage.List(ctx, "prefix1/")
	if err != nil {
		t.Errorf("List with prefix failed: %v", err)
	}
	
	if len(prefix1Keys) != 2 {
		t.Errorf("List prefix1/ returned %d keys, want 2", len(prefix1Keys))
	}
	
	// Verify the keys match the prefix
	for _, key := range prefix1Keys {
		if !strings.HasPrefix(key, "prefix1/") {
			t.Errorf("Key %s should have prefix 'prefix1/'", key)
		}
	}
	
	// Test list with non-existent prefix
	noKeys, err := storage.List(ctx, "nonexistent/")
	if err != nil {
		t.Errorf("List non-existent prefix failed: %v", err)
	}
	
	if len(noKeys) != 0 {
		t.Errorf("List non-existent prefix returned %d keys, want 0", len(noKeys))
	}
}

func TestStorageExists(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	key := "exists-test"
	
	// Initially should not exist
	exists, err := storage.Exists(ctx, key)
	if err != nil {
		t.Errorf("Exists check failed: %v", err)
	}
	if exists {
		t.Error("Key should not exist initially")
	}
	
	// Store content
	content := &Content{Data: []byte("test")}
	err = storage.Store(ctx, key, content)
	if err != nil {
		t.Errorf("Store failed: %v", err)
	}
	
	// Should exist now
	exists, err = storage.Exists(ctx, key)
	if err != nil {
		t.Errorf("Exists check after store failed: %v", err)
	}
	if !exists {
		t.Error("Key should exist after storing")
	}
	
	// Delete content
	err = storage.Delete(ctx, key)
	if err != nil {
		t.Errorf("Delete failed: %v", err)
	}
	
	// Should not exist anymore
	exists, err = storage.Exists(ctx, key)
	if err != nil {
		t.Errorf("Exists check after delete failed: %v", err)
	}
	if exists {
		t.Error("Key should not exist after deletion")
	}
}

func TestStorageClear(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	
	// Store multiple pieces of content
	for i := 0; i < 5; i++ {
		key := fmt.Sprintf("key-%d", i)
		content := &Content{Data: []byte(fmt.Sprintf("content-%d", i))}
		err = storage.Store(ctx, key, content)
		if err != nil {
			t.Errorf("Store failed for key %s: %v", key, err)
		}
	}
	
	// Verify content exists
	keys, err := storage.List(ctx, "")
	if err != nil {
		t.Errorf("List failed: %v", err)
	}
	if len(keys) != 5 {
		t.Errorf("Expected 5 keys before clear, got %d", len(keys))
	}
	
	// Clear all content
	err = storage.Clear(ctx)
	if err != nil {
		t.Errorf("Clear failed: %v", err)
	}
	
	// Verify all content is gone
	keys, err = storage.List(ctx, "")
	if err != nil {
		t.Errorf("List after clear failed: %v", err)
	}
	if len(keys) != 0 {
		t.Errorf("Expected 0 keys after clear, got %d", len(keys))
	}
	
	// Verify directory still exists but is empty
	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
		t.Error("Base directory should still exist after clear")
	}
	
	entries, err := ioutil.ReadDir(tempDir)
	if err != nil {
		t.Errorf("Failed to read directory after clear: %v", err)
	}
	if len(entries) != 0 {
		t.Errorf("Directory should be empty after clear, found %d entries", len(entries))
	}
}

func TestStorageGetStorageStats(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	
	// Initially should have no files
	stats, err := storage.GetStorageStats()
	if err != nil {
		t.Errorf("GetStorageStats failed: %v", err)
	}
	
	if stats["file_count"].(int) != 0 {
		t.Errorf("Initial file count = %d, want 0", stats["file_count"])
	}
	
	if stats["total_size"].(int64) != 0 {
		t.Errorf("Initial total size = %d, want 0", stats["total_size"])
	}
	
	if stats["base_path"].(string) != tempDir {
		t.Errorf("Base path = %s, want %s", stats["base_path"], tempDir)
	}
	
	// Store some content
	for i := 0; i < 3; i++ {
		key := fmt.Sprintf("stats-key-%d", i)
		content := &Content{Data: []byte(fmt.Sprintf("test content %d", i))}
		err = storage.Store(ctx, key, content)
		if err != nil {
			t.Errorf("Store failed: %v", err)
		}
	}
	
	// Check stats again
	stats, err = storage.GetStorageStats()
	if err != nil {
		t.Errorf("GetStorageStats after store failed: %v", err)
	}
	
	if stats["file_count"].(int) != 3 {
		t.Errorf("File count after storing = %d, want 3", stats["file_count"])
	}
	
	if stats["total_size"].(int64) <= 0 {
		t.Error("Total size should be greater than 0 after storing content")
	}
}

func TestStorageGetFilePath(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	tests := []struct {
		name         string
		key          string
		shouldContain []string
		shouldNotContain []string
	}{
		{
			name: "simple key",
			key:  "simple-key",
			shouldContain: []string{"simple-key.json"},
			shouldNotContain: []string{":"},
		},
		{
			name: "key with colons",
			key:  "agent:role",
			shouldContain: []string{"agent_role.json"},
			shouldNotContain: []string{":"},
		},
		{
			name: "key with at symbol",
			key:  "agent@project",
			shouldContain: []string{"agent_at_project.json"},
			shouldNotContain: []string{"@"},
		},
		{
			name: "key with slashes",
			key:  "path/to/resource",
			shouldContain: []string{".json"},
			// Should not contain the original slash as literal
		},
	}
	
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			filePath := storage.getFilePath(tt.key)
			
			// Should always start with base path
			if !strings.HasPrefix(filePath, tempDir) {
				t.Errorf("File path should start with base path")
			}
			
			// Should always end with .json
			if !strings.HasSuffix(filePath, ".json") {
				t.Errorf("File path should end with .json")
			}
			
			// Check required substrings
			for _, required := range tt.shouldContain {
				if !strings.Contains(filePath, required) {
					t.Errorf("File path should contain '%s', got: %s", required, filePath)
				}
			}
			
			// Check forbidden substrings
			for _, forbidden := range tt.shouldNotContain {
				if strings.Contains(filePath, forbidden) {
					t.Errorf("File path should not contain '%s', got: %s", forbidden, filePath)
				}
			}
		})
	}
}

func TestStorageErrorCases(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	
	// Test empty key
	content := &Content{Data: []byte("test")}
	
	err = storage.Store(ctx, "", content)
	if err == nil {
		t.Error("Store with empty key should fail")
	}
	
	_, err = storage.Retrieve(ctx, "")
	if err == nil {
		t.Error("Retrieve with empty key should fail")
	}
	
	err = storage.Delete(ctx, "")
	if err == nil {
		t.Error("Delete with empty key should fail")
	}
	
	_, err = storage.Exists(ctx, "")
	if err == nil {
		t.Error("Exists with empty key should fail")
	}
	
	// Test nil content
	err = storage.Store(ctx, "test-key", nil)
	if err == nil {
		t.Error("Store with nil content should fail")
	}
	
	// Test retrieve non-existent key
	_, err = storage.Retrieve(ctx, "non-existent-key")
	if err == nil {
		t.Error("Retrieve non-existent key should fail")
	}
}

// Test concurrent access to storage
func TestStorageConcurrency(t *testing.T) {
	tempDir := createTempStorageDir(t)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		t.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	done := make(chan bool, 10)
	
	// Run multiple goroutines that store, retrieve, and delete content
	for i := 0; i < 10; i++ {
		go func(id int) {
			defer func() { done <- true }()
			
			key := fmt.Sprintf("concurrent-key-%d", id)
			content := &Content{Data: []byte(fmt.Sprintf("content-%d", id))}
			
			// Store
			if err := storage.Store(ctx, key, content); err != nil {
				t.Errorf("Goroutine %d store failed: %v", id, err)
				return
			}
			
			// Retrieve
			if _, err := storage.Retrieve(ctx, key); err != nil {
				t.Errorf("Goroutine %d retrieve failed: %v", id, err)
				return
			}
			
			// Delete
			if err := storage.Delete(ctx, key); err != nil {
				t.Errorf("Goroutine %d delete failed: %v", id, err)
				return
			}
		}(i)
	}
	
	// Wait for all goroutines to complete
	for i := 0; i < 10; i++ {
		<-done
	}
	
	// Verify final state - all content should be deleted
	keys, err := storage.List(ctx, "")
	if err != nil {
		t.Errorf("List after concurrent operations failed: %v", err)
	}
	
	if len(keys) != 0 {
		t.Errorf("Expected 0 keys after concurrent operations, got %d", len(keys))
	}
}

// Benchmark tests
func BenchmarkStorageStore(b *testing.B) {
	tempDir := createTempStorageDirForBench(b)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		b.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	content := &Content{
		Data:        []byte("benchmark test content"),
		ContentType: "text/plain",
	}
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		key := fmt.Sprintf("benchmark-key-%d", i)
		storage.Store(ctx, key, content)
	}
}

func BenchmarkStorageRetrieve(b *testing.B) {
	tempDir := createTempStorageDirForBench(b)
	defer os.RemoveAll(tempDir)
	
	storage, err := NewBasicContentStorage(tempDir)
	if err != nil {
		b.Fatalf("Failed to create storage: %v", err)
	}
	
	ctx := context.Background()
	content := &Content{
		Data:        []byte("benchmark test content"),
		ContentType: "text/plain",
	}
	
	// Pre-populate storage
	keys := make([]string, 1000)
	for i := 0; i < 1000; i++ {
		keys[i] = fmt.Sprintf("benchmark-key-%d", i)
		storage.Store(ctx, keys[i], content)
	}
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		key := keys[i%1000]
		storage.Retrieve(ctx, key)
	}
}

// Helper function for benchmark that creates temp directory
func createTempStorageDirForBench(t testing.TB) string {
	dir, err := ioutil.TempDir("", "ucxi-storage-test-*")
	if err != nil {
		t.Fatalf("Failed to create temp directory: %v", err)
	}
	return dir
}