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