1// Unless explicitly stated otherwise all files in this repository are licensed 2// under the Apache License Version 2.0. 3// This product includes software developed at Datadog (https://www.datadoghq.com/). 4// Copyright 2016 Datadog, Inc. 5 6package mgo 7 8import ( 9 "context" 10 "fmt" 11 "os" 12 "testing" 13 14 "github.com/globalsign/mgo" 15 "github.com/globalsign/mgo/bson" 16 "github.com/stretchr/testify/assert" 17 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" 18 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" 19 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" 20 "gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig" 21) 22 23func TestMain(m *testing.M) { 24 _, ok := os.LookupEnv("INTEGRATION") 25 if !ok { 26 fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable") 27 os.Exit(0) 28 } 29 os.Exit(m.Run()) 30} 31 32func testMongoCollectionCommand(assert *assert.Assertions, command func(*Collection)) []mocktracer.Span { 33 mt := mocktracer.Start() 34 defer mt.Stop() 35 36 parentSpan, ctx := tracer.StartSpanFromContext( 37 context.Background(), 38 "mgo-unittest", 39 tracer.SpanType("app"), 40 tracer.ResourceName("insert-test"), 41 ) 42 43 session, err := Dial("localhost:27017", WithServiceName("unit-tests"), WithContext(ctx)) 44 defer session.Close() 45 46 assert.NotNil(session) 47 assert.Nil(err) 48 49 db := session.DB("my_db") 50 collection := db.C("MyCollection") 51 52 command(collection) 53 54 parentSpan.Finish() 55 56 spans := mt.FinishedSpans() 57 return spans 58} 59 60func TestCollection_Insert(t *testing.T) { 61 assert := assert.New(t) 62 63 entity := bson.D{ 64 bson.DocElem{ 65 Name: "entity", 66 Value: bson.DocElem{ 67 Name: "index", 68 Value: 0}}} 69 70 insert := func(collection *Collection) { 71 collection.Insert(entity) 72 } 73 74 spans := testMongoCollectionCommand(assert, insert) 75 assert.Equal(2, len(spans)) 76 assert.Equal("mongodb.query", spans[0].OperationName()) 77} 78 79func TestCollection_Update(t *testing.T) { 80 assert := assert.New(t) 81 82 entity := bson.D{ 83 bson.DocElem{ 84 Name: "entity", 85 Value: bson.DocElem{ 86 Name: "index", 87 Value: 0}}} 88 89 insert := func(collection *Collection) { 90 collection.Insert(entity) 91 collection.Update(entity, entity) 92 } 93 94 spans := testMongoCollectionCommand(assert, insert) 95 assert.Equal(3, len(spans)) 96 assert.Equal("mongodb.query", spans[1].OperationName()) 97} 98 99func TestCollection_UpdateId(t *testing.T) { 100 assert := assert.New(t) 101 102 entity := bson.D{ 103 bson.DocElem{ 104 Name: "entity", 105 Value: bson.DocElem{ 106 Name: "index", 107 Value: 0}}} 108 109 insert := func(collection *Collection) { 110 collection.Insert(entity) 111 var r bson.D 112 collection.Find(entity).Iter().Next(&r) 113 collection.UpdateId(r.Map()["_id"], entity) 114 } 115 116 spans := testMongoCollectionCommand(assert, insert) 117 assert.Equal(5, len(spans)) 118 assert.Equal("mongodb.query", spans[3].OperationName()) 119} 120 121func TestIssue874(t *testing.T) { 122 // regression test for DataDog/dd-trace-go#873 123 assert := assert.New(t) 124 125 entity := bson.D{ 126 bson.DocElem{ 127 Name: "entity", 128 Value: bson.DocElem{ 129 Name: "index", 130 Value: 0}}} 131 132 insert := func(collection *Collection) { 133 collection.Insert(entity) 134 var r bson.D 135 collection.Find(entity).All(&r) 136 collection.Find(entity).Apply(mgo.Change{Update: entity}, &r) 137 collection.Find(entity).Count() 138 collection.Find(entity).Distinct("index", &r) 139 collection.Find(entity).Explain(&r) 140 collection.Find(entity).One(&r) 141 collection.UpdateId(r.Map()["_id"], entity) 142 } 143 144 spans := testMongoCollectionCommand(assert, insert) 145 assert.Equal(9, len(spans)) 146} 147 148func TestCollection_Upsert(t *testing.T) { 149 assert := assert.New(t) 150 151 entity := bson.D{ 152 bson.DocElem{ 153 Name: "entity", 154 Value: bson.DocElem{ 155 Name: "index", 156 Value: 0}}} 157 158 insert := func(collection *Collection) { 159 collection.Insert(entity) 160 collection.Upsert(entity, entity) 161 var r bson.D 162 collection.Find(entity).Iter().Next(&r) 163 collection.UpsertId(r.Map()["_id"], entity) 164 } 165 166 spans := testMongoCollectionCommand(assert, insert) 167 assert.Equal(6, len(spans)) 168 assert.Equal("mongodb.query", spans[1].OperationName()) 169 assert.Equal("mongodb.query", spans[4].OperationName()) 170} 171 172func TestCollection_UpdateAll(t *testing.T) { 173 assert := assert.New(t) 174 175 entity := bson.D{ 176 bson.DocElem{ 177 Name: "entity", 178 Value: bson.DocElem{ 179 Name: "index", 180 Value: 0}}} 181 182 insert := func(collection *Collection) { 183 collection.Insert(entity) 184 collection.UpdateAll(entity, entity) 185 } 186 187 spans := testMongoCollectionCommand(assert, insert) 188 assert.Equal(3, len(spans)) 189 assert.Equal("mongodb.query", spans[1].OperationName()) 190} 191 192func TestCollection_FindId(t *testing.T) { 193 assert := assert.New(t) 194 195 entity := bson.D{ 196 bson.DocElem{ 197 Name: "entity", 198 Value: bson.DocElem{ 199 Name: "index", 200 Value: 0}}} 201 202 insert := func(collection *Collection) { 203 collection.Insert(entity) 204 var r bson.D 205 collection.Find(entity).Iter().Next(&r) 206 var r2 bson.D 207 collection.FindId(r.Map()["_id"]).Iter().Next(&r2) 208 } 209 210 spans := testMongoCollectionCommand(assert, insert) 211 assert.Equal(6, len(spans)) 212} 213 214func TestCollection_Remove(t *testing.T) { 215 assert := assert.New(t) 216 217 entity := bson.D{ 218 bson.DocElem{ 219 Name: "entity", 220 Value: bson.DocElem{ 221 Name: "index", 222 Value: 0}}} 223 224 insert := func(collection *Collection) { 225 collection.Insert(entity) 226 collection.Remove(entity) 227 } 228 229 spans := testMongoCollectionCommand(assert, insert) 230 assert.Equal(3, len(spans)) 231 assert.Equal("mongodb.query", spans[1].OperationName()) 232} 233 234func TestCollection_RemoveId(t *testing.T) { 235 assert := assert.New(t) 236 237 entity := bson.D{ 238 bson.DocElem{ 239 Name: "entity", 240 Value: bson.DocElem{ 241 Name: "index", 242 Value: 0}}} 243 244 removeByID := func(collection *Collection) { 245 collection.Insert(entity) 246 query := collection.Find(entity) 247 iter := query.Iter() 248 var r bson.D 249 iter.Next(&r) 250 id := r.Map()["_id"] 251 err := collection.RemoveId(id) 252 assert.NoError(err) 253 } 254 255 spans := testMongoCollectionCommand(assert, removeByID) 256 assert.Equal(5, len(spans)) 257 assert.Equal("mongodb.query", spans[3].OperationName()) 258} 259 260func TestCollection_RemoveAll(t *testing.T) { 261 assert := assert.New(t) 262 263 entity := bson.D{ 264 bson.DocElem{ 265 Name: "entity", 266 Value: bson.DocElem{ 267 Name: "index", 268 Value: 0}}} 269 270 insert := func(collection *Collection) { 271 collection.Insert(entity) 272 collection.RemoveAll(entity) 273 } 274 275 spans := testMongoCollectionCommand(assert, insert) 276 assert.Equal(3, len(spans)) 277 assert.Equal("mongodb.query", spans[1].OperationName()) 278} 279 280func TestCollection_DropCollection(t *testing.T) { 281 assert := assert.New(t) 282 283 insert := func(collection *Collection) { 284 collection.DropCollection() 285 } 286 287 spans := testMongoCollectionCommand(assert, insert) 288 assert.Equal(2, len(spans)) 289 assert.Equal("mongodb.query", spans[0].OperationName()) 290} 291 292func TestCollection_Create(t *testing.T) { 293 assert := assert.New(t) 294 295 insert := func(collection *Collection) { 296 collection.Create(&mgo.CollectionInfo{}) 297 } 298 299 spans := testMongoCollectionCommand(assert, insert) 300 assert.Equal(2, len(spans)) 301 assert.Equal("mongodb.query", spans[0].OperationName()) 302} 303 304func TestCollection_Count(t *testing.T) { 305 assert := assert.New(t) 306 307 insert := func(collection *Collection) { 308 collection.Count() 309 } 310 311 spans := testMongoCollectionCommand(assert, insert) 312 assert.Equal(2, len(spans)) 313 assert.Equal("mongodb.query", spans[0].OperationName()) 314} 315 316func TestCollection_IndexCommands(t *testing.T) { 317 assert := assert.New(t) 318 319 indexTest := func(collection *Collection) { 320 indexes, _ := collection.Indexes() 321 collection.DropIndex("_id_") 322 collection.DropIndexName("_id_") 323 collection.EnsureIndex(indexes[0]) 324 collection.EnsureIndexKey("_id_") 325 } 326 327 spans := testMongoCollectionCommand(assert, indexTest) 328 assert.Equal(6, len(spans)) 329 assert.Equal("mongodb.query", spans[0].OperationName()) 330 assert.Equal("mongodb.query", spans[1].OperationName()) 331 assert.Equal("mongodb.query", spans[2].OperationName()) 332 assert.Equal("mongodb.query", spans[3].OperationName()) 333 assert.Equal("mongodb.query", spans[4].OperationName()) 334 assert.Equal("mgo-unittest", spans[5].OperationName()) 335} 336 337func TestCollection_FindAndIter(t *testing.T) { 338 assert := assert.New(t) 339 340 entity := bson.D{ 341 bson.DocElem{ 342 Name: "entity", 343 Value: bson.DocElem{ 344 Name: "index", 345 Value: 0}}} 346 347 insert := func(collection *Collection) { 348 collection.Insert(entity) 349 collection.Insert(entity) 350 collection.Insert(entity) 351 352 query := collection.Find(nil) 353 iter := query.Iter() 354 var r bson.D 355 iter.Next(&r) 356 var all []bson.D 357 iter.All(&all) 358 iter.Close() 359 } 360 361 spans := testMongoCollectionCommand(assert, insert) 362 assert.Equal(8, len(spans)) 363 assert.Equal("mongodb.query", spans[3].OperationName()) 364 assert.Equal("mongodb.query", spans[4].OperationName()) 365 assert.Equal("mongodb.query", spans[5].OperationName()) 366 assert.Equal("mongodb.query", spans[6].OperationName()) 367} 368 369func TestCollection_Bulk(t *testing.T) { 370 assert := assert.New(t) 371 372 entity := bson.D{ 373 bson.DocElem{ 374 Name: "entity", 375 Value: bson.DocElem{ 376 Name: "index", 377 Value: 0}}} 378 379 insert := func(collection *Collection) { 380 bulk := collection.Bulk() 381 bulk.Insert(entity) 382 bulk.Run() 383 } 384 385 spans := testMongoCollectionCommand(assert, insert) 386 assert.Equal(2, len(spans)) 387 assert.Equal("mongodb.query", spans[0].OperationName()) 388} 389 390func TestAnalyticsSettings(t *testing.T) { 391 assertRate := func(t *testing.T, mt mocktracer.Tracer, rate interface{}, opts ...DialOption) { 392 assert := assert.New(t) 393 394 session, err := Dial("localhost:27017", opts...) 395 assert.NoError(err) 396 defer session.Close() 397 398 db := session.DB("my_db") 399 collection := db.C("MyCollection") 400 bulk := collection.Bulk() 401 bulk.Insert(bson.D{ 402 bson.DocElem{ 403 Name: "entity", 404 Value: bson.DocElem{ 405 Name: "index", 406 Value: 0, 407 }, 408 }, 409 }) 410 bulk.Run() 411 412 spans := mt.FinishedSpans() 413 assert.Len(spans, 1) 414 s := spans[0] 415 assert.Equal(rate, s.Tag(ext.EventSampleRate)) 416 } 417 418 t.Run("defaults", func(t *testing.T) { 419 mt := mocktracer.Start() 420 defer mt.Stop() 421 422 assertRate(t, mt, nil) 423 }) 424 425 t.Run("global", func(t *testing.T) { 426 t.Skip("global flag disabled") 427 mt := mocktracer.Start() 428 defer mt.Stop() 429 430 rate := globalconfig.AnalyticsRate() 431 defer globalconfig.SetAnalyticsRate(rate) 432 globalconfig.SetAnalyticsRate(0.4) 433 434 assertRate(t, mt, 0.4) 435 }) 436 437 t.Run("enabled", func(t *testing.T) { 438 mt := mocktracer.Start() 439 defer mt.Stop() 440 441 assertRate(t, mt, 1.0, WithAnalytics(true)) 442 }) 443 444 t.Run("disabled", func(t *testing.T) { 445 mt := mocktracer.Start() 446 defer mt.Stop() 447 448 assertRate(t, mt, nil, WithAnalytics(false)) 449 }) 450 451 t.Run("override", func(t *testing.T) { 452 mt := mocktracer.Start() 453 defer mt.Stop() 454 455 rate := globalconfig.AnalyticsRate() 456 defer globalconfig.SetAnalyticsRate(rate) 457 globalconfig.SetAnalyticsRate(0.4) 458 459 assertRate(t, mt, 0.23, WithAnalyticsRate(0.23)) 460 }) 461} 462