mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-22 17:31:03 +01:00
dc59599509
- Deleted the db module, consolidating database-related functionality into the core/mongo package for better organization and maintainability. - Updated all import paths across the codebase to replace references to the removed db module with core/mongo. - Cleaned up unused code and dependencies, enhancing overall project clarity and reducing complexity. - This refactor improves the structure of the codebase by centralizing database operations and simplifying module management.
538 lines
12 KiB
Go
538 lines
12 KiB
Go
package mongo
|
|
|
|
import (
|
|
"fmt"
|
|
"strconv"
|
|
"testing"
|
|
|
|
"github.com/spf13/viper"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/bson/primitive"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
)
|
|
|
|
type ColTestObject struct {
|
|
dbName string
|
|
colName string
|
|
col *Col
|
|
}
|
|
|
|
type TestDocument struct {
|
|
Key string `bson:"key"`
|
|
Value int `bson:"value"`
|
|
Tags []string `bson:"tags"`
|
|
}
|
|
|
|
type TestAggregateResult struct {
|
|
Id string `bson:"_id"`
|
|
Count int `bson:"count"`
|
|
Value int `bson:"value"`
|
|
}
|
|
|
|
func setupColTest() (to *ColTestObject, err error) {
|
|
dbName := "test_db"
|
|
colName := "test_col"
|
|
viper.Set("mongo.db", dbName)
|
|
col := GetMongoCol(colName)
|
|
if err := col.db.Drop(col.ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
return &ColTestObject{
|
|
dbName: dbName,
|
|
colName: colName,
|
|
col: col,
|
|
}, nil
|
|
}
|
|
|
|
func cleanupColTest(to *ColTestObject) {
|
|
_ = to.col.db.Drop(to.col.ctx)
|
|
}
|
|
|
|
func TestGetMongoCol(t *testing.T) {
|
|
// Test getting a MongoDB collection with default database
|
|
colName := "test_col"
|
|
|
|
// Get collection and verify the name matches
|
|
col := GetMongoCol(colName)
|
|
require.Equal(t, colName, col.c.Name())
|
|
}
|
|
|
|
func TestGetMongoColWithDb(t *testing.T) {
|
|
// Test getting a MongoDB collection with specified database
|
|
dbName := "test_db"
|
|
colName := "test_col"
|
|
|
|
// Get database instance first
|
|
db := GetMongoDb(dbName)
|
|
// Get collection with specific database and verify both names
|
|
col := GetMongoColWithDb(colName, db)
|
|
require.Equal(t, colName, col.c.Name())
|
|
require.Equal(t, dbName, col.db.Name())
|
|
}
|
|
|
|
func TestCol_Insert(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Insert a single document
|
|
id, err := to.col.Insert(bson.M{"key": "value"})
|
|
require.Nil(t, err)
|
|
require.IsType(t, primitive.ObjectID{}, id)
|
|
|
|
// Verify the inserted document
|
|
var doc map[string]string
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
require.Equal(t, doc["key"], "value")
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_InsertMany(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Create test data: 10 documents with incrementing values
|
|
n := 10
|
|
var docs []interface{}
|
|
for i := 0; i < n; i++ {
|
|
docs = append(docs, bson.M{"key": fmt.Sprintf("value-%d", i)})
|
|
}
|
|
|
|
// Insert multiple documents at once
|
|
ids, err := to.col.InsertMany(docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(ids))
|
|
|
|
// Verify all documents were inserted correctly
|
|
// Retrieve all documents sorted by _id
|
|
var resDocs []map[string]string
|
|
err = to.col.Find(nil, &FindOptions{Sort: bson.D{{"_id", 1}}}).All(&resDocs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(resDocs))
|
|
|
|
// Verify each document has correct value
|
|
for i, doc := range resDocs {
|
|
require.Equal(t, fmt.Sprintf("value-%d", i), doc["key"])
|
|
}
|
|
|
|
// Cleanup test environment
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_UpdateId(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Insert initial document
|
|
id, err := to.col.Insert(bson.M{"key": "old-value"})
|
|
require.Nil(t, err)
|
|
|
|
// Update document by ID using $set operator
|
|
err = to.col.UpdateId(id, bson.M{
|
|
"$set": bson.M{
|
|
"key": "new-value",
|
|
},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
// Verify the update
|
|
var doc map[string]string
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
require.Equal(t, "new-value", doc["key"])
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_Update(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Prepare test documents
|
|
n := 10
|
|
var docs []interface{}
|
|
for i := 0; i < n; i++ {
|
|
docs = append(docs, bson.M{"key": fmt.Sprintf("old-value-%d", i)})
|
|
}
|
|
|
|
// Update all documents (nil query means update all)
|
|
err = to.col.Update(nil, bson.M{
|
|
"$set": bson.M{
|
|
"key": "new-value",
|
|
},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
// Verify all documents were updated
|
|
var resDocs []map[string]string
|
|
err = to.col.Find(nil, &FindOptions{Sort: bson.D{{"_id", 1}}}).All(&resDocs)
|
|
require.Nil(t, err)
|
|
for _, doc := range resDocs {
|
|
require.Equal(t, "new-value", doc["key"])
|
|
}
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_ReplaceId(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Insert initial document
|
|
id, err := to.col.Insert(bson.M{"key": "old-value"})
|
|
require.Nil(t, err)
|
|
|
|
// Get the document and modify it
|
|
var doc map[string]interface{}
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
doc["key"] = "new-value"
|
|
|
|
// Replace the entire document
|
|
err = to.col.ReplaceId(id, doc)
|
|
require.Nil(t, err)
|
|
|
|
// Verify the replacement
|
|
err = to.col.FindId(id).One(doc)
|
|
require.Nil(t, err)
|
|
require.Equal(t, "new-value", doc["key"])
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_Replace(t *testing.T) {
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
id, err := to.col.Insert(bson.M{"key": "old-value"})
|
|
require.Nil(t, err)
|
|
|
|
var doc map[string]interface{}
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
doc["key"] = "new-value"
|
|
|
|
err = to.col.Replace(bson.M{"key": "old-value"}, doc)
|
|
require.Nil(t, err)
|
|
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
require.Equal(t, "new-value", doc["key"])
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_DeleteId(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Insert a document to be deleted
|
|
id, err := to.col.Insert(bson.M{"key": "value"})
|
|
require.Nil(t, err)
|
|
|
|
// Delete the document by ID
|
|
err = to.col.DeleteId(id)
|
|
require.Nil(t, err)
|
|
|
|
// Verify deletion by checking collection count
|
|
total, err := to.col.Count(nil)
|
|
require.Nil(t, err)
|
|
require.Equal(t, 0, total)
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_Delete(t *testing.T) {
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
n := 10
|
|
var docs []interface{}
|
|
for i := 0; i < n; i++ {
|
|
docs = append(docs, bson.M{"key": fmt.Sprintf("value-%d", i)})
|
|
}
|
|
ids, err := to.col.InsertMany(docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(ids))
|
|
|
|
err = to.col.Delete(bson.M{"key": "value-0"})
|
|
require.Nil(t, err)
|
|
|
|
total, err := to.col.Count(nil)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n-1, total)
|
|
|
|
err = to.col.Delete(nil)
|
|
require.Nil(t, err)
|
|
|
|
total, err = to.col.Count(nil)
|
|
require.Nil(t, err)
|
|
require.Equal(t, 0, total)
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_FindId(t *testing.T) {
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
id, err := to.col.Insert(bson.M{"key": "value"})
|
|
require.Nil(t, err)
|
|
|
|
var doc map[string]string
|
|
err = to.col.FindId(id).One(&doc)
|
|
require.Nil(t, err)
|
|
require.Equal(t, "value", doc["key"])
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_Find(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Insert test documents
|
|
n := 10
|
|
var docs []interface{}
|
|
for i := 0; i < n; i++ {
|
|
docs = append(docs, TestDocument{
|
|
Key: fmt.Sprintf("value-%d", i),
|
|
Tags: []string{"test tag"},
|
|
})
|
|
}
|
|
ids, err := to.col.InsertMany(docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(ids))
|
|
|
|
// Test 1: Find all documents
|
|
err = to.col.Find(nil, nil).All(&docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(docs))
|
|
|
|
// Test 2: Find documents with conditional query ($gte)
|
|
err = to.col.Find(bson.M{"key": bson.M{"$gte": fmt.Sprintf("value-%d", 5)}}, nil).All(&docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n-5, len(docs))
|
|
|
|
// Test 3: Test skip functionality
|
|
err = to.col.Find(nil, &FindOptions{
|
|
Skip: 5,
|
|
}).All(&docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n-5, len(docs))
|
|
|
|
// Test 4: Test limit functionality
|
|
err = to.col.Find(nil, &FindOptions{
|
|
Limit: 5,
|
|
}).All(&docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, 5, len(docs))
|
|
|
|
// Test 5: Test ascending sort
|
|
var resDocs []TestDocument
|
|
err = to.col.Find(nil, &FindOptions{
|
|
Sort: bson.D{{"key", 1}},
|
|
}).All(&resDocs)
|
|
require.Nil(t, err)
|
|
require.Greater(t, len(resDocs), 0)
|
|
require.Equal(t, "value-0", resDocs[0].Key)
|
|
|
|
// Test 6: Test descending sort
|
|
err = to.col.Find(nil, &FindOptions{
|
|
Sort: bson.D{{"key", -1}},
|
|
}).All(&resDocs)
|
|
require.Nil(t, err)
|
|
require.Greater(t, len(resDocs), 0)
|
|
require.Equal(t, fmt.Sprintf("value-%d", n-1), resDocs[0].Key)
|
|
|
|
// Test 7: Test array query with $in operator
|
|
var resDocs2 []TestDocument
|
|
err = to.col.Find(bson.M{"tags": bson.M{"$in": []string{"test tag"}}}, nil).All(&resDocs2)
|
|
require.Nil(t, err)
|
|
require.Greater(t, len(resDocs2), 0)
|
|
|
|
// Cleanup test environment
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_CreateIndex(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Create a single index on the 'key' field
|
|
err = to.col.CreateIndex(mongo.IndexModel{
|
|
Keys: bson.D{{"key", 1}},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
// Verify index creation
|
|
// Should have 2 indexes: default _id index and our new index
|
|
indexes, err := to.col.ListIndexes()
|
|
require.Nil(t, err)
|
|
require.Equal(t, 2, len(indexes))
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_Aggregate(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Define aggregation pipeline
|
|
var results []TestAggregateResult
|
|
pipeline := mongo.Pipeline{
|
|
{
|
|
{
|
|
"$group",
|
|
bson.D{
|
|
{"_id", "$key"},
|
|
{
|
|
"count",
|
|
bson.D{{"$sum", 1}},
|
|
},
|
|
{
|
|
"value",
|
|
bson.D{{"$sum", "$value"}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
{
|
|
"$sort",
|
|
bson.D{{"_id", 1}},
|
|
},
|
|
},
|
|
}
|
|
|
|
// Test empty collection aggregation
|
|
err = to.col.Aggregate(pipeline, nil).All(&results)
|
|
assert.Nil(t, err)
|
|
|
|
// Insert test documents
|
|
n := 10 // total documents
|
|
v := 2 // value for each document
|
|
var docs []interface{}
|
|
for i := 0; i < n; i++ {
|
|
docs = append(docs, TestDocument{
|
|
Key: fmt.Sprintf("%d", i%2), // alternates between "0" and "1"
|
|
Value: v,
|
|
})
|
|
}
|
|
ids, err := to.col.InsertMany(docs)
|
|
require.Nil(t, err)
|
|
require.Equal(t, n, len(ids))
|
|
|
|
// Run aggregation on populate
|
|
err = to.col.Aggregate(pipeline, nil).All(&results)
|
|
require.Nil(t, err)
|
|
require.Equal(t, 2, len(results))
|
|
|
|
for i, r := range results {
|
|
require.Equal(t, strconv.Itoa(i), r.Id)
|
|
require.Equal(t, n/2, r.Count)
|
|
require.Equal(t, n*v/2, r.Value)
|
|
}
|
|
}
|
|
|
|
func TestCol_CreateIndexes(t *testing.T) {
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
indexes, err := to.col.ListIndexes()
|
|
assert.Nil(t, err)
|
|
|
|
err = to.col.CreateIndexes([]mongo.IndexModel{
|
|
{
|
|
Keys: bson.D{{"key", 1}},
|
|
},
|
|
{
|
|
Keys: bson.D{{"empty-key", 1}},
|
|
},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
indexes, err = to.col.ListIndexes()
|
|
require.Nil(t, err)
|
|
require.Equal(t, 3, len(indexes))
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_DeleteIndex(t *testing.T) {
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
err = to.col.DeleteIndex("non-existent-index")
|
|
assert.Nil(t, err)
|
|
|
|
err = to.col.CreateIndex(mongo.IndexModel{
|
|
Keys: bson.D{{"key", 1}},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
indexes, err := to.col.ListIndexes()
|
|
require.Nil(t, err)
|
|
assert.Equal(t, 2, len(indexes))
|
|
for _, index := range indexes {
|
|
name, ok := index["name"].(string)
|
|
assert.True(t, ok)
|
|
|
|
if name != "_id_" {
|
|
err = to.col.DeleteIndex(name)
|
|
assert.Nil(t, err)
|
|
break
|
|
}
|
|
}
|
|
|
|
indexes, err = to.col.ListIndexes()
|
|
require.Nil(t, err)
|
|
assert.Equal(t, 1, len(indexes))
|
|
|
|
cleanupColTest(to)
|
|
}
|
|
|
|
func TestCol_DeleteIndexes(t *testing.T) {
|
|
// Setup test environment
|
|
to, err := setupColTest()
|
|
require.Nil(t, err)
|
|
|
|
// Delete all existing indexes
|
|
err = to.col.DeleteAllIndexes()
|
|
assert.Nil(t, err)
|
|
|
|
// Create multiple test indexes
|
|
err = to.col.CreateIndexes([]mongo.IndexModel{
|
|
{
|
|
Keys: bson.D{{"key", 1}},
|
|
},
|
|
{
|
|
Keys: bson.D{{"empty-key", 1}},
|
|
},
|
|
})
|
|
require.Nil(t, err)
|
|
|
|
// Delete all indexes except the default _id index
|
|
err = to.col.DeleteAllIndexes()
|
|
require.Nil(t, err)
|
|
|
|
// Verify only default _id index remains
|
|
indexes, err := to.col.ListIndexes()
|
|
require.Nil(t, err)
|
|
require.Equal(t, 1, len(indexes))
|
|
|
|
cleanupColTest(to)
|
|
}
|