protect.go 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. package environment
  2. import (
  3. "context"
  4. "testing"
  5. "github.com/docker/docker/api/types"
  6. "github.com/docker/docker/api/types/filters"
  7. dclient "github.com/docker/docker/client"
  8. "gotest.tools/v3/assert"
  9. )
  10. var frozenImages = []string{"busybox:latest", "busybox:glibc", "hello-world:frozen", "debian:bullseye-slim"}
  11. type protectedElements struct {
  12. containers map[string]struct{}
  13. images map[string]struct{}
  14. networks map[string]struct{}
  15. plugins map[string]struct{}
  16. volumes map[string]struct{}
  17. }
  18. func newProtectedElements() protectedElements {
  19. return protectedElements{
  20. containers: map[string]struct{}{},
  21. images: map[string]struct{}{},
  22. networks: map[string]struct{}{},
  23. plugins: map[string]struct{}{},
  24. volumes: map[string]struct{}{},
  25. }
  26. }
  27. // ProtectAll protects the existing environment (containers, images, networks,
  28. // volumes, and, on Linux, plugins) from being cleaned up at the end of test
  29. // runs
  30. func ProtectAll(t testing.TB, testEnv *Execution) {
  31. t.Helper()
  32. ProtectContainers(t, testEnv)
  33. ProtectImages(t, testEnv)
  34. ProtectNetworks(t, testEnv)
  35. ProtectVolumes(t, testEnv)
  36. if testEnv.OSType == "linux" {
  37. ProtectPlugins(t, testEnv)
  38. }
  39. }
  40. // ProtectContainer adds the specified container(s) to be protected in case of
  41. // clean
  42. func (e *Execution) ProtectContainer(t testing.TB, containers ...string) {
  43. t.Helper()
  44. for _, container := range containers {
  45. e.protectedElements.containers[container] = struct{}{}
  46. }
  47. }
  48. // ProtectContainers protects existing containers from being cleaned up at the
  49. // end of test runs
  50. func ProtectContainers(t testing.TB, testEnv *Execution) {
  51. t.Helper()
  52. containers := getExistingContainers(t, testEnv)
  53. testEnv.ProtectContainer(t, containers...)
  54. }
  55. func getExistingContainers(t testing.TB, testEnv *Execution) []string {
  56. t.Helper()
  57. client := testEnv.APIClient()
  58. containerList, err := client.ContainerList(context.Background(), types.ContainerListOptions{
  59. All: true,
  60. })
  61. assert.NilError(t, err, "failed to list containers")
  62. var containers []string
  63. for _, container := range containerList {
  64. containers = append(containers, container.ID)
  65. }
  66. return containers
  67. }
  68. // ProtectImage adds the specified image(s) to be protected in case of clean
  69. func (e *Execution) ProtectImage(t testing.TB, images ...string) {
  70. t.Helper()
  71. for _, image := range images {
  72. e.protectedElements.images[image] = struct{}{}
  73. }
  74. }
  75. // ProtectImages protects existing images and on linux frozen images from being
  76. // cleaned up at the end of test runs
  77. func ProtectImages(t testing.TB, testEnv *Execution) {
  78. t.Helper()
  79. images := getExistingImages(t, testEnv)
  80. if testEnv.OSType == "linux" {
  81. images = append(images, frozenImages...)
  82. }
  83. testEnv.ProtectImage(t, images...)
  84. }
  85. func getExistingImages(t testing.TB, testEnv *Execution) []string {
  86. t.Helper()
  87. client := testEnv.APIClient()
  88. filter := filters.NewArgs()
  89. filter.Add("dangling", "false")
  90. imageList, err := client.ImageList(context.Background(), types.ImageListOptions{
  91. All: true,
  92. Filters: filter,
  93. })
  94. assert.NilError(t, err, "failed to list images")
  95. var images []string
  96. for _, image := range imageList {
  97. images = append(images, tagsFromImageSummary(image)...)
  98. }
  99. return images
  100. }
  101. func tagsFromImageSummary(image types.ImageSummary) []string {
  102. var result []string
  103. for _, tag := range image.RepoTags {
  104. if tag != "<none>:<none>" {
  105. result = append(result, tag)
  106. }
  107. }
  108. for _, digest := range image.RepoDigests {
  109. if digest != "<none>@<none>" {
  110. result = append(result, digest)
  111. }
  112. }
  113. return result
  114. }
  115. // ProtectNetwork adds the specified network(s) to be protected in case of
  116. // clean
  117. func (e *Execution) ProtectNetwork(t testing.TB, networks ...string) {
  118. t.Helper()
  119. for _, network := range networks {
  120. e.protectedElements.networks[network] = struct{}{}
  121. }
  122. }
  123. // ProtectNetworks protects existing networks from being cleaned up at the end
  124. // of test runs
  125. func ProtectNetworks(t testing.TB, testEnv *Execution) {
  126. t.Helper()
  127. networks := getExistingNetworks(t, testEnv)
  128. testEnv.ProtectNetwork(t, networks...)
  129. }
  130. func getExistingNetworks(t testing.TB, testEnv *Execution) []string {
  131. t.Helper()
  132. client := testEnv.APIClient()
  133. networkList, err := client.NetworkList(context.Background(), types.NetworkListOptions{})
  134. assert.NilError(t, err, "failed to list networks")
  135. var networks []string
  136. for _, network := range networkList {
  137. networks = append(networks, network.ID)
  138. }
  139. return networks
  140. }
  141. // ProtectPlugin adds the specified plugin(s) to be protected in case of clean
  142. func (e *Execution) ProtectPlugin(t testing.TB, plugins ...string) {
  143. t.Helper()
  144. for _, plugin := range plugins {
  145. e.protectedElements.plugins[plugin] = struct{}{}
  146. }
  147. }
  148. // ProtectPlugins protects existing plugins from being cleaned up at the end of
  149. // test runs
  150. func ProtectPlugins(t testing.TB, testEnv *Execution) {
  151. t.Helper()
  152. plugins := getExistingPlugins(t, testEnv)
  153. testEnv.ProtectPlugin(t, plugins...)
  154. }
  155. func getExistingPlugins(t testing.TB, testEnv *Execution) []string {
  156. t.Helper()
  157. client := testEnv.APIClient()
  158. pluginList, err := client.PluginList(context.Background(), filters.Args{})
  159. // Docker EE does not allow cluster-wide plugin management.
  160. if dclient.IsErrNotImplemented(err) {
  161. return []string{}
  162. }
  163. assert.NilError(t, err, "failed to list plugins")
  164. var plugins []string
  165. for _, plugin := range pluginList {
  166. plugins = append(plugins, plugin.Name)
  167. }
  168. return plugins
  169. }
  170. // ProtectVolume adds the specified volume(s) to be protected in case of clean
  171. func (e *Execution) ProtectVolume(t testing.TB, volumes ...string) {
  172. t.Helper()
  173. for _, volume := range volumes {
  174. e.protectedElements.volumes[volume] = struct{}{}
  175. }
  176. }
  177. // ProtectVolumes protects existing volumes from being cleaned up at the end of
  178. // test runs
  179. func ProtectVolumes(t testing.TB, testEnv *Execution) {
  180. t.Helper()
  181. volumes := getExistingVolumes(t, testEnv)
  182. testEnv.ProtectVolume(t, volumes...)
  183. }
  184. func getExistingVolumes(t testing.TB, testEnv *Execution) []string {
  185. t.Helper()
  186. client := testEnv.APIClient()
  187. volumeList, err := client.VolumeList(context.Background(), filters.Args{})
  188. assert.NilError(t, err, "failed to list volumes")
  189. var volumes []string
  190. for _, volume := range volumeList.Volumes {
  191. volumes = append(volumes, volume.Name)
  192. }
  193. return volumes
  194. }