controller_stub.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. package exec
  2. import (
  3. "github.com/docker/swarmkit/api"
  4. "golang.org/x/net/context"
  5. "runtime"
  6. "strings"
  7. )
  8. // StubController implements the Controller interface,
  9. // but allows you to specify behaviors for each of its methods.
  10. type StubController struct {
  11. Controller
  12. UpdateFn func(ctx context.Context, t *api.Task) error
  13. PrepareFn func(ctx context.Context) error
  14. StartFn func(ctx context.Context) error
  15. WaitFn func(ctx context.Context) error
  16. ShutdownFn func(ctx context.Context) error
  17. TerminateFn func(ctx context.Context) error
  18. RemoveFn func(ctx context.Context) error
  19. CloseFn func() error
  20. calls map[string]int
  21. cstatus *api.ContainerStatus
  22. }
  23. // NewStubController returns an initialized StubController
  24. func NewStubController() *StubController {
  25. return &StubController{
  26. calls: make(map[string]int),
  27. }
  28. }
  29. // If function A calls updateCountsForSelf,
  30. // The callCount[A] value will be incremented
  31. func (sc *StubController) called() {
  32. pc, _, _, ok := runtime.Caller(1)
  33. if !ok {
  34. panic("Failed to find caller of function")
  35. }
  36. // longName looks like 'github.com/docker/swarmkit/agent/exec.(*StubController).Prepare:1'
  37. longName := runtime.FuncForPC(pc).Name()
  38. parts := strings.Split(longName, ".")
  39. tail := strings.Split(parts[len(parts)-1], ":")
  40. sc.calls[tail[0]]++
  41. }
  42. // Update is part of the Controller interface
  43. func (sc *StubController) Update(ctx context.Context, t *api.Task) error {
  44. sc.called()
  45. return sc.UpdateFn(ctx, t)
  46. }
  47. // Prepare is part of the Controller interface
  48. func (sc *StubController) Prepare(ctx context.Context) error { sc.called(); return sc.PrepareFn(ctx) }
  49. // Start is part of the Controller interface
  50. func (sc *StubController) Start(ctx context.Context) error { sc.called(); return sc.StartFn(ctx) }
  51. // Wait is part of the Controller interface
  52. func (sc *StubController) Wait(ctx context.Context) error { sc.called(); return sc.WaitFn(ctx) }
  53. // Shutdown is part of the Controller interface
  54. func (sc *StubController) Shutdown(ctx context.Context) error { sc.called(); return sc.ShutdownFn(ctx) }
  55. // Terminate is part of the Controller interface
  56. func (sc *StubController) Terminate(ctx context.Context) error {
  57. sc.called()
  58. return sc.TerminateFn(ctx)
  59. }
  60. // Remove is part of the Controller interface
  61. func (sc *StubController) Remove(ctx context.Context) error { sc.called(); return sc.RemoveFn(ctx) }
  62. // Close is part of the Controller interface
  63. func (sc *StubController) Close() error { sc.called(); return sc.CloseFn() }