From a966ad83eeeaa0d765752f0fd398a008c1022ace Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Mon, 22 Dec 2025 09:19:23 +0800
Subject: [PATCH] Enhance Cleanup Logic in Auth Integration Tests

- Updated the cleanup process in `cleanupAuthCollections` to include a waiting mechanism for Qdrant to fully process deletions, improving reliability of test setups.
- Removed unnecessary sleep calls and added logging to warn if collections still exist after cleanup, ensuring better visibility during test execution.
- Refactored comments for clarity regarding the cleanup process in both `TestAuthSearchSetup` and `ensureAuthTestData` functions.
---
 .../assistant/search_auth_integration_test.go | 28 +++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/agent/assistant/search_auth_integration_test.go b/agent/assistant/search_auth_integration_test.go
index 24cdab51..f7b3722a 100644
--- a/agent/assistant/search_auth_integration_test.go
+++ b/agent/assistant/search_auth_integration_test.go
@@ -60,10 +60,9 @@ func TestAuthSearchSetup(t *testing.T) {
 		return
 	}
 
-	// Cleanup existing
+	// Cleanup existing (includes waiting for deletion to complete)
 	t.Log("Cleaning up existing collections...")
 	cleanupAuthCollections(ctx, t)
-	time.Sleep(1 * time.Second)
 
 	// Create Team1 collection (owned by UserA, Team1)
 	t.Log("Creating Team1 collection...")
@@ -461,10 +460,9 @@ func ensureAuthTestData(t *testing.T) {
 
 // createAuthTestData creates the test collections and documents
 func createAuthTestData(t *testing.T, ctx context.Context) {
-	// Cleanup existing
+	// Cleanup existing (includes waiting for deletion to complete)
 	t.Log("Cleaning up existing collections...")
 	cleanupAuthCollections(ctx, t)
-	time.Sleep(1 * time.Second)
 
 	// Create Team1 collection (owned by UserA, Team1)
 	t.Log("Creating Team1 collection...")
@@ -529,6 +527,24 @@ func cleanupAuthCollections(ctx context.Context, t *testing.T) {
 			t.Logf("  Removed: %s", id)
 		}
 	}
+
+	// Wait for Qdrant to fully process deletions
+	// This is necessary because Qdrant may take time to propagate deletions
+	maxRetries := 10
+	for retry := 0; retry < maxRetries; retry++ {
+		allGone := true
+		for _, id := range collections {
+			if exists, err := kb.API.CollectionExists(ctx, id); err == nil && exists.Exists {
+				allGone = false
+				break
+			}
+		}
+		if allGone {
+			return
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+	t.Log("  Warning: Some collections may still exist after cleanup")
 }
 
 func createAuthCollection(ctx context.Context, t *testing.T, id, userID, teamID string, public bool, share string) {
@@ -597,10 +613,6 @@ func sanitizeForID(s string) string {
 	return result
 }
 
-func executeKBSearch(t *testing.T, collectionID, query string, metadata map[string]interface{}) *searchTypes.Result {
-	return executeKBSearchOnCollections(t, []string{collectionID}, query)
-}
-
 func executeKBSearchOnCollections(t *testing.T, collections []string, query string) *searchTypes.Result {
 	if len(collections) == 0 {
 		return &searchTypes.Result{Items: []*searchTypes.ResultItem{}}