diff --git a/integration-cli/docker_cli_update_test.go b/integration-cli/docker_cli_update_test.go
deleted file mode 100644
index dce137ebf8..0000000000
--- a/integration-cli/docker_cli_update_test.go
+++ /dev/null
@@ -1,43 +0,0 @@
-package main
-
-import (
-	"strings"
-	"time"
-
-	"github.com/docker/docker/integration-cli/checker"
-	"github.com/docker/docker/integration-cli/cli"
-	"github.com/go-check/check"
-	"github.com/gotestyourself/gotestyourself/icmd"
-)
-
-func (s *DockerSuite) TestUpdateRestartPolicy(c *check.C) {
-	out := cli.DockerCmd(c, "run", "-d", "--restart=on-failure:3", "busybox", "sh", "-c", "sleep 1 && false").Combined()
-	timeout := 60 * time.Second
-	if testEnv.OSType == "windows" {
-		timeout = 180 * time.Second
-	}
-
-	id := strings.TrimSpace(string(out))
-
-	// update restart policy to on-failure:5
-	cli.DockerCmd(c, "update", "--restart=on-failure:5", id)
-
-	cli.WaitExited(c, id, timeout)
-
-	count := inspectField(c, id, "RestartCount")
-	c.Assert(count, checker.Equals, "5")
-
-	maximumRetryCount := inspectField(c, id, "HostConfig.RestartPolicy.MaximumRetryCount")
-	c.Assert(maximumRetryCount, checker.Equals, "5")
-}
-
-func (s *DockerSuite) TestUpdateRestartWithAutoRemoveFlag(c *check.C) {
-	out := runSleepingContainer(c, "--rm")
-	id := strings.TrimSpace(out)
-
-	// update restart policy for an AutoRemove container
-	cli.Docker(cli.Args("update", "--restart=always", id)).Assert(c, icmd.Expected{
-		ExitCode: 1,
-		Err:      "Restart policy cannot be updated because AutoRemove is enabled for the container",
-	})
-}
diff --git a/integration/container/update_test.go b/integration/container/update_test.go
new file mode 100644
index 0000000000..aee95804a1
--- /dev/null
+++ b/integration/container/update_test.go
@@ -0,0 +1,65 @@
+package container // import "github.com/docker/docker/integration/container"
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	containertypes "github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/integration/internal/container"
+	"github.com/docker/docker/integration/internal/request"
+	"github.com/docker/docker/internal/testutil"
+	"github.com/gotestyourself/gotestyourself/poll"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUpdateRestartPolicy(t *testing.T) {
+	defer setupTest(t)()
+	client := request.NewAPIClient(t)
+	ctx := context.Background()
+
+	cID := container.Run(t, ctx, client, container.WithCmd("sh", "-c", "sleep 1 && false"), func(c *container.TestContainerConfig) {
+		c.HostConfig.RestartPolicy = containertypes.RestartPolicy{
+			Name:              "on-failure",
+			MaximumRetryCount: 3,
+		}
+	})
+
+	_, err := client.ContainerUpdate(ctx, cID, containertypes.UpdateConfig{
+		RestartPolicy: containertypes.RestartPolicy{
+			Name:              "on-failure",
+			MaximumRetryCount: 5,
+		},
+	})
+	require.NoError(t, err)
+
+	timeout := 60 * time.Second
+	if testEnv.OSType == "windows" {
+		timeout = 180 * time.Second
+	}
+
+	poll.WaitOn(t, containerIsInState(ctx, client, cID, "exited"), poll.WithDelay(100*time.Millisecond), poll.WithTimeout(timeout))
+
+	inspect, err := client.ContainerInspect(ctx, cID)
+	require.NoError(t, err)
+	assert.Equal(t, inspect.RestartCount, 5)
+	assert.Equal(t, inspect.HostConfig.RestartPolicy.MaximumRetryCount, 5)
+}
+
+func TestUpdateRestartWithAutoRemove(t *testing.T) {
+	defer setupTest(t)()
+	client := request.NewAPIClient(t)
+	ctx := context.Background()
+
+	cID := container.Run(t, ctx, client, func(c *container.TestContainerConfig) {
+		c.HostConfig.AutoRemove = true
+	})
+
+	_, err := client.ContainerUpdate(ctx, cID, containertypes.UpdateConfig{
+		RestartPolicy: containertypes.RestartPolicy{
+			Name: "always",
+		},
+	})
+	testutil.ErrorContains(t, err, "Restart policy cannot be updated because AutoRemove is enabled for the container")
+}