ops.go 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. package container
  2. import (
  3. "fmt"
  4. "strings"
  5. containertypes "github.com/docker/docker/api/types/container"
  6. mounttypes "github.com/docker/docker/api/types/mount"
  7. networktypes "github.com/docker/docker/api/types/network"
  8. "github.com/docker/docker/api/types/strslice"
  9. "github.com/docker/go-connections/nat"
  10. specs "github.com/opencontainers/image-spec/specs-go/v1"
  11. )
  12. // WithName sets the name of the container
  13. func WithName(name string) func(*TestContainerConfig) {
  14. return func(c *TestContainerConfig) {
  15. c.Name = name
  16. }
  17. }
  18. // WithLinks sets the links of the container
  19. func WithLinks(links ...string) func(*TestContainerConfig) {
  20. return func(c *TestContainerConfig) {
  21. c.HostConfig.Links = links
  22. }
  23. }
  24. // WithImage sets the image of the container
  25. func WithImage(image string) func(*TestContainerConfig) {
  26. return func(c *TestContainerConfig) {
  27. c.Config.Image = image
  28. }
  29. }
  30. // WithCmd sets the comannds of the container
  31. func WithCmd(cmds ...string) func(*TestContainerConfig) {
  32. return func(c *TestContainerConfig) {
  33. c.Config.Cmd = strslice.StrSlice(cmds)
  34. }
  35. }
  36. // WithNetworkMode sets the network mode of the container
  37. func WithNetworkMode(mode string) func(*TestContainerConfig) {
  38. return func(c *TestContainerConfig) {
  39. c.HostConfig.NetworkMode = containertypes.NetworkMode(mode)
  40. }
  41. }
  42. // WithExposedPorts sets the exposed ports of the container
  43. func WithExposedPorts(ports ...string) func(*TestContainerConfig) {
  44. return func(c *TestContainerConfig) {
  45. c.Config.ExposedPorts = map[nat.Port]struct{}{}
  46. for _, port := range ports {
  47. c.Config.ExposedPorts[nat.Port(port)] = struct{}{}
  48. }
  49. }
  50. }
  51. // WithTty sets the TTY mode of the container
  52. func WithTty(tty bool) func(*TestContainerConfig) {
  53. return func(c *TestContainerConfig) {
  54. c.Config.Tty = tty
  55. }
  56. }
  57. // WithWorkingDir sets the working dir of the container
  58. func WithWorkingDir(dir string) func(*TestContainerConfig) {
  59. return func(c *TestContainerConfig) {
  60. c.Config.WorkingDir = dir
  61. }
  62. }
  63. // WithMount adds an mount
  64. func WithMount(m mounttypes.Mount) func(*TestContainerConfig) {
  65. return func(c *TestContainerConfig) {
  66. c.HostConfig.Mounts = append(c.HostConfig.Mounts, m)
  67. }
  68. }
  69. // WithVolume sets the volume of the container
  70. func WithVolume(target string) func(*TestContainerConfig) {
  71. return func(c *TestContainerConfig) {
  72. if c.Config.Volumes == nil {
  73. c.Config.Volumes = map[string]struct{}{}
  74. }
  75. c.Config.Volumes[target] = struct{}{}
  76. }
  77. }
  78. // WithBind sets the bind mount of the container
  79. func WithBind(src, target string) func(*TestContainerConfig) {
  80. return func(c *TestContainerConfig) {
  81. c.HostConfig.Binds = append(c.HostConfig.Binds, fmt.Sprintf("%s:%s", src, target))
  82. }
  83. }
  84. // WithTmpfs sets a target path in the container to a tmpfs
  85. func WithTmpfs(target string) func(config *TestContainerConfig) {
  86. return func(c *TestContainerConfig) {
  87. if c.HostConfig.Tmpfs == nil {
  88. c.HostConfig.Tmpfs = make(map[string]string)
  89. }
  90. spec := strings.SplitN(target, ":", 2)
  91. var opts string
  92. if len(spec) > 1 {
  93. opts = spec[1]
  94. }
  95. c.HostConfig.Tmpfs[spec[0]] = opts
  96. }
  97. }
  98. // WithIPv4 sets the specified ip for the specified network of the container
  99. func WithIPv4(network, ip string) func(*TestContainerConfig) {
  100. return func(c *TestContainerConfig) {
  101. if c.NetworkingConfig.EndpointsConfig == nil {
  102. c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
  103. }
  104. if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
  105. c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
  106. }
  107. if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
  108. c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
  109. }
  110. c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv4Address = ip
  111. }
  112. }
  113. // WithIPv6 sets the specified ip6 for the specified network of the container
  114. func WithIPv6(network, ip string) func(*TestContainerConfig) {
  115. return func(c *TestContainerConfig) {
  116. if c.NetworkingConfig.EndpointsConfig == nil {
  117. c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
  118. }
  119. if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
  120. c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
  121. }
  122. if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
  123. c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
  124. }
  125. c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv6Address = ip
  126. }
  127. }
  128. // WithLogDriver sets the log driver to use for the container
  129. func WithLogDriver(driver string) func(*TestContainerConfig) {
  130. return func(c *TestContainerConfig) {
  131. c.HostConfig.LogConfig.Type = driver
  132. }
  133. }
  134. // WithAutoRemove sets the container to be removed on exit
  135. func WithAutoRemove(c *TestContainerConfig) {
  136. c.HostConfig.AutoRemove = true
  137. }
  138. // WithPidsLimit sets the container's "pids-limit
  139. func WithPidsLimit(limit *int64) func(*TestContainerConfig) {
  140. return func(c *TestContainerConfig) {
  141. if c.HostConfig == nil {
  142. c.HostConfig = &containertypes.HostConfig{}
  143. }
  144. c.HostConfig.PidsLimit = limit
  145. }
  146. }
  147. // WithRestartPolicy sets container's restart policy
  148. func WithRestartPolicy(policy string) func(c *TestContainerConfig) {
  149. return func(c *TestContainerConfig) {
  150. c.HostConfig.RestartPolicy.Name = policy
  151. }
  152. }
  153. // WithUser sets the user
  154. func WithUser(user string) func(c *TestContainerConfig) {
  155. return func(c *TestContainerConfig) {
  156. c.Config.User = user
  157. }
  158. }
  159. // WithPrivileged sets privileged mode for the container
  160. func WithPrivileged(privileged bool) func(*TestContainerConfig) {
  161. return func(c *TestContainerConfig) {
  162. if c.HostConfig == nil {
  163. c.HostConfig = &containertypes.HostConfig{}
  164. }
  165. c.HostConfig.Privileged = privileged
  166. }
  167. }
  168. // WithCgroupnsMode sets the cgroup namespace mode for the container
  169. func WithCgroupnsMode(mode string) func(*TestContainerConfig) {
  170. return func(c *TestContainerConfig) {
  171. if c.HostConfig == nil {
  172. c.HostConfig = &containertypes.HostConfig{}
  173. }
  174. c.HostConfig.CgroupnsMode = containertypes.CgroupnsMode(mode)
  175. }
  176. }
  177. // WithExtraHost sets the user defined IP:Host mappings in the container's
  178. // /etc/hosts file
  179. func WithExtraHost(extraHost string) func(*TestContainerConfig) {
  180. return func(c *TestContainerConfig) {
  181. c.HostConfig.ExtraHosts = append(c.HostConfig.ExtraHosts, extraHost)
  182. }
  183. }
  184. // WithPlatform specifies the desired platform the image should have.
  185. func WithPlatform(p *specs.Platform) func(*TestContainerConfig) {
  186. return func(c *TestContainerConfig) {
  187. c.Platform = p
  188. }
  189. }