mock_registry.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. // Copyright 2015 CoreOS, 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 subnet
  15. import (
  16. "fmt"
  17. "time"
  18. "github.com/coreos/flannel/Godeps/_workspace/src/github.com/coreos/go-etcd/etcd"
  19. "github.com/coreos/flannel/Godeps/_workspace/src/golang.org/x/net/context"
  20. )
  21. type mockSubnetRegistry struct {
  22. config *etcd.Node
  23. subnets *etcd.Node
  24. events chan *etcd.Response
  25. index uint64
  26. ttl uint64
  27. }
  28. func newMockRegistry(ttlOverride uint64, config string, initialSubnets []*etcd.Node) *mockSubnetRegistry {
  29. index := uint64(0)
  30. for _, n := range initialSubnets {
  31. if n.ModifiedIndex > index {
  32. index = n.ModifiedIndex
  33. }
  34. }
  35. return &mockSubnetRegistry{
  36. config: &etcd.Node{
  37. Value: config,
  38. },
  39. subnets: &etcd.Node{
  40. Nodes: initialSubnets,
  41. },
  42. events: make(chan *etcd.Response, 1000),
  43. index: index + 1,
  44. ttl: ttlOverride,
  45. }
  46. }
  47. func (msr *mockSubnetRegistry) getConfig(ctx context.Context, network string) (*etcd.Response, error) {
  48. return &etcd.Response{
  49. EtcdIndex: msr.index,
  50. Node: msr.config,
  51. }, nil
  52. }
  53. func (msr *mockSubnetRegistry) setConfig(config string) {
  54. msr.config = &etcd.Node{
  55. Key: "config",
  56. Value: config,
  57. }
  58. }
  59. func (msr *mockSubnetRegistry) getSubnets(ctx context.Context, network string) (*etcd.Response, error) {
  60. return &etcd.Response{
  61. Node: msr.subnets,
  62. EtcdIndex: msr.index,
  63. }, nil
  64. }
  65. func (msr *mockSubnetRegistry) createSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error) {
  66. msr.index += 1
  67. if msr.ttl > 0 {
  68. ttl = msr.ttl
  69. }
  70. // add squared durations :)
  71. exp := time.Now().Add(time.Duration(ttl) * time.Second)
  72. node := &etcd.Node{
  73. Key: sn,
  74. Value: data,
  75. ModifiedIndex: msr.index,
  76. Expiration: &exp,
  77. }
  78. msr.subnets.Nodes = append(msr.subnets.Nodes, node)
  79. msr.events <- &etcd.Response{
  80. Action: "add",
  81. Node: node,
  82. }
  83. return &etcd.Response{
  84. Node: node,
  85. EtcdIndex: msr.index,
  86. }, nil
  87. }
  88. func (msr *mockSubnetRegistry) updateSubnet(ctx context.Context, network, sn, data string, ttl uint64) (*etcd.Response, error) {
  89. msr.index += 1
  90. // add squared durations :)
  91. exp := time.Now().Add(time.Duration(ttl) * time.Second)
  92. for _, n := range msr.subnets.Nodes {
  93. if n.Key == sn {
  94. n.Value = data
  95. n.ModifiedIndex = msr.index
  96. n.Expiration = &exp
  97. msr.events <- &etcd.Response{
  98. Action: "add",
  99. Node: n,
  100. }
  101. return &etcd.Response{
  102. Node: n,
  103. EtcdIndex: msr.index,
  104. }, nil
  105. }
  106. }
  107. return nil, fmt.Errorf("Subnet not found")
  108. }
  109. func (msr *mockSubnetRegistry) deleteSubnet(ctx context.Context, network, sn string) (*etcd.Response, error) {
  110. msr.index += 1
  111. for i, n := range msr.subnets.Nodes {
  112. if n.Key == sn {
  113. msr.subnets.Nodes[i] = msr.subnets.Nodes[len(msr.subnets.Nodes)-1]
  114. msr.subnets.Nodes = msr.subnets.Nodes[:len(msr.subnets.Nodes)-1]
  115. msr.events <- &etcd.Response{
  116. Action: "delete",
  117. Node: n,
  118. }
  119. return &etcd.Response{
  120. Node: n,
  121. EtcdIndex: msr.index,
  122. }, nil
  123. }
  124. }
  125. return nil, fmt.Errorf("Subnet not found")
  126. }
  127. func (msr *mockSubnetRegistry) watchSubnets(ctx context.Context, network string, since uint64) (*etcd.Response, error) {
  128. for {
  129. select {
  130. case <-ctx.Done():
  131. return nil, ctx.Err()
  132. case r := <-msr.events:
  133. if r.Node.ModifiedIndex < since {
  134. continue
  135. }
  136. return r, nil
  137. }
  138. }
  139. }
  140. func (msr *mockSubnetRegistry) hasSubnet(sn string) bool {
  141. for _, n := range msr.subnets.Nodes {
  142. if n.Key == sn {
  143. return true
  144. }
  145. }
  146. return false
  147. }
  148. func (msr *mockSubnetRegistry) expireSubnet(sn string) {
  149. for i, n := range msr.subnets.Nodes {
  150. if n.Key == sn {
  151. msr.index += 1
  152. msr.subnets.Nodes[i] = msr.subnets.Nodes[len(msr.subnets.Nodes)-1]
  153. msr.subnets.Nodes = msr.subnets.Nodes[:len(msr.subnets.Nodes)-2]
  154. n.ModifiedIndex = msr.index
  155. msr.events <- &etcd.Response{
  156. Action: "expire",
  157. Node: n,
  158. }
  159. return
  160. }
  161. }
  162. }