From dd51b97d9d2e72b8d6d556a3c7467ff647f0a28b Mon Sep 17 00:00:00 2001
From: Fabricio Voznika <fvoznika@google.com>
Date: Mon, 15 Apr 2024 15:35:36 -0700
Subject: [PATCH] Add compression variant for checkpoint tests

PiperOrigin-RevId: 625104713
---
 pkg/state/statefile/statefile.go  |  2 ++
 runsc/cmd/checkpoint.go           |  2 +-
 runsc/container/container_test.go | 22 +++++++++++++++-------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/pkg/state/statefile/statefile.go b/pkg/state/statefile/statefile.go
index 66d2990bd..31b91d4a7 100644
--- a/pkg/state/statefile/statefile.go
+++ b/pkg/state/statefile/statefile.go
@@ -99,6 +99,8 @@ const (
 	CompressionLevelFlateBestSpeed = CompressionLevel("flate-best-speed")
 	// CompressionLevelNone represents the absence of any compression on an image.
 	CompressionLevelNone = CompressionLevel("none")
+	// CompressionLevelDefault represents the default compression level.
+	CompressionLevelDefault = CompressionLevelFlateBestSpeed
 )
 
 // Options is statefile options.
diff --git a/runsc/cmd/checkpoint.go b/runsc/cmd/checkpoint.go
index c25e4bdca..48775ebbb 100644
--- a/runsc/cmd/checkpoint.go
+++ b/runsc/cmd/checkpoint.go
@@ -54,7 +54,7 @@ func (*Checkpoint) Usage() string {
 func (c *Checkpoint) SetFlags(f *flag.FlagSet) {
 	f.StringVar(&c.imagePath, "image-path", "", "directory path to saved container image")
 	f.BoolVar(&c.leaveRunning, "leave-running", false, "restart the container after checkpointing")
-	f.Var(newCheckpointCompressionValue(statefile.CompressionLevelFlateBestSpeed, &c.compression), "compression", "compress checkpoint image on disk. Values: none|flate-best-speed.")
+	f.Var(newCheckpointCompressionValue(statefile.CompressionLevelDefault, &c.compression), "compression", "compress checkpoint image on disk. Values: none|flate-best-speed.")
 
 	// Unimplemented flags necessary for compatibility with docker.
 	var wp string
diff --git a/runsc/container/container_test.go b/runsc/container/container_test.go
index 61bf71a44..d896d9d89 100644
--- a/runsc/container/container_test.go
+++ b/runsc/container/container_test.go
@@ -1023,7 +1023,7 @@ func TestKillPid(t *testing.T) {
 // recorded. Then, it is restored in two new containers and the first number
 // printed from these containers is checked. Both should be the next consecutive
 // number after the last number from the checkpointed container.
-func testCheckpointRestore(t *testing.T, conf *config.Config, newSpecWithScript func(string) *specs.Spec) {
+func testCheckpointRestore(t *testing.T, conf *config.Config, compression statefile.CompressionLevel, newSpecWithScript func(string) *specs.Spec) {
 	dir, err := ioutil.TempDir(testutil.TmpDir(), "checkpoint-test")
 	if err != nil {
 		t.Fatalf("ioutil.TempDir failed: %v", err)
@@ -1069,7 +1069,7 @@ func testCheckpointRestore(t *testing.T, conf *config.Config, newSpecWithScript
 	}
 
 	// Checkpoint running container; save state into new file.
-	if err := cont.Checkpoint(dir, statefile.Options{Compression: statefile.CompressionLevelFlateBestSpeed}); err != nil {
+	if err := cont.Checkpoint(dir, statefile.Options{Compression: compression}); err != nil {
 		t.Fatalf("error checkpointing container to empty file: %v", err)
 	}
 
@@ -1171,9 +1171,17 @@ func TestCheckpointRestore(t *testing.T) {
 	// Skip overlay because test requires writing to host file.
 	for name, conf := range configs(t, true /* noOverlay */) {
 		t.Run(name, func(t *testing.T) {
-			testCheckpointRestore(t, conf, func(script string) *specs.Spec {
-				return testutil.NewSpecWithArgs("bash", "-c", script)
-			})
+			compressionLevels := []statefile.CompressionLevel{
+				statefile.CompressionLevelNone,
+				statefile.CompressionLevelFlateBestSpeed,
+			}
+			for _, compression := range compressionLevels {
+				t.Run(string(compression), func(t *testing.T) {
+					testCheckpointRestore(t, conf, compression, func(script string) *specs.Spec {
+						return testutil.NewSpecWithArgs("bash", "-c", script)
+					})
+				})
+			}
 		})
 	}
 }
@@ -1337,7 +1345,7 @@ func TestUnixDomainSockets(t *testing.T) {
 			}
 
 			// Checkpoint running container; save state into new file.
-			if err := cont.Checkpoint(dir, statefile.Options{Compression: statefile.CompressionLevelFlateBestSpeed}); err != nil {
+			if err := cont.Checkpoint(dir, statefile.Options{Compression: statefile.CompressionLevelDefault}); err != nil {
 				t.Fatalf("error checkpointing container to empty file: %v", err)
 			}
 
@@ -3401,7 +3409,7 @@ func TestCheckpointRestoreEROFS(t *testing.T) {
 	// Skip overlay because test requires writing to host file.
 	for name, conf := range configs(t, true /* noOverlay */) {
 		t.Run(name, func(t *testing.T) {
-			testCheckpointRestore(t, conf, func(script string) *specs.Spec {
+			testCheckpointRestore(t, conf, statefile.CompressionLevelDefault, func(script string) *specs.Spec {
 				spec := testutil.NewSpecWithArgs("/busybox", "sh", "-c", script)
 				spec.Root = &specs.Root{
 					Path:     rootfsDir,