diff --git a/integration/container/kill_test.go b/integration/container/kill_test.go
index 2c73026e18..60f4f179bd 100644
--- a/integration/container/kill_test.go
+++ b/integration/container/kill_test.go
@@ -5,7 +5,6 @@ import (
 	"testing"
 	"time"
 
-	containertypes "github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/client"
 	"github.com/docker/docker/integration/internal/container"
 	"github.com/docker/docker/internal/test/request"
@@ -96,12 +95,11 @@ func TestKillWithStopSignalAndRestartPolicies(t *testing.T) {
 		tc := tc
 		t.Run(tc.doc, func(t *testing.T) {
 			ctx := context.Background()
-			id := container.Run(t, ctx, client, func(c *container.TestContainerConfig) {
-				c.Config.StopSignal = tc.stopsignal
-				c.HostConfig.RestartPolicy = containertypes.RestartPolicy{
-					Name: "always",
-				}
-			})
+			id := container.Run(t, ctx, client,
+				container.WithRestartPolicy("always"),
+				func(c *container.TestContainerConfig) {
+					c.Config.StopSignal = tc.stopsignal
+				})
 			err := client.ContainerKill(ctx, id, "TERM")
 			assert.NilError(t, err)
 
diff --git a/integration/container/stop_test.go b/integration/container/stop_test.go
index 43d729b7b1..8a95f32f13 100644
--- a/integration/container/stop_test.go
+++ b/integration/container/stop_test.go
@@ -17,9 +17,11 @@ func TestStopContainerWithRestartPolicyAlways(t *testing.T) {
 
 	names := []string{"verifyRestart1-" + t.Name(), "verifyRestart2-" + t.Name()}
 	for _, name := range names {
-		container.Run(t, ctx, client, container.WithName(name), container.WithCmd("false"), func(c *container.TestContainerConfig) {
-			c.HostConfig.RestartPolicy.Name = "always"
-		})
+		container.Run(t, ctx, client,
+			container.WithName(name),
+			container.WithCmd("false"),
+			container.WithRestartPolicy("always"),
+		)
 	}
 
 	for _, name := range names {
diff --git a/integration/internal/container/ops.go b/integration/internal/container/ops.go
index d1243f7b1a..7bbf6cf0a2 100644
--- a/integration/internal/container/ops.go
+++ b/integration/internal/container/ops.go
@@ -136,3 +136,10 @@ func WithLogDriver(driver string) func(*TestContainerConfig) {
 func WithAutoRemove(c *TestContainerConfig) {
 	c.HostConfig.AutoRemove = true
 }
+
+// WithRestartPolicy sets container's restart policy
+func WithRestartPolicy(policy string) func(c *TestContainerConfig) {
+	return func(c *TestContainerConfig) {
+		c.HostConfig.RestartPolicy.Name = policy
+	}
+}