standard/mongodb/tool/mongo_test.go

package tool

import (
	"context"
	"fmt"
	"testing"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

// 配置测试用的 MongoDB 连接信息
var cfg = Config{
	Uri:      "mongodb://localhost:27017", // 请根据实际情况修改
	Database: "test",
	Username: "test",
	Password: "123456",
}

func TestInitMongo(t *testing.T) {

	// 调用 InitMongo 方法
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("InitMongo failed: %v", err)
	}
	defer func() {
		// 测试结束后断开连接
		err := repo.DB.Client().Disconnect(context.TODO())
		if err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	// 验证连接是否成功
	err = repo.DB.Client().Ping(context.TODO(), readpref.Primary())
	if err != nil {
		t.Errorf("Failed to ping MongoDB: %v", err)
	}

	// 验证数据库实例是否正确
	if repo.DB.Name() != cfg.Database {
		t.Errorf("Expected database name %s, but got %s", cfg.Database, repo.DB.Name())
	}
}

// TestInsertOne 测试 InsertOne 方法
func TestInsertOne(t *testing.T) {
	// 初始化 MongoDB 连接，使用全局的 cfg
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		// 测试结束后断开连接
		err := repo.DB.Client().Disconnect(context.TODO())
		if err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	// 验证连接是否成功
	err = repo.DB.Client().Ping(context.TODO(), readpref.Primary())
	if err != nil {
		t.Fatalf("Failed to ping MongoDB: %v", err)
	}

	// 准备测试数据
	collectionName := "testCollection"
	testDocument := bson.M{
		"name":  "Test User",
		"age":   25,
		"email": "test@example.com",
	}

	// 调用 InsertOne 方法
	result, err := repo.InsertOne(context.TODO(), collectionName, testDocument)
	if err != nil {
		t.Fatalf("InsertOne failed: %v", err)
	}

	// 验证插入结果
	if result.InsertedID == nil {
		t.Errorf("InsertedID is nil, insert may have failed")
	}

	// 验证文档是否实际插入到数据库
	var insertedDoc bson.M
	err = repo.DB.Collection(collectionName).FindOne(context.TODO(), bson.M{"_id": result.InsertedID}).Decode(&insertedDoc)
	if err != nil {
		t.Errorf("Failed to find inserted document: %v", err)
	}
	// 打印查询出来的插入数据
	fmt.Printf("查询到的插入数据: %+v\n", insertedDoc)
	// 清理测试数据
	_, err = repo.DB.Collection(collectionName).DeleteOne(context.TODO(), bson.M{"_id": result.InsertedID})
	if err != nil {
		t.Errorf("Failed to clean up test data: %v", err)
	}
}

func TestUpsert(t *testing.T) {
	// 使用全局 cfg 初始化 MongoDB 连接
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		// 测试结束后断开连接
		err = repo.DB.Client().Disconnect(context.TODO())
		if err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	// 验证连接是否成功
	err = repo.DB.Client().Ping(context.TODO(), readpref.Primary())
	if err != nil {
		t.Fatalf("Failed to ping MongoDB: %v", err)
	}

	collectionName := "testUpsertCollection"

	// 测试文档不存在时的 Upsert
	t.Run("Upsert when document does not exist", func(t *testing.T) {
		filter := bson.M{"name": "newUser"}
		update := bson.M{"$set": bson.M{"age": 25}}

		result, err := repo.Upsert(context.TODO(), collectionName, filter, update)
		if err != nil {
			t.Fatalf("Upsert failed: %v", err)
		}

		if result.MatchedCount != 0 {
			t.Errorf("Expected matched count 0, got %d", result.MatchedCount)
		}
		if result.ModifiedCount != 0 {
			t.Errorf("Expected modified count 0, got %d", result.ModifiedCount)
		}
		if result.UpsertedID == nil {
			t.Errorf("Expected non-nil upserted ID, got nil")
		}

		// 查询插入的文档
		var insertedDoc bson.M
		err = repo.DB.Collection(collectionName).FindOne(context.TODO(), bson.M{"_id": result.UpsertedID}).Decode(&insertedDoc)
		if err != nil {
			t.Errorf("Failed to find inserted document: %v", err)
		}

		// 打印插入的文档
		fmt.Printf("插入的文档: %+v\n", insertedDoc)

		// 清理测试数据
		_, err = repo.DB.Collection(collectionName).DeleteOne(context.TODO(), bson.M{"_id": result.UpsertedID})
		if err != nil {
			t.Errorf("Failed to clean up test data: %v", err)
		}
	})

	// 测试文档存在时的 Upsert
	t.Run("Upsert when document exists", func(t *testing.T) {
		// 先插入一个文档
		insertResult, err := repo.InsertOne(context.TODO(), collectionName, bson.M{"name": "existingUser", "age": 20})
		if err != nil {
			t.Fatalf("Failed to insert test document: %v", err)
		}

		filter := bson.M{"name": "existingUser"}
		update := bson.M{"$set": bson.M{"age": 30}}

		result, err := repo.Upsert(context.TODO(), collectionName, filter, update)
		if err != nil {
			t.Fatalf("Upsert failed: %v", err)
		}

		if result.MatchedCount != 1 {
			t.Errorf("Expected matched count 1, got %d", result.MatchedCount)
		}
		if result.ModifiedCount != 1 {
			t.Errorf("Expected modified count 1, got %d", result.ModifiedCount)
		}
		if result.UpsertedID != nil {
			t.Errorf("Expected nil upserted ID, got %v", result.UpsertedID)
		}

		// 查询更新后的文档
		var updatedDoc bson.M
		err = repo.DB.Collection(collectionName).FindOne(context.TODO(), bson.M{"_id": insertResult.InsertedID}).Decode(&updatedDoc)
		if err != nil {
			t.Errorf("Failed to find updated document: %v", err)
		}

		// 打印更新后的文档
		fmt.Printf("更新后的文档: %+v\n", updatedDoc)

		// 清理测试数据
		_, err = repo.DB.Collection(collectionName).DeleteOne(context.TODO(), bson.M{"_id": insertResult.InsertedID})
		if err != nil {
			t.Errorf("Failed to clean up test data: %v", err)
		}
	})
}

// TestNestedQuery 测试 NestedQuery 方法
func TestNestedQuery(t *testing.T) {
	// 使用全局 cfg 初始化 MongoDB 连接
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		// 测试结束后断开连接
		err = repo.DB.Client().Disconnect(context.TODO())
		if err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	// 验证连接是否成功
	err = repo.DB.Client().Ping(context.TODO(), readpref.Primary())
	if err != nil {
		t.Fatalf("Failed to ping MongoDB: %v", err)
	}

	collectionName := "testNestedQueryCollection"

	// 准备测试数据
	testDocuments := []interface{}{
		bson.M{
			"name": "Alice",
			"address": bson.M{
				"city": "New York",
				"zip":  "10001",
			},
		},
		bson.M{
			"name": "Bob",
			"address": bson.M{
				"city": "Los Angeles",
				"zip":  "90001",
			},
		},
	}

	// 插入测试数据
	_, err = repo.DB.Collection(collectionName).InsertMany(context.TODO(), testDocuments)
	if err != nil {
		t.Fatalf("Failed to insert test documents: %v", err)
	}

	// 定义嵌套查询条件
	nestedFilter := bson.M{
		"address.city": "New York",
	}

	// 执行嵌套查询
	cursor, err := repo.NestedQuery(context.TODO(), collectionName, nestedFilter)
	if err != nil {
		t.Fatalf("NestedQuery failed: %v", err)
	}
	defer cursor.Close(context.TODO())

	// 遍历游标并验证结果
	var results []bson.M
	for cursor.Next(context.TODO()) {
		var result bson.M
		err = cursor.Decode(&result)
		if err != nil {
			t.Errorf("Failed to decode document: %v", err)
			continue
		}
		results = append(results, result)
	}

	if err = cursor.Err(); err != nil {
		t.Errorf("Cursor error: %v", err)
	}

	// 验证查询结果数量
	if len(results) != 1 {
		t.Errorf("Expected 1 result, got %d", len(results))
	}

	// 打印查询结果
	fmt.Printf("查询结果: %+v\n", results)

	// 验证查询结果内容
	if results[0]["name"] != "Alice" {
		t.Errorf("Expected name 'Alice', got %v", results[0]["name"])
	}

	// 清理测试数据
	_, err = repo.DB.Collection(collectionName).DeleteMany(context.TODO(), bson.M{})
	if err != nil {
		t.Errorf("Failed to clean up test data: %v", err)
	}
}

// 测试 QueryArrayContains 方法
func TestQueryArrayContains(t *testing.T) {
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		if err := repo.DB.Client().Disconnect(context.TODO()); err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	collectionName := "testArrayContains"
	// 准备测试数据
	testDocuments := []interface{}{
		bson.M{"name": "doc1", "tags": []string{"apple", "banana"}},
		bson.M{"name": "doc2", "tags": []string{"cherry", "date"}},
	}

	// 插入测试数据
	_, err = repo.DB.Collection(collectionName).InsertMany(context.TODO(), testDocuments)
	if err != nil {
		t.Fatalf("Failed to insert test documents: %v", err)
	}

	// 执行查询
	elements := []interface{}{"apple"}
	cursor, err := repo.QueryArrayContains(context.TODO(), collectionName, "tags", elements)
	if err != nil {
		t.Fatalf("QueryArrayContains failed: %v", err)
	}
	defer cursor.Close(context.TODO())

	// 验证结果
	var results []bson.M
	for cursor.Next(context.TODO()) {
		var result bson.M
		if err := cursor.Decode(&result); err != nil {
			t.Errorf("Failed to decode document: %v", err)
			continue
		}
		results = append(results, result)
	}
	// 打印查询结果
	fmt.Printf("查询结果: %+v\n", results)

	if len(results) != 1 || results[0]["name"] != "doc1" {
		t.Errorf("Expected 1 result, got %d", len(results))
	}

	// 清理数据
	_, err = repo.DB.Collection(collectionName).DeleteMany(context.TODO(), bson.M{})
	if err != nil {
		t.Errorf("Failed to clean up test data: %v", err)
	}
}

func TestQueryArrayWithComplexFilter(t *testing.T) {
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		if err := repo.DB.Client().Disconnect(context.TODO()); err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	collectionName := "testArrayWithComplexFilter"
	// 准备测试数据
	testDocuments := []interface{}{
		bson.M{
			"name": "product1",
			"reviews": []bson.M{
				{"rating": 4, "author": "Alice"},
				{"rating": 2, "author": "Bob"},
			},
		},
		bson.M{
			"name": "product2",
			"reviews": []bson.M{
				{"rating": 1, "author": "Charlie"},
				{"rating": 3, "author": "David"},
			},
		},
	}

	// 插入测试数据
	_, err = repo.DB.Collection(collectionName).InsertMany(context.TODO(), testDocuments)
	if err != nil {
		t.Fatalf("Failed to insert test documents: %v", err)
	}

	// 定义复合筛选条件
	complexFilter := bson.M{
		"rating": bson.M{"$gte": 3},
		"author": "Alice",
	}

	// 执行查询
	cursor, err := repo.QueryArrayWithComplexFilter(context.TODO(), collectionName, "reviews", complexFilter)
	if err != nil {
		t.Fatalf("QueryArrayWithComplexFilter failed: %v", err)
	}
	defer cursor.Close(context.TODO())

	// 验证结果
	var results []bson.M
	for cursor.Next(context.TODO()) {
		var result bson.M
		if err := cursor.Decode(&result); err != nil {
			t.Errorf("Failed to decode document: %v", err)
			continue
		}
		results = append(results, result)
	}
	fmt.Printf("查询结果: %+v\n", results)
	if len(results) != 1 || results[0]["name"] != "product1" {
		t.Errorf("Expected 1 result, got %d", len(results))
	}

	// 清理数据
	_, err = repo.DB.Collection(collectionName).DeleteMany(context.TODO(), bson.M{})
	if err != nil {
		t.Errorf("Failed to clean up test data: %v", err)
	}
}

func TestQueryArrayElementByIndex(t *testing.T) {
	repo, err := InitMongo(cfg)
	if err != nil {
		t.Fatalf("Failed to initialize MongoDB: %v", err)
	}
	defer func() {
		if err := repo.DB.Client().Disconnect(context.TODO()); err != nil {
			t.Errorf("Failed to disconnect from MongoDB: %v", err)
		}
	}()

	// 验证连接
	if err := repo.DB.Client().Ping(context.TODO(), readpref.Primary()); err != nil {
		t.Fatalf("Failed to ping MongoDB: %v", err)
	}

	collectionName := "testQueryArrayElementByIndex"
	// 准备测试数据
	testDocuments := []interface{}{
		bson.M{
			"name":    "doc1",
			"numbers": []int{10, 20, 30},
		},
		bson.M{
			"name":    "doc2",
			"numbers": []int{40},
		},
	}

	// 插入测试数据
	_, err = repo.DB.Collection(collectionName).InsertMany(context.TODO(), testDocuments)
	if err != nil {
		t.Fatalf("Failed to insert test documents: %v", err)
	}

	// 按索引位置查询元素
	index := 1
	cursor, err := repo.QueryArrayElementByIndex(context.TODO(), collectionName, "numbers", index)
	if err != nil {
		t.Fatalf("QueryArrayElementByIndex failed: %v", err)
	}
	defer cursor.Close(context.TODO())

	// 验证结果
	var results []bson.M
	for cursor.Next(context.TODO()) {
		var result bson.M
		if err := cursor.Decode(&result); err != nil {
			t.Errorf("Failed to decode document: %v", err)
			continue
		}
		results = append(results, result)
	}

	if len(results) != 1 {
		t.Errorf("Expected 1 result, got %d", len(results))
	}

	// 清理数据
	_, err = repo.DB.Collection(collectionName).DeleteMany(context.TODO(), bson.M{})
	if err != nil {
		t.Errorf("Failed to clean up test data: %v", err)
	}
}