@@ -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)

@@ -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 {

@@ -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

+	}

+}