iter_test.go 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264
  1. package mongo
  2. import (
  3. "errors"
  4. "testing"
  5. "github.com/globalsign/mgo"
  6. "github.com/golang/mock/gomock"
  7. "github.com/stretchr/testify/assert"
  8. "github.com/tal-tech/go-zero/core/breaker"
  9. "github.com/tal-tech/go-zero/core/stringx"
  10. "github.com/tal-tech/go-zero/core/syncx"
  11. )
  12. func TestClosableIter_Close(t *testing.T) {
  13. errs := []error{
  14. nil,
  15. mgo.ErrNotFound,
  16. }
  17. for _, err := range errs {
  18. t.Run(stringx.RandId(), func(t *testing.T) {
  19. ctrl := gomock.NewController(t)
  20. defer ctrl.Finish()
  21. cleaned := syncx.NewAtomicBool()
  22. iter := NewMockIter(ctrl)
  23. iter.EXPECT().Close().Return(err)
  24. ci := ClosableIter{
  25. Iter: iter,
  26. Cleanup: func() {
  27. cleaned.Set(true)
  28. },
  29. }
  30. assert.Equal(t, err, ci.Close())
  31. assert.True(t, cleaned.True())
  32. })
  33. }
  34. }
  35. func TestPromisedIter_AllAndClose(t *testing.T) {
  36. tests := []struct {
  37. err error
  38. accepted bool
  39. reason string
  40. }{
  41. {
  42. err: nil,
  43. accepted: true,
  44. reason: "",
  45. },
  46. {
  47. err: mgo.ErrNotFound,
  48. accepted: true,
  49. reason: "",
  50. },
  51. {
  52. err: errors.New("any"),
  53. accepted: false,
  54. reason: "any",
  55. },
  56. }
  57. for _, test := range tests {
  58. t.Run(stringx.RandId(), func(t *testing.T) {
  59. ctrl := gomock.NewController(t)
  60. defer ctrl.Finish()
  61. iter := NewMockIter(ctrl)
  62. iter.EXPECT().All(gomock.Any()).Return(test.err)
  63. promise := new(mockPromise)
  64. pi := promisedIter{
  65. Iter: iter,
  66. promise: keepablePromise{
  67. promise: promise,
  68. log: func(error) {},
  69. },
  70. }
  71. assert.Equal(t, test.err, pi.All(nil))
  72. assert.Equal(t, test.accepted, promise.accepted)
  73. assert.Equal(t, test.reason, promise.reason)
  74. })
  75. }
  76. for _, test := range tests {
  77. t.Run(stringx.RandId(), func(t *testing.T) {
  78. ctrl := gomock.NewController(t)
  79. defer ctrl.Finish()
  80. iter := NewMockIter(ctrl)
  81. iter.EXPECT().Close().Return(test.err)
  82. promise := new(mockPromise)
  83. pi := promisedIter{
  84. Iter: iter,
  85. promise: keepablePromise{
  86. promise: promise,
  87. log: func(error) {},
  88. },
  89. }
  90. assert.Equal(t, test.err, pi.Close())
  91. assert.Equal(t, test.accepted, promise.accepted)
  92. assert.Equal(t, test.reason, promise.reason)
  93. })
  94. }
  95. }
  96. func TestPromisedIter_Err(t *testing.T) {
  97. errs := []error{
  98. nil,
  99. mgo.ErrNotFound,
  100. }
  101. for _, err := range errs {
  102. t.Run(stringx.RandId(), func(t *testing.T) {
  103. ctrl := gomock.NewController(t)
  104. defer ctrl.Finish()
  105. iter := NewMockIter(ctrl)
  106. iter.EXPECT().Err().Return(err)
  107. promise := new(mockPromise)
  108. pi := promisedIter{
  109. Iter: iter,
  110. promise: keepablePromise{
  111. promise: promise,
  112. log: func(error) {},
  113. },
  114. }
  115. assert.Equal(t, err, pi.Err())
  116. })
  117. }
  118. }
  119. func TestPromisedIter_For(t *testing.T) {
  120. tests := []struct {
  121. err error
  122. accepted bool
  123. reason string
  124. }{
  125. {
  126. err: nil,
  127. accepted: true,
  128. reason: "",
  129. },
  130. {
  131. err: mgo.ErrNotFound,
  132. accepted: true,
  133. reason: "",
  134. },
  135. {
  136. err: errors.New("any"),
  137. accepted: false,
  138. reason: "any",
  139. },
  140. }
  141. for _, test := range tests {
  142. t.Run(stringx.RandId(), func(t *testing.T) {
  143. ctrl := gomock.NewController(t)
  144. defer ctrl.Finish()
  145. iter := NewMockIter(ctrl)
  146. iter.EXPECT().For(gomock.Any(), gomock.Any()).Return(test.err)
  147. promise := new(mockPromise)
  148. pi := promisedIter{
  149. Iter: iter,
  150. promise: keepablePromise{
  151. promise: promise,
  152. log: func(error) {},
  153. },
  154. }
  155. assert.Equal(t, test.err, pi.For(nil, nil))
  156. assert.Equal(t, test.accepted, promise.accepted)
  157. assert.Equal(t, test.reason, promise.reason)
  158. })
  159. }
  160. }
  161. func TestRejectedIter_All(t *testing.T) {
  162. assert.Equal(t, breaker.ErrServiceUnavailable, new(rejectedIter).All(nil))
  163. }
  164. func TestRejectedIter_Close(t *testing.T) {
  165. assert.Equal(t, breaker.ErrServiceUnavailable, new(rejectedIter).Close())
  166. }
  167. func TestRejectedIter_Done(t *testing.T) {
  168. assert.False(t, new(rejectedIter).Done())
  169. }
  170. func TestRejectedIter_Err(t *testing.T) {
  171. assert.Equal(t, breaker.ErrServiceUnavailable, new(rejectedIter).Err())
  172. }
  173. func TestRejectedIter_For(t *testing.T) {
  174. assert.Equal(t, breaker.ErrServiceUnavailable, new(rejectedIter).For(nil, nil))
  175. }
  176. func TestRejectedIter_Next(t *testing.T) {
  177. assert.False(t, new(rejectedIter).Next(nil))
  178. }
  179. func TestRejectedIter_State(t *testing.T) {
  180. n, raw := new(rejectedIter).State()
  181. assert.Equal(t, int64(0), n)
  182. assert.Nil(t, raw)
  183. }
  184. func TestRejectedIter_Timeout(t *testing.T) {
  185. assert.False(t, new(rejectedIter).Timeout())
  186. }
  187. func TestIter_Done(t *testing.T) {
  188. ctrl := gomock.NewController(t)
  189. defer ctrl.Finish()
  190. iter := NewMockIter(ctrl)
  191. iter.EXPECT().Done().Return(true)
  192. ci := ClosableIter{
  193. Iter: iter,
  194. Cleanup: nil,
  195. }
  196. assert.True(t, ci.Done())
  197. }
  198. func TestIter_Next(t *testing.T) {
  199. ctrl := gomock.NewController(t)
  200. defer ctrl.Finish()
  201. iter := NewMockIter(ctrl)
  202. iter.EXPECT().Next(gomock.Any()).Return(true)
  203. ci := ClosableIter{
  204. Iter: iter,
  205. Cleanup: nil,
  206. }
  207. assert.True(t, ci.Next(nil))
  208. }
  209. func TestIter_State(t *testing.T) {
  210. ctrl := gomock.NewController(t)
  211. defer ctrl.Finish()
  212. iter := NewMockIter(ctrl)
  213. iter.EXPECT().State().Return(int64(1), nil)
  214. ci := ClosableIter{
  215. Iter: iter,
  216. Cleanup: nil,
  217. }
  218. n, raw := ci.State()
  219. assert.Equal(t, int64(1), n)
  220. assert.Nil(t, raw)
  221. }
  222. func TestIter_Timeout(t *testing.T) {
  223. ctrl := gomock.NewController(t)
  224. defer ctrl.Finish()
  225. iter := NewMockIter(ctrl)
  226. iter.EXPECT().Timeout().Return(true)
  227. ci := ClosableIter{
  228. Iter: iter,
  229. Cleanup: nil,
  230. }
  231. assert.True(t, ci.Timeout())
  232. }