yao/kb/api/addurl_test.go

package api_test

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	graphragtypes "github.com/yaoapp/gou/graphrag/types"
	"github.com/yaoapp/yao/kb"
	"github.com/yaoapp/yao/kb/api"
)

// Note: TestMain is defined in collection_test.go, which handles environment setup
// Run tests with: source env.local.sh && go test -v ./kb/api/...

// createTestCollectionForURL is a helper to create a test collection for URL tests
func createTestCollectionForURL(t *testing.T, ctx context.Context) string {
	if kb.API == nil {
		t.Skip("KB API not initialized")
	}

	collectionID := fmt.Sprintf("test_url_%d", time.Now().UnixNano())

	params := &api.CreateCollectionParams{
		ID: collectionID,
		Metadata: map[string]interface{}{
			"name":        "Test URL Collection",
			"description": "Collection for AddURL tests",
		},
		EmbeddingProviderID: "__yao.openai",
		EmbeddingOptionID:   "text-embedding-3-small",
		Locale:              "en",
		Config: &graphragtypes.CreateCollectionOptions{
			Distance:  "cosine",
			IndexType: "hnsw",
		},
	}

	_, err := kb.API.CreateCollection(ctx, params)
	if err != nil {
		t.Fatalf("Failed to create test collection: %v", err)
	}

	return collectionID
}

// cleanupTestCollectionForURL removes a test collection
func cleanupTestCollectionForURL(ctx context.Context, collectionID string) {
	if kb.API != nil {
		_, _ = kb.API.RemoveCollection(ctx, collectionID)
	}
}

// ========== AddURL Tests ==========

func TestAddURL(t *testing.T) {
	if kb.API == nil {
		t.Skip("KB API not initialized")
	}

	ctx := context.Background()
	collectionID := createTestCollectionForURL(t, ctx)
	defer cleanupTestCollectionForURL(ctx, collectionID)

	t.Run("AddURLSuccess", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go",
			Locale:       "en",
			Metadata: map[string]interface{}{
				"title":       "Yao Agent Caller",
				"description": "A test URL document",
			},
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
			AuthScope: map[string]interface{}{
				"__yao_created_by": "test_user",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		if result != nil {
			assert.Equal(t, collectionID, result.CollectionID)
			assert.NotEmpty(t, result.DocID)
			assert.Equal(t, "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go", result.URL)
			assert.Contains(t, result.Message, "successfully")
			t.Logf("Added URL document: %s", result.DocID)
		}

		// Verify document was created
		if result != nil {
			doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
			assert.NoError(t, err)
			assert.NotNil(t, doc)
			assert.Equal(t, "url", doc["type"])
			assert.Equal(t, "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go", doc["url"])
			t.Logf("✅ URL Document verified: type=%v, url=%v, status=%v", doc["type"], doc["url"], doc["status"])
		}
	})

	t.Run("AddURLMissingCollectionID", func(t *testing.T) {
		params := &api.AddURLParams{
			URL: "https://example.com",
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "collection_id is required")
	})

	t.Run("AddURLMissingURL", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "url is required")
	})

	t.Run("AddURLMissingChunking", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com",
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "chunking configuration is required")
	})

	t.Run("AddURLMissingEmbedding", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com",
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "embedding configuration is required")
	})

	t.Run("AddURLWithCustomDocID", func(t *testing.T) {
		customDocID := fmt.Sprintf("custom_url_doc_%d", time.Now().UnixNano())
		params := &api.AddURLParams{
			CollectionID: collectionID,
			DocID:        customDocID,
			URL:          "https://example.com", // Use root URL which always exists
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		if result != nil {
			assert.Equal(t, customDocID, result.DocID)
			t.Logf("Added URL with custom DocID: %s", result.DocID)
		}
	})

	t.Run("AddURLWithTitleFromMetadata", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com", // Use root URL which always exists
			Metadata: map[string]interface{}{
				"title": "Custom URL Title From Metadata",
			},
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
		}

		result, err := kb.API.AddURL(ctx, params)
		assert.NoError(t, err)
		assert.NotNil(t, result)

		if result != nil {
			doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
			assert.NoError(t, err)
			assert.Equal(t, "Custom URL Title From Metadata", doc["name"])
			t.Logf("✅ Title from metadata verified: %v", doc["name"])
		}
	})
}

// ========== AddURLAsync Tests ==========

func TestAddURLAsync(t *testing.T) {
	if kb.API == nil {
		t.Skip("KB API not initialized")
	}

	ctx := context.Background()
	collectionID := createTestCollectionForURL(t, ctx)
	defer cleanupTestCollectionForURL(ctx, collectionID)

	t.Run("AddURLAsyncSuccess", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com",
			Locale:       "en",
			Metadata: map[string]interface{}{
				"title": "Async URL Document",
			},
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
			Job: &api.JobOptionsParams{
				Name:        "Test Async URL Job",
				Description: "Testing async URL processing",
				Category:    "Test",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		if result != nil {
			assert.NotEmpty(t, result.JobID)
			assert.NotEmpty(t, result.DocID)
			t.Logf("Created async URL job: %s for document: %s", result.JobID, result.DocID)
		}

		// Verify document was created
		if result != nil {
			doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
			assert.NoError(t, err)
			assert.NotNil(t, doc)
			assert.Equal(t, "url", doc["type"])
			assert.Equal(t, result.JobID, doc["job_id"])
			t.Logf("✅ Async URL document created: status=%v, job_id=%v", doc["status"], doc["job_id"])

			// Wait for job to complete (max 30 seconds)
			maxWait := 30 * time.Second
			pollInterval := 500 * time.Millisecond
			startTime := time.Now()
			var finalStatus string

			for time.Since(startTime) < maxWait {
				doc, err = kb.API.GetDocument(ctx, result.DocID, nil)
				if err != nil {
					t.Logf("Error getting document: %v", err)
					break
				}
				finalStatus, _ = doc["status"].(string)
				if finalStatus == "completed" || finalStatus == "error" {
					break
				}
				time.Sleep(pollInterval)
			}

			t.Logf("✅ Job completed: final status=%s, elapsed=%v", finalStatus, time.Since(startTime))
			assert.Equal(t, "completed", finalStatus, "Job should complete successfully")
		}
	})

	t.Run("AddURLAsyncMissingCollectionID", func(t *testing.T) {
		params := &api.AddURLParams{
			URL: "https://example.com",
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "collection_id is required")
	})

	t.Run("AddURLAsyncMissingURL", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "url is required")
	})

	t.Run("AddURLAsyncMissingChunking", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com",
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "chunking configuration is required")
	})

	t.Run("AddURLAsyncMissingEmbedding", func(t *testing.T) {
		params := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com",
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.Error(t, err)
		assert.Nil(t, result)
		assert.Contains(t, err.Error(), "embedding configuration is required")
	})

	t.Run("AddURLAsyncWithCustomDocID", func(t *testing.T) {
		customDocID := fmt.Sprintf("async_custom_url_%d", time.Now().UnixNano())
		params := &api.AddURLParams{
			CollectionID: collectionID,
			DocID:        customDocID,
			URL:          "https://example.com/async-custom",
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
		}

		result, err := kb.API.AddURLAsync(ctx, params)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		if result != nil {
			assert.Equal(t, customDocID, result.DocID)
			t.Logf("Created async URL with custom DocID: %s", result.DocID)
		}
	})
}

// ========== AddURL Integration Test ==========

func TestAddURLIntegration(t *testing.T) {
	if kb.API == nil {
		t.Skip("KB API not initialized")
	}

	ctx := context.Background()
	collectionID := createTestCollectionForURL(t, ctx)
	defer cleanupTestCollectionForURL(ctx, collectionID)

	t.Run("FullURLLifecycle", func(t *testing.T) {
		// 1. Add URL Document
		addParams := &api.AddURLParams{
			CollectionID: collectionID,
			URL:          "https://example.com", // Use root URL which always exists
			Locale:       "en",
			Metadata: map[string]interface{}{
				"title":       "URL Lifecycle Test",
				"description": "Full lifecycle integration test",
			},
			Chunking: &api.ProviderConfigParams{
				ProviderID: "__yao.structured",
				OptionID:   "standard",
			},
			Embedding: &api.ProviderConfigParams{
				ProviderID: "__yao.openai",
				OptionID:   "text-embedding-3-small",
			},
			Fetcher: &api.ProviderConfigParams{
				ProviderID: "__yao.http",
				OptionID:   "http",
			},
			AuthScope: map[string]interface{}{
				"__yao_created_by": "integration_test",
			},
		}

		result, err := kb.API.AddURL(ctx, addParams)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		if result == nil {
			t.Fatalf("Failed to create URL document: result is nil")
		}
		t.Logf("1. Created URL document: %s", result.DocID)

		// 2. Get Document
		doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
		assert.NoError(t, err)
		assert.NotNil(t, doc)
		assert.Equal(t, "URL Lifecycle Test", doc["name"])
		assert.Equal(t, "url", doc["type"])
		assert.Equal(t, "https://example.com", doc["url"])
		t.Logf("2. Retrieved document: name=%v, url=%v, status=%v", doc["name"], doc["url"], doc["status"])

		// 3. List Documents
		listFilter := &api.ListDocumentsFilter{
			Page:         1,
			PageSize:     20,
			CollectionID: collectionID,
			Keywords:     "URL Lifecycle",
		}
		listResult, err := kb.API.ListDocuments(ctx, listFilter)
		assert.NoError(t, err)
		assert.GreaterOrEqual(t, len(listResult.Data), 1)
		t.Logf("3. Found document in list: %d documents", len(listResult.Data))

		// 4. Remove Document
		removeParams := &api.RemoveDocumentsParams{
			DocumentIDs: []string{result.DocID},
		}
		removeResult, err := kb.API.RemoveDocuments(ctx, removeParams)
		assert.NoError(t, err)
		assert.NotNil(t, removeResult)
		t.Logf("4. Removed document: %d deleted", removeResult.DeletedCount)

		// 5. Verify Removal
		_, err = kb.API.GetDocument(ctx, result.DocID, nil)
		assert.Error(t, err)
		t.Logf("5. Verified document removal")

		t.Logf("✅ Full URL lifecycle test completed successfully")
	})
}