collection_test.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667
  1. package mon
  2. import (
  3. "context"
  4. "errors"
  5. "strings"
  6. "testing"
  7. "github.com/stretchr/testify/assert"
  8. "github.com/zeromicro/go-zero/core/breaker"
  9. "github.com/zeromicro/go-zero/core/logx"
  10. "github.com/zeromicro/go-zero/core/stringx"
  11. "github.com/zeromicro/go-zero/core/timex"
  12. "go.mongodb.org/mongo-driver/bson"
  13. "go.mongodb.org/mongo-driver/bson/primitive"
  14. "go.mongodb.org/mongo-driver/mongo"
  15. "go.mongodb.org/mongo-driver/mongo/integration/mtest"
  16. mopt "go.mongodb.org/mongo-driver/mongo/options"
  17. )
  18. var errDummy = errors.New("dummy")
  19. func init() {
  20. logx.Disable()
  21. }
  22. func TestKeepPromise_accept(t *testing.T) {
  23. p := new(mockPromise)
  24. kp := keepablePromise{
  25. promise: p,
  26. log: func(error) {},
  27. }
  28. assert.Nil(t, kp.accept(nil))
  29. assert.Equal(t, ErrNotFound, kp.accept(ErrNotFound))
  30. }
  31. func TestKeepPromise_keep(t *testing.T) {
  32. tests := []struct {
  33. err error
  34. accepted bool
  35. reason string
  36. }{
  37. {
  38. err: nil,
  39. accepted: true,
  40. reason: "",
  41. },
  42. {
  43. err: ErrNotFound,
  44. accepted: true,
  45. reason: "",
  46. },
  47. {
  48. err: errors.New("any"),
  49. accepted: false,
  50. reason: "any",
  51. },
  52. }
  53. for _, test := range tests {
  54. t.Run(stringx.RandId(), func(t *testing.T) {
  55. p := new(mockPromise)
  56. kp := keepablePromise{
  57. promise: p,
  58. log: func(error) {},
  59. }
  60. assert.Equal(t, test.err, kp.keep(test.err))
  61. assert.Equal(t, test.accepted, p.accepted)
  62. assert.Equal(t, test.reason, p.reason)
  63. })
  64. }
  65. }
  66. func TestNewCollection(t *testing.T) {
  67. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  68. defer mt.Close()
  69. mt.Run("test", func(mt *mtest.T) {
  70. coll := mt.Coll
  71. assert.NotNil(t, coll)
  72. col := newCollection(coll, breaker.GetBreaker("localhost"))
  73. assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
  74. })
  75. }
  76. func TestCollection_Aggregate(t *testing.T) {
  77. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  78. defer mt.Close()
  79. mt.Run("test", func(mt *mtest.T) {
  80. coll := mt.Coll
  81. assert.NotNil(t, coll)
  82. col := newCollection(coll, breaker.GetBreaker("localhost"))
  83. ns := mt.Coll.Database().Name() + "." + mt.Coll.Name()
  84. aggRes := mtest.CreateCursorResponse(1, ns, mtest.FirstBatch)
  85. mt.AddMockResponses(aggRes)
  86. assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
  87. cursor, err := col.Aggregate(context.Background(), mongo.Pipeline{}, mopt.Aggregate())
  88. assert.Nil(t, err)
  89. cursor.Close(context.Background())
  90. })
  91. }
  92. func TestCollection_BulkWrite(t *testing.T) {
  93. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  94. defer mt.Close()
  95. mt.Run("test", func(mt *mtest.T) {
  96. c := decoratedCollection{
  97. Collection: mt.Coll,
  98. brk: breaker.NewBreaker(),
  99. }
  100. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
  101. res, err := c.BulkWrite(context.Background(), []mongo.WriteModel{
  102. mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}})},
  103. )
  104. assert.Nil(t, err)
  105. assert.NotNil(t, res)
  106. c.brk = new(dropBreaker)
  107. _, err = c.BulkWrite(context.Background(), []mongo.WriteModel{
  108. mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}})},
  109. )
  110. assert.Equal(t, errDummy, err)
  111. })
  112. }
  113. func TestCollection_CountDocuments(t *testing.T) {
  114. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  115. defer mt.Close()
  116. mt.Run("test", func(mt *mtest.T) {
  117. c := decoratedCollection{
  118. Collection: mt.Coll,
  119. brk: breaker.NewBreaker(),
  120. }
  121. mt.AddMockResponses(mtest.CreateCursorResponse(
  122. 1,
  123. "DBName.CollectionName",
  124. mtest.FirstBatch,
  125. bson.D{
  126. {Key: "n", Value: 1},
  127. }))
  128. res, err := c.CountDocuments(context.Background(), bson.D{})
  129. assert.Nil(t, err)
  130. assert.Equal(t, int64(1), res)
  131. c.brk = new(dropBreaker)
  132. _, err = c.CountDocuments(context.Background(), bson.D{{Key: "foo", Value: 1}})
  133. assert.Equal(t, errDummy, err)
  134. })
  135. }
  136. func TestDecoratedCollection_DeleteMany(t *testing.T) {
  137. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  138. defer mt.Close()
  139. mt.Run("test", func(mt *mtest.T) {
  140. c := decoratedCollection{
  141. Collection: mt.Coll,
  142. brk: breaker.NewBreaker(),
  143. }
  144. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  145. res, err := c.DeleteMany(context.Background(), bson.D{})
  146. assert.Nil(t, err)
  147. assert.Equal(t, int64(1), res.DeletedCount)
  148. c.brk = new(dropBreaker)
  149. _, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: 1}})
  150. assert.Equal(t, errDummy, err)
  151. })
  152. }
  153. func TestCollection_Distinct(t *testing.T) {
  154. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  155. defer mt.Close()
  156. mt.Run("test", func(mt *mtest.T) {
  157. c := decoratedCollection{
  158. Collection: mt.Coll,
  159. brk: breaker.NewBreaker(),
  160. }
  161. mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "values", Value: []int{1}}})
  162. resp, err := c.Distinct(context.Background(), "foo", bson.D{})
  163. assert.Nil(t, err)
  164. assert.Equal(t, 1, len(resp))
  165. c.brk = new(dropBreaker)
  166. _, err = c.Distinct(context.Background(), "foo", bson.D{{Key: "foo", Value: 1}})
  167. assert.Equal(t, errDummy, err)
  168. })
  169. }
  170. func TestCollection_EstimatedDocumentCount(t *testing.T) {
  171. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  172. defer mt.Close()
  173. mt.Run("test", func(mt *mtest.T) {
  174. c := decoratedCollection{
  175. Collection: mt.Coll,
  176. brk: breaker.NewBreaker(),
  177. }
  178. mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "n", Value: 1}})
  179. res, err := c.EstimatedDocumentCount(context.Background())
  180. assert.Nil(t, err)
  181. assert.Equal(t, int64(1), res)
  182. c.brk = new(dropBreaker)
  183. _, err = c.EstimatedDocumentCount(context.Background())
  184. assert.Equal(t, errDummy, err)
  185. })
  186. }
  187. func TestCollectionFind(t *testing.T) {
  188. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  189. defer mt.Close()
  190. mt.Run("test", func(mt *mtest.T) {
  191. c := decoratedCollection{
  192. Collection: mt.Coll,
  193. brk: breaker.NewBreaker(),
  194. }
  195. find := mtest.CreateCursorResponse(
  196. 1,
  197. "DBName.CollectionName",
  198. mtest.FirstBatch,
  199. bson.D{
  200. {Key: "name", Value: "John"},
  201. })
  202. getMore := mtest.CreateCursorResponse(
  203. 1,
  204. "DBName.CollectionName",
  205. mtest.NextBatch,
  206. bson.D{
  207. {Key: "name", Value: "Mary"},
  208. })
  209. killCursors := mtest.CreateCursorResponse(
  210. 0,
  211. "DBName.CollectionName",
  212. mtest.NextBatch)
  213. mt.AddMockResponses(find, getMore, killCursors)
  214. filter := bson.D{{Key: "x", Value: 1}}
  215. cursor, err := c.Find(context.Background(), filter, mopt.Find())
  216. assert.Nil(t, err)
  217. defer cursor.Close(context.Background())
  218. var val []struct {
  219. ID primitive.ObjectID `bson:"_id"`
  220. Name string `bson:"name"`
  221. }
  222. assert.Nil(t, cursor.All(context.Background(), &val))
  223. assert.Equal(t, 2, len(val))
  224. assert.Equal(t, "John", val[0].Name)
  225. assert.Equal(t, "Mary", val[1].Name)
  226. c.brk = new(dropBreaker)
  227. _, err = c.Find(context.Background(), filter, mopt.Find())
  228. assert.Equal(t, errDummy, err)
  229. })
  230. }
  231. func TestCollectionFindOne(t *testing.T) {
  232. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  233. defer mt.Close()
  234. mt.Run("test", func(mt *mtest.T) {
  235. c := decoratedCollection{
  236. Collection: mt.Coll,
  237. brk: breaker.NewBreaker(),
  238. }
  239. find := mtest.CreateCursorResponse(
  240. 1,
  241. "DBName.CollectionName",
  242. mtest.FirstBatch,
  243. bson.D{
  244. {Key: "name", Value: "John"},
  245. })
  246. getMore := mtest.CreateCursorResponse(
  247. 1,
  248. "DBName.CollectionName",
  249. mtest.NextBatch,
  250. bson.D{
  251. {Key: "name", Value: "Mary"},
  252. })
  253. killCursors := mtest.CreateCursorResponse(
  254. 0,
  255. "DBName.CollectionName",
  256. mtest.NextBatch)
  257. mt.AddMockResponses(find, getMore, killCursors)
  258. filter := bson.D{{Key: "x", Value: 1}}
  259. resp, err := c.FindOne(context.Background(), filter)
  260. assert.Nil(t, err)
  261. var val struct {
  262. ID primitive.ObjectID `bson:"_id"`
  263. Name string `bson:"name"`
  264. }
  265. assert.Nil(t, resp.Decode(&val))
  266. assert.Equal(t, "John", val.Name)
  267. c.brk = new(dropBreaker)
  268. _, err = c.FindOne(context.Background(), filter)
  269. assert.Equal(t, errDummy, err)
  270. })
  271. }
  272. func TestCollection_FindOneAndDelete(t *testing.T) {
  273. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  274. defer mt.Close()
  275. mt.Run("test", func(mt *mtest.T) {
  276. c := decoratedCollection{
  277. Collection: mt.Coll,
  278. brk: breaker.NewBreaker(),
  279. }
  280. filter := bson.D{}
  281. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
  282. _, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
  283. assert.Equal(t, mongo.ErrNoDocuments, err)
  284. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
  285. {Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
  286. }...))
  287. resp, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
  288. assert.Nil(t, err)
  289. var val struct {
  290. Name string `bson:"name"`
  291. }
  292. assert.Nil(t, resp.Decode(&val))
  293. assert.Equal(t, "John", val.Name)
  294. c.brk = new(dropBreaker)
  295. _, err = c.FindOneAndDelete(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  296. assert.Equal(t, errDummy, err)
  297. })
  298. }
  299. func TestCollection_FindOneAndReplace(t *testing.T) {
  300. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  301. defer mt.Close()
  302. mt.Run("test", func(mt *mtest.T) {
  303. c := decoratedCollection{
  304. Collection: mt.Coll,
  305. brk: breaker.NewBreaker(),
  306. }
  307. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
  308. filter := bson.D{{Key: "x", Value: 1}}
  309. replacement := bson.D{{Key: "x", Value: 2}}
  310. opts := mopt.FindOneAndReplace().SetUpsert(true)
  311. _, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
  312. assert.Equal(t, mongo.ErrNoDocuments, err)
  313. mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
  314. {Key: "name", Value: "John"},
  315. }}})
  316. resp, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
  317. assert.Nil(t, err)
  318. var val struct {
  319. Name string `bson:"name"`
  320. }
  321. assert.Nil(t, resp.Decode(&val))
  322. assert.Equal(t, "John", val.Name)
  323. c.brk = new(dropBreaker)
  324. _, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
  325. assert.Equal(t, errDummy, err)
  326. })
  327. }
  328. func TestCollection_FindOneAndUpdate(t *testing.T) {
  329. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  330. defer mt.Close()
  331. mt.Run("test", func(mt *mtest.T) {
  332. c := decoratedCollection{
  333. Collection: mt.Coll,
  334. brk: breaker.NewBreaker(),
  335. }
  336. mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}})
  337. filter := bson.D{{Key: "x", Value: 1}}
  338. update := bson.D{{Key: "$x", Value: 2}}
  339. opts := mopt.FindOneAndUpdate().SetUpsert(true)
  340. _, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
  341. assert.Equal(t, mongo.ErrNoDocuments, err)
  342. mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
  343. {Key: "name", Value: "John"},
  344. }}})
  345. resp, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
  346. assert.Nil(t, err)
  347. var val struct {
  348. Name string `bson:"name"`
  349. }
  350. assert.Nil(t, resp.Decode(&val))
  351. assert.Equal(t, "John", val.Name)
  352. c.brk = new(dropBreaker)
  353. _, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
  354. assert.Equal(t, errDummy, err)
  355. })
  356. }
  357. func TestCollection_InsertOne(t *testing.T) {
  358. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  359. defer mt.Close()
  360. mt.Run("test", func(mt *mtest.T) {
  361. c := decoratedCollection{
  362. Collection: mt.Coll,
  363. brk: breaker.NewBreaker(),
  364. }
  365. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
  366. res, err := c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  367. assert.Nil(t, err)
  368. assert.NotNil(t, res)
  369. c.brk = new(dropBreaker)
  370. _, err = c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  371. assert.Equal(t, errDummy, err)
  372. })
  373. }
  374. func TestCollection_InsertMany(t *testing.T) {
  375. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  376. defer mt.Close()
  377. mt.Run("test", func(mt *mtest.T) {
  378. c := decoratedCollection{
  379. Collection: mt.Coll,
  380. brk: breaker.NewBreaker(),
  381. }
  382. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
  383. res, err := c.InsertMany(context.Background(), []interface{}{
  384. bson.D{{Key: "foo", Value: "bar"}},
  385. bson.D{{Key: "foo", Value: "baz"}},
  386. })
  387. assert.Nil(t, err)
  388. assert.NotNil(t, res)
  389. assert.Equal(t, 2, len(res.InsertedIDs))
  390. c.brk = new(dropBreaker)
  391. _, err = c.InsertMany(context.Background(), []interface{}{bson.D{{Key: "foo", Value: "bar"}}})
  392. assert.Equal(t, errDummy, err)
  393. })
  394. }
  395. func TestCollection_Remove(t *testing.T) {
  396. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  397. defer mt.Close()
  398. mt.Run("test", func(mt *mtest.T) {
  399. c := decoratedCollection{
  400. Collection: mt.Coll,
  401. brk: breaker.NewBreaker(),
  402. }
  403. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  404. res, err := c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  405. assert.Nil(t, err)
  406. assert.Equal(t, int64(1), res.DeletedCount)
  407. c.brk = new(dropBreaker)
  408. _, err = c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  409. assert.Equal(t, errDummy, err)
  410. })
  411. }
  412. func TestCollectionRemoveAll(t *testing.T) {
  413. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  414. defer mt.Close()
  415. mt.Run("test", func(mt *mtest.T) {
  416. c := decoratedCollection{
  417. Collection: mt.Coll,
  418. brk: breaker.NewBreaker(),
  419. }
  420. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  421. res, err := c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  422. assert.Nil(t, err)
  423. assert.Equal(t, int64(1), res.DeletedCount)
  424. c.brk = new(dropBreaker)
  425. _, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
  426. assert.Equal(t, errDummy, err)
  427. })
  428. }
  429. func TestCollection_ReplaceOne(t *testing.T) {
  430. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  431. defer mt.Close()
  432. mt.Run("test", func(mt *mtest.T) {
  433. c := decoratedCollection{
  434. Collection: mt.Coll,
  435. brk: breaker.NewBreaker(),
  436. }
  437. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  438. res, err := c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  439. bson.D{{Key: "foo", Value: "baz"}},
  440. )
  441. assert.Nil(t, err)
  442. assert.Equal(t, int64(1), res.MatchedCount)
  443. c.brk = new(dropBreaker)
  444. _, err = c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  445. bson.D{{Key: "foo", Value: "baz"}})
  446. assert.Equal(t, errDummy, err)
  447. })
  448. }
  449. func TestCollection_UpdateOne(t *testing.T) {
  450. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  451. defer mt.Close()
  452. mt.Run("test", func(mt *mtest.T) {
  453. c := decoratedCollection{
  454. Collection: mt.Coll,
  455. brk: breaker.NewBreaker(),
  456. }
  457. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  458. resp, err := c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  459. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  460. assert.Nil(t, err)
  461. assert.Equal(t, int64(1), resp.MatchedCount)
  462. c.brk = new(dropBreaker)
  463. _, err = c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  464. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  465. assert.Equal(t, errDummy, err)
  466. })
  467. }
  468. func TestCollection_UpdateByID(t *testing.T) {
  469. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  470. defer mt.Close()
  471. mt.Run("test", func(mt *mtest.T) {
  472. c := decoratedCollection{
  473. Collection: mt.Coll,
  474. brk: breaker.NewBreaker(),
  475. }
  476. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  477. resp, err := c.UpdateByID(context.Background(), primitive.NewObjectID(),
  478. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  479. assert.Nil(t, err)
  480. assert.Equal(t, int64(1), resp.MatchedCount)
  481. c.brk = new(dropBreaker)
  482. _, err = c.UpdateByID(context.Background(), primitive.NewObjectID(),
  483. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  484. assert.Equal(t, errDummy, err)
  485. })
  486. }
  487. func TestCollection_UpdateMany(t *testing.T) {
  488. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  489. defer mt.Close()
  490. mt.Run("test", func(mt *mtest.T) {
  491. c := decoratedCollection{
  492. Collection: mt.Coll,
  493. brk: breaker.NewBreaker(),
  494. }
  495. mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
  496. resp, err := c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  497. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  498. assert.Nil(t, err)
  499. assert.Equal(t, int64(1), resp.MatchedCount)
  500. c.brk = new(dropBreaker)
  501. _, err = c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
  502. bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
  503. assert.Equal(t, errDummy, err)
  504. })
  505. }
  506. func Test_DecoratedCollectionLogDuration(t *testing.T) {
  507. mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
  508. defer mt.Close()
  509. c := decoratedCollection{
  510. Collection: mt.Coll,
  511. brk: breaker.NewBreaker(),
  512. }
  513. var buf strings.Builder
  514. w := logx.NewWriter(&buf)
  515. o := logx.Reset()
  516. logx.SetWriter(w)
  517. defer func() {
  518. logx.Reset()
  519. logx.SetWriter(o)
  520. }()
  521. buf.Reset()
  522. c.logDuration(context.Background(), "foo", timex.Now(), nil, "bar")
  523. assert.Contains(t, buf.String(), "foo")
  524. assert.Contains(t, buf.String(), "bar")
  525. buf.Reset()
  526. c.logDuration(context.Background(), "foo", timex.Now(), errors.New("bar"), make(chan int))
  527. assert.Contains(t, buf.String(), "foo")
  528. assert.Contains(t, buf.String(), "bar")
  529. buf.Reset()
  530. c.logDuration(context.Background(), "foo", timex.Now(), nil, make(chan int))
  531. assert.Contains(t, buf.String(), "foo")
  532. buf.Reset()
  533. c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
  534. nil, make(chan int))
  535. assert.Contains(t, buf.String(), "foo")
  536. assert.Contains(t, buf.String(), "slowcall")
  537. buf.Reset()
  538. c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
  539. errors.New("bar"), make(chan int))
  540. assert.Contains(t, buf.String(), "foo")
  541. assert.Contains(t, buf.String(), "bar")
  542. assert.Contains(t, buf.String(), "slowcall")
  543. buf.Reset()
  544. c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
  545. errors.New("bar"))
  546. assert.Contains(t, buf.String(), "foo")
  547. assert.Contains(t, buf.String(), "slowcall")
  548. buf.Reset()
  549. c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2, nil)
  550. assert.Contains(t, buf.String(), "foo")
  551. assert.Contains(t, buf.String(), "slowcall")
  552. }
  553. type mockPromise struct {
  554. accepted bool
  555. reason string
  556. }
  557. func (p *mockPromise) Accept() {
  558. p.accepted = true
  559. }
  560. func (p *mockPromise) Reject(reason string) {
  561. p.reason = reason
  562. }
  563. type dropBreaker struct{}
  564. func (d *dropBreaker) Name() string {
  565. return "dummy"
  566. }
  567. func (d *dropBreaker) Allow() (breaker.Promise, error) {
  568. return nil, errDummy
  569. }
  570. func (d *dropBreaker) Do(_ func() error) error {
  571. return nil
  572. }
  573. func (d *dropBreaker) DoWithAcceptable(_ func() error, _ breaker.Acceptable) error {
  574. return errDummy
  575. }
  576. func (d *dropBreaker) DoWithFallback(_ func() error, _ func(err error) error) error {
  577. return nil
  578. }
  579. func (d *dropBreaker) DoWithFallbackAcceptable(_ func() error, _ func(err error) error,
  580. _ breaker.Acceptable) error {
  581. return nil
  582. }