controller.go 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336
  1. // Copyright 2010 Google Inc.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. // Package gomock is a mock framework for Go.
  15. //
  16. // Standard usage:
  17. // (1) Define an interface that you wish to mock.
  18. // type MyInterface interface {
  19. // SomeMethod(x int64, y string)
  20. // }
  21. // (2) Use mockgen to generate a mock from the interface.
  22. // (3) Use the mock in a test:
  23. // func TestMyThing(t *testing.T) {
  24. // mockCtrl := gomock.NewController(t)
  25. // defer mockCtrl.Finish()
  26. //
  27. // mockObj := something.NewMockMyInterface(mockCtrl)
  28. // mockObj.EXPECT().SomeMethod(4, "blah")
  29. // // pass mockObj to a real object and play with it.
  30. // }
  31. //
  32. // By default, expected calls are not enforced to run in any particular order.
  33. // Call order dependency can be enforced by use of InOrder and/or Call.After.
  34. // Call.After can create more varied call order dependencies, but InOrder is
  35. // often more convenient.
  36. //
  37. // The following examples create equivalent call order dependencies.
  38. //
  39. // Example of using Call.After to chain expected call order:
  40. //
  41. // firstCall := mockObj.EXPECT().SomeMethod(1, "first")
  42. // secondCall := mockObj.EXPECT().SomeMethod(2, "second").After(firstCall)
  43. // mockObj.EXPECT().SomeMethod(3, "third").After(secondCall)
  44. //
  45. // Example of using InOrder to declare expected call order:
  46. //
  47. // gomock.InOrder(
  48. // mockObj.EXPECT().SomeMethod(1, "first"),
  49. // mockObj.EXPECT().SomeMethod(2, "second"),
  50. // mockObj.EXPECT().SomeMethod(3, "third"),
  51. // )
  52. package gomock
  53. import (
  54. "context"
  55. "fmt"
  56. "reflect"
  57. "runtime"
  58. "sync"
  59. )
  60. // A TestReporter is something that can be used to report test failures. It
  61. // is satisfied by the standard library's *testing.T.
  62. type TestReporter interface {
  63. Errorf(format string, args ...interface{})
  64. Fatalf(format string, args ...interface{})
  65. }
  66. // TestHelper is a TestReporter that has the Helper method. It is satisfied
  67. // by the standard library's *testing.T.
  68. type TestHelper interface {
  69. TestReporter
  70. Helper()
  71. }
  72. // cleanuper is used to check if TestHelper also has the `Cleanup` method. A
  73. // common pattern is to pass in a `*testing.T` to
  74. // `NewController(t TestReporter)`. In Go 1.14+, `*testing.T` has a cleanup
  75. // method. This can be utilized to call `Finish()` so the caller of this library
  76. // does not have to.
  77. type cleanuper interface {
  78. Cleanup(func())
  79. }
  80. // A Controller represents the top-level control of a mock ecosystem. It
  81. // defines the scope and lifetime of mock objects, as well as their
  82. // expectations. It is safe to call Controller's methods from multiple
  83. // goroutines. Each test should create a new Controller and invoke Finish via
  84. // defer.
  85. //
  86. // func TestFoo(t *testing.T) {
  87. // ctrl := gomock.NewController(t)
  88. // defer ctrl.Finish()
  89. // // ..
  90. // }
  91. //
  92. // func TestBar(t *testing.T) {
  93. // t.Run("Sub-Test-1", st) {
  94. // ctrl := gomock.NewController(st)
  95. // defer ctrl.Finish()
  96. // // ..
  97. // })
  98. // t.Run("Sub-Test-2", st) {
  99. // ctrl := gomock.NewController(st)
  100. // defer ctrl.Finish()
  101. // // ..
  102. // })
  103. // })
  104. type Controller struct {
  105. // T should only be called within a generated mock. It is not intended to
  106. // be used in user code and may be changed in future versions. T is the
  107. // TestReporter passed in when creating the Controller via NewController.
  108. // If the TestReporter does not implement a TestHelper it will be wrapped
  109. // with a nopTestHelper.
  110. T TestHelper
  111. mu sync.Mutex
  112. expectedCalls *callSet
  113. finished bool
  114. }
  115. // NewController returns a new Controller. It is the preferred way to create a
  116. // Controller.
  117. //
  118. // New in go1.14+, if you are passing a *testing.T into this function you no
  119. // longer need to call ctrl.Finish() in your test methods.
  120. func NewController(t TestReporter) *Controller {
  121. h, ok := t.(TestHelper)
  122. if !ok {
  123. h = &nopTestHelper{t}
  124. }
  125. ctrl := &Controller{
  126. T: h,
  127. expectedCalls: newCallSet(),
  128. }
  129. if c, ok := isCleanuper(ctrl.T); ok {
  130. c.Cleanup(func() {
  131. ctrl.T.Helper()
  132. ctrl.finish(true, nil)
  133. })
  134. }
  135. return ctrl
  136. }
  137. type cancelReporter struct {
  138. t TestHelper
  139. cancel func()
  140. }
  141. func (r *cancelReporter) Errorf(format string, args ...interface{}) {
  142. r.t.Errorf(format, args...)
  143. }
  144. func (r *cancelReporter) Fatalf(format string, args ...interface{}) {
  145. defer r.cancel()
  146. r.t.Fatalf(format, args...)
  147. }
  148. func (r *cancelReporter) Helper() {
  149. r.t.Helper()
  150. }
  151. // WithContext returns a new Controller and a Context, which is cancelled on any
  152. // fatal failure.
  153. func WithContext(ctx context.Context, t TestReporter) (*Controller, context.Context) {
  154. h, ok := t.(TestHelper)
  155. if !ok {
  156. h = &nopTestHelper{t: t}
  157. }
  158. ctx, cancel := context.WithCancel(ctx)
  159. return NewController(&cancelReporter{t: h, cancel: cancel}), ctx
  160. }
  161. type nopTestHelper struct {
  162. t TestReporter
  163. }
  164. func (h *nopTestHelper) Errorf(format string, args ...interface{}) {
  165. h.t.Errorf(format, args...)
  166. }
  167. func (h *nopTestHelper) Fatalf(format string, args ...interface{}) {
  168. h.t.Fatalf(format, args...)
  169. }
  170. func (h nopTestHelper) Helper() {}
  171. // RecordCall is called by a mock. It should not be called by user code.
  172. func (ctrl *Controller) RecordCall(receiver interface{}, method string, args ...interface{}) *Call {
  173. ctrl.T.Helper()
  174. recv := reflect.ValueOf(receiver)
  175. for i := 0; i < recv.Type().NumMethod(); i++ {
  176. if recv.Type().Method(i).Name == method {
  177. return ctrl.RecordCallWithMethodType(receiver, method, recv.Method(i).Type(), args...)
  178. }
  179. }
  180. ctrl.T.Fatalf("gomock: failed finding method %s on %T", method, receiver)
  181. panic("unreachable")
  182. }
  183. // RecordCallWithMethodType is called by a mock. It should not be called by user code.
  184. func (ctrl *Controller) RecordCallWithMethodType(receiver interface{}, method string, methodType reflect.Type, args ...interface{}) *Call {
  185. ctrl.T.Helper()
  186. call := newCall(ctrl.T, receiver, method, methodType, args...)
  187. ctrl.mu.Lock()
  188. defer ctrl.mu.Unlock()
  189. ctrl.expectedCalls.Add(call)
  190. return call
  191. }
  192. // Call is called by a mock. It should not be called by user code.
  193. func (ctrl *Controller) Call(receiver interface{}, method string, args ...interface{}) []interface{} {
  194. ctrl.T.Helper()
  195. // Nest this code so we can use defer to make sure the lock is released.
  196. actions := func() []func([]interface{}) []interface{} {
  197. ctrl.T.Helper()
  198. ctrl.mu.Lock()
  199. defer ctrl.mu.Unlock()
  200. expected, err := ctrl.expectedCalls.FindMatch(receiver, method, args)
  201. if err != nil {
  202. // callerInfo's skip should be updated if the number of calls between the user's test
  203. // and this line changes, i.e. this code is wrapped in another anonymous function.
  204. // 0 is us, 1 is controller.Call(), 2 is the generated mock, and 3 is the user's test.
  205. origin := callerInfo(3)
  206. ctrl.T.Fatalf("Unexpected call to %T.%v(%v) at %s because: %s", receiver, method, args, origin, err)
  207. }
  208. // Two things happen here:
  209. // * the matching call no longer needs to check prerequite calls,
  210. // * and the prerequite calls are no longer expected, so remove them.
  211. preReqCalls := expected.dropPrereqs()
  212. for _, preReqCall := range preReqCalls {
  213. ctrl.expectedCalls.Remove(preReqCall)
  214. }
  215. actions := expected.call()
  216. if expected.exhausted() {
  217. ctrl.expectedCalls.Remove(expected)
  218. }
  219. return actions
  220. }()
  221. var rets []interface{}
  222. for _, action := range actions {
  223. if r := action(args); r != nil {
  224. rets = r
  225. }
  226. }
  227. return rets
  228. }
  229. // Finish checks to see if all the methods that were expected to be called
  230. // were called. It should be invoked for each Controller. It is not idempotent
  231. // and therefore can only be invoked once.
  232. //
  233. // New in go1.14+, if you are passing a *testing.T into NewController function you no
  234. // longer need to call ctrl.Finish() in your test methods.
  235. func (ctrl *Controller) Finish() {
  236. // If we're currently panicking, probably because this is a deferred call.
  237. // This must be recovered in the deferred function.
  238. err := recover()
  239. ctrl.finish(false, err)
  240. }
  241. func (ctrl *Controller) finish(cleanup bool, panicErr interface{}) {
  242. ctrl.T.Helper()
  243. ctrl.mu.Lock()
  244. defer ctrl.mu.Unlock()
  245. if ctrl.finished {
  246. if _, ok := isCleanuper(ctrl.T); !ok {
  247. ctrl.T.Fatalf("Controller.Finish was called more than once. It has to be called exactly once.")
  248. }
  249. return
  250. }
  251. ctrl.finished = true
  252. // Short-circuit, pass through the panic.
  253. if panicErr != nil {
  254. panic(panicErr)
  255. }
  256. // Check that all remaining expected calls are satisfied.
  257. failures := ctrl.expectedCalls.Failures()
  258. for _, call := range failures {
  259. ctrl.T.Errorf("missing call(s) to %v", call)
  260. }
  261. if len(failures) != 0 {
  262. if !cleanup {
  263. ctrl.T.Fatalf("aborting test due to missing call(s)")
  264. return
  265. }
  266. ctrl.T.Errorf("aborting test due to missing call(s)")
  267. }
  268. }
  269. // callerInfo returns the file:line of the call site. skip is the number
  270. // of stack frames to skip when reporting. 0 is callerInfo's call site.
  271. func callerInfo(skip int) string {
  272. if _, file, line, ok := runtime.Caller(skip + 1); ok {
  273. return fmt.Sprintf("%s:%d", file, line)
  274. }
  275. return "unknown file"
  276. }
  277. // isCleanuper checks it if t's base TestReporter has a Cleanup method.
  278. func isCleanuper(t TestReporter) (cleanuper, bool) {
  279. tr := unwrapTestReporter(t)
  280. c, ok := tr.(cleanuper)
  281. return c, ok
  282. }
  283. // unwrapTestReporter unwraps TestReporter to the base implementation.
  284. func unwrapTestReporter(t TestReporter) TestReporter {
  285. tr := t
  286. switch nt := t.(type) {
  287. case *cancelReporter:
  288. tr = nt.t
  289. if h, check := tr.(*nopTestHelper); check {
  290. tr = h.t
  291. }
  292. case *nopTestHelper:
  293. tr = nt.t
  294. default:
  295. // not wrapped
  296. }
  297. return tr
  298. }