feat(kubernetes): advanced node drain options with agent failover [C9S-334] (#3346)

2026-08-07 11:14:49 +00:00 · 2026-08-05 12:09:44 +12:00
parent b2009d1937
commit 5d4dbb5a67
18 changed files with 1135 additions and 34 deletions
@@ -0,0 +1,175 @@
+package libkubectl
+
+import (
+	"context"
+	"time"
+
+	"github.com/rs/zerolog/log"
+	corev1 "k8s.io/api/core/v1"
+	k8serrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/kubectl/pkg/drain"
+)
+
+const (
+	// Namespace and Deployment name used by Portainer's default Kubernetes
+	// agent manifests. Custom installs using different values won't be
+	// detected, and the drain proceeds with no special agent handling.
+	portainerAgentNamespace      = "portainer"
+	portainerAgentDeploymentName = "portainer-agent"
+
+	// How often to poll for a replacement agent pod.
+	agentFailoverInterval = 2 * time.Second
+
+	// How long to wait for a replacement agent pod before giving up.
+	agentFailoverTimeout = 3 * time.Minute
+)
+
+// findAgentPodsOnNode returns the Portainer agent's pods that are scheduled
+// on nodeName. Returns nil if the agent Deployment isn't found or its pods
+// can't be listed — callers treat that as "nothing special to do" rather
+// than fail the drain over this best-effort heuristic.
+func findAgentPodsOnNode(ctx context.Context, clientset kubernetes.Interface, nodeName string) []corev1.Pod {
+	deployment, err := clientset.AppsV1().Deployments(portainerAgentNamespace).Get(ctx, portainerAgentDeploymentName, metav1.GetOptions{})
+	if k8serrors.IsNotFound(err) {
+		return nil
+	}
+	if err != nil {
+		log.Warn().
+			Str("context", "libkubectl").
+			Str("node_name", nodeName).
+			Err(err).
+			Msg("Unable to check for a Portainer agent deployment before draining; proceeding without special agent handling")
+		return nil
+	}
+
+	selector, err := metav1.LabelSelectorAsSelector(deployment.Spec.Selector)
+	if err != nil {
+		log.Warn().
+			Str("context", "libkubectl").
+			Str("node_name", nodeName).
+			Err(err).
+			Msg("Unable to parse the Portainer agent deployment's selector; proceeding without special agent handling")
+		return nil
+	}
+
+	agentPods, err := clientset.CoreV1().Pods(portainerAgentNamespace).List(ctx, metav1.ListOptions{
+		LabelSelector: selector.String(),
+	})
+	if err != nil {
+		log.Warn().
+			Str("context", "libkubectl").
+			Str("node_name", nodeName).
+			Err(err).
+			Msg("Unable to list Portainer agent pods before draining; proceeding without special agent handling")
+		return nil
+	}
+
+	return podsOnNode(agentPods.Items, nodeName)
+}
+
+// skipPodsFilter returns a drain.PodFilter that excludes the given pods
+// (matched by namespace+name) from a drain.Helper's normal pod-deletion
+// pass, so they can be handled separately.
+//
+// Note: drain.Helper runs its own base filters before AdditionalFilters like
+// this one, and stops at the first rejection. A pod that fails a base filter
+// (e.g. unreplicatedFilter, for a pod with no controller) would fail the
+// drain before ever reaching this filter. A real portainer-agent pod is
+// always owned by a ReplicaSet, so this doesn't happen in practice.
+func skipPodsFilter(pods []corev1.Pod) drain.PodFilter {
+	skip := make(map[string]bool, len(pods))
+	for _, pod := range pods {
+		skip[pod.Namespace+"/"+pod.Name] = true
+	}
+
+	return func(pod corev1.Pod) drain.PodDeleteStatus {
+		if skip[pod.Namespace+"/"+pod.Name] {
+			return drain.MakePodDeleteStatusSkip()
+		}
+
+		return drain.MakePodDeleteStatusOkay()
+	}
+}
+
+// evictAgentAndWaitForFailover evicts the given Portainer agent pods
+// (expected to already be excluded from the node's normal drain pass via
+// skipPodsFilter) and waits for a replacement to become Running elsewhere.
+//
+// Errors here are logged as warnings rather than returned: evicting the
+// agent pod can disrupt the very connection used to confirm that eviction,
+// which can look like a failure even when the eviction succeeded. Since this
+// runs after the rest of the node has already drained successfully, it's not
+// worth failing the overall drain over.
+func evictAgentAndWaitForFailover(ctx context.Context, drainer *drain.Helper, clientset kubernetes.Interface, nodeName string, agentPodsOnNode []corev1.Pod) {
+	if len(agentPodsOnNode) == 0 {
+		return
+	}
+
+	log.Info().
+		Str("context", "libkubectl").
+		Str("node_name", nodeName).
+		Msg("Evicting the Portainer agent pod on the drained node and waiting for a replacement on another node")
+
+	if err := drainer.DeleteOrEvictPods(agentPodsOnNode); err != nil {
+		log.Warn().
+			Str("context", "libkubectl").
+			Str("node_name", nodeName).
+			Err(err).
+			Msg("Evicting the portainer-agent pod reported an error; it may still have been evicted successfully despite the error")
+	}
+
+	err := wait.PollUntilContextTimeout(ctx, agentFailoverInterval, agentFailoverTimeout, true, func(ctx context.Context) (bool, error) {
+		selector, err := clientset.AppsV1().Deployments(portainerAgentNamespace).Get(ctx, portainerAgentDeploymentName, metav1.GetOptions{})
+		if err != nil {
+			return false, nil
+		}
+
+		labelSelector, err := metav1.LabelSelectorAsSelector(selector.Spec.Selector)
+		if err != nil {
+			return false, nil
+		}
+
+		pods, err := clientset.CoreV1().Pods(portainerAgentNamespace).List(ctx, metav1.ListOptions{
+			LabelSelector: labelSelector.String(),
+		})
+		if err != nil {
+			return false, nil
+		}
+
+		for _, pod := range pods.Items {
+			if pod.Spec.NodeName != nodeName && pod.Spec.NodeName != "" && pod.Status.Phase == corev1.PodRunning {
+				return true, nil
+			}
+		}
+
+		return false, nil
+	})
+	if err != nil {
+		log.Warn().
+			Str("context", "libkubectl").
+			Str("node_name", nodeName).
+			Err(err).
+			Msg("Timed out waiting for a new portainer-agent instance to start on another node; the node has otherwise been drained successfully")
+		return
+	}
+
+	log.Info().
+		Str("context", "libkubectl").
+		Str("node_name", nodeName).
+		Msg("Replacement Portainer agent pod is running on another node")
+}
+
+// podsOnNode returns the subset of pods scheduled on nodeName.
+func podsOnNode(pods []corev1.Pod, nodeName string) []corev1.Pod {
+	var onNode []corev1.Pod
+	for _, pod := range pods {
+		if pod.Spec.NodeName == nodeName {
+			onNode = append(onNode, pod)
+		}
+	}
+
+	return onNode
+}
@@ -0,0 +1,213 @@
+package libkubectl
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	kubetesting "k8s.io/client-go/testing"
+
+	kfake "k8s.io/client-go/kubernetes/fake"
+	"k8s.io/kubectl/pkg/cmd/util"
+	"k8s.io/kubectl/pkg/drain"
+)
+
+var errUnknownServerError = errors.New(`an error on the server ("unknown") has prevented the request from succeeding`)
+
+func agentDeployment() *appsv1.Deployment {
+	return &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      portainerAgentDeploymentName,
+			Namespace: portainerAgentNamespace,
+		},
+		Spec: appsv1.DeploymentSpec{
+			Selector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{"app": "portainer-agent"},
+			},
+		},
+	}
+}
+
+func agentPod(name, nodeName string, phase corev1.PodPhase) *corev1.Pod {
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: portainerAgentNamespace,
+			Labels:    map[string]string{"app": "portainer-agent"},
+		},
+		Spec: corev1.PodSpec{
+			NodeName: nodeName,
+		},
+		Status: corev1.PodStatus{
+			Phase: phase,
+		},
+	}
+}
+
+// testDrainer returns a *drain.Helper wired to clientset. It also registers
+// a bare "v1" API resource list on the fake clientset with no eviction
+// subresource, since drain.CheckEvictionSupport queries server discovery for
+// "v1" and errors out if the fake clientset has no resources registered at
+// all (rather than reporting "no eviction support") — this makes
+// DeleteOrEvictPods fall back to plain pod deletion, as it would against a
+// real, older/eviction-less cluster.
+func testDrainer(clientset *kfake.Clientset) *drain.Helper {
+	clientset.Resources = append(clientset.Resources, &metav1.APIResourceList{GroupVersion: "v1"})
+
+	buf := new(bytes.Buffer)
+	return &drain.Helper{
+		Ctx:            context.Background(),
+		Client:         clientset,
+		Timeout:        30 * time.Second,
+		Out:            buf,
+		ErrOut:         buf,
+		DryRunStrategy: util.DryRunNone,
+	}
+}
+
+func TestFindAgentPodsOnNode_NoAgentDeployment(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset()
+
+	pods := findAgentPodsOnNode(context.Background(), clientset, "node-1")
+	require.Empty(t, pods)
+}
+
+func TestFindAgentPodsOnNode_AgentNotOnDrainedNode(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		agentDeployment(),
+		agentPod("portainer-agent-abc", "node-2", corev1.PodRunning),
+	)
+
+	pods := findAgentPodsOnNode(context.Background(), clientset, "node-1")
+	require.Empty(t, pods)
+}
+
+func TestFindAgentPodsOnNode_AgentOnDrainedNode(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		agentDeployment(),
+		agentPod("portainer-agent-abc", "node-1", corev1.PodRunning),
+	)
+
+	pods := findAgentPodsOnNode(context.Background(), clientset, "node-1")
+	require.Len(t, pods, 1)
+	require.Equal(t, "portainer-agent-abc", pods[0].Name)
+}
+
+func TestSkipPodsFilter_SkipsMatchingPodAndKeepsOthers(t *testing.T) {
+	t.Parallel()
+
+	skip := agentPod("portainer-agent-abc", "node-1", corev1.PodRunning)
+	filter := skipPodsFilter([]corev1.Pod{*skip})
+
+	status := filter(*skip)
+	require.False(t, status.Delete)
+	require.Equal(t, drain.PodDeleteStatusTypeSkip, status.Reason)
+
+	other := corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "other-pod", Namespace: "default"}}
+	otherStatus := filter(other)
+	require.True(t, otherStatus.Delete)
+}
+
+// TestEvictAgentAndWaitForFailover_EvictsThenWaitsForReplacement asserts
+// that evicting the agent pod is what triggers the wait to succeed, rather
+// than the function passively waiting for external state it never caused.
+func TestEvictAgentAndWaitForFailover_EvictsThenWaitsForReplacement(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		agentDeployment(),
+		agentPod("portainer-agent-old", "node-1", corev1.PodRunning),
+	)
+
+	clientset.PrependReactor("delete", "pods", func(action kubetesting.Action) (bool, runtime.Object, error) {
+		deleteAction, ok := action.(kubetesting.DeleteAction)
+		if ok && deleteAction.GetName() == "portainer-agent-old" {
+			go func() {
+				_, _ = clientset.CoreV1().Pods(portainerAgentNamespace).Create(
+					context.Background(),
+					agentPod("portainer-agent-new", "node-2", corev1.PodRunning),
+					metav1.CreateOptions{},
+				)
+			}()
+		}
+		return false, nil, nil
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	agentPodsOnNode := findAgentPodsOnNode(ctx, clientset, "node-1")
+	require.Len(t, agentPodsOnNode, 1)
+
+	evictAgentAndWaitForFailover(ctx, testDrainer(clientset), clientset, "node-1", agentPodsOnNode)
+
+	_, err := clientset.CoreV1().Pods(portainerAgentNamespace).Get(context.Background(), "portainer-agent-old", metav1.GetOptions{})
+	require.Error(t, err, "the original agent pod on the drained node should have been deleted")
+}
+
+// TestEvictAgentAndWaitForFailover_TimeoutIsNonFatal asserts that when no
+// replacement ever appears, the function does not panic or block forever —
+// it simply gives up after the bound and returns, since by the time this
+// runs the rest of the node has already drained successfully and the
+// agent's own fate is best-effort.
+func TestEvictAgentAndWaitForFailover_TimeoutIsNonFatal(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		agentDeployment(),
+		agentPod("portainer-agent-old", "node-1", corev1.PodRunning),
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*agentFailoverInterval)
+	defer cancel()
+
+	agentPodsOnNode := findAgentPodsOnNode(ctx, clientset, "node-1")
+	require.Len(t, agentPodsOnNode, 1)
+
+	// Must return (not panic, not hang) even though no replacement appears.
+	evictAgentAndWaitForFailover(ctx, testDrainer(clientset), clientset, "node-1", agentPodsOnNode)
+}
+
+// TestEvictAgentAndWaitForFailover_EvictionErrorIsNonFatal reproduces the
+// live-cluster regression: evicting the agent pod can itself return an error
+// (e.g. the "unknown" transient error seen when the agent pod being deleted
+// is also the connection's own proxy) even though the pod is genuinely
+// evicted. The function must not panic or otherwise fail hard — the caller
+// (DrainNode) no longer propagates any error from this step.
+func TestEvictAgentAndWaitForFailover_EvictionErrorIsNonFatal(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		agentDeployment(),
+		agentPod("portainer-agent-old", "node-1", corev1.PodRunning),
+	)
+
+	clientset.PrependReactor("delete", "pods", func(action kubetesting.Action) (bool, runtime.Object, error) {
+		deleteAction, ok := action.(kubetesting.DeleteAction)
+		if ok && deleteAction.GetName() == "portainer-agent-old" {
+			return true, nil, errUnknownServerError
+		}
+		return false, nil, nil
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*agentFailoverInterval)
+	defer cancel()
+
+	agentPodsOnNode := findAgentPodsOnNode(ctx, clientset, "node-1")
+	require.Len(t, agentPodsOnNode, 1)
+
+	evictAgentAndWaitForFailover(ctx, testDrainer(clientset), clientset, "node-1", agentPodsOnNode)
+}
@@ -8,19 +8,52 @@ import (

 	"github.com/rs/zerolog/log"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/kubectl/pkg/cmd/util"
 	"k8s.io/kubectl/pkg/drain"
 )

+// DrainOptions controls the behaviour of a node drain operation, mirroring the
+// flags exposed by "kubectl drain".
+type DrainOptions struct {
+	// Force allows deletion of standalone pods not managed by a controller.
+	Force bool
+	// Timeout is the overall time to wait for the drain to complete.
+	Timeout time.Duration
+	// GracePeriodSeconds overrides each pod's termination grace period. -1 uses
+	// the pod's own grace period.
+	GracePeriodSeconds int
+	// IgnoreAllDaemonSets skips DaemonSet-managed pods, which would otherwise
+	// block the drain since they are recreated by their controller.
+	IgnoreAllDaemonSets bool
+	// DeleteEmptyDirData allows eviction of pods using emptyDir volumes, whose
+	// data is lost once the pod is deleted.
+	DeleteEmptyDirData bool
+	// DisableEviction forces the use of direct pod deletion instead of the
+	// eviction API, ignoring any configured PodDisruptionBudgets.
+	DisableEviction bool
+}
+
+// DefaultDrainOptions returns the defaults applied when a caller omits the
+// drain request payload.
+func DefaultDrainOptions() DrainOptions {
+	return DrainOptions{
+		Force:               false,
+		GracePeriodSeconds:  -1,
+		IgnoreAllDaemonSets: true,
+		Timeout:             60 * time.Second,
+		DeleteEmptyDirData:  true,
+		DisableEviction:     false,
+	}
+}
+
 // DrainNode drains a node from the cluster
-func (c *Client) DrainNode(nodeName string) (string, error) {
+func (c *Client) DrainNode(nodeName string, opts DrainOptions) (string, error) {
 	log.Debug().
 		Str("context", "libkubectl").
 		Str("node_name", nodeName).
 		Msg("Starting node drain operation")

-	buf := new(bytes.Buffer)
-
 	// Get clientset from factory
 	clientset, err := c.factory.KubernetesClientSet()
 	if err != nil {
@@ -37,15 +70,25 @@ func (c *Client) DrainNode(nodeName string) (string, error) {
 		Str("node_name", nodeName).
 		Msg("Successfully obtained kubernetes clientset")

+	return drainNodeWithClient(context.Background(), clientset, nodeName, opts)
+}
+
+// drainNodeWithClient runs the cordon-drain-agent-failover sequence against
+// an already-resolved clientset. Split out from DrainNode so it can be
+// unit-tested with a fake clientset, since KubernetesClientSet() returns a
+// concrete type that can't be faked directly.
+func drainNodeWithClient(ctx context.Context, clientset kubernetes.Interface, nodeName string, opts DrainOptions) (string, error) {
+	buf := new(bytes.Buffer)
+
 	drainer := &drain.Helper{
-		Ctx:                 context.Background(),
+		Ctx:                 ctx,
 		Client:              clientset,
-		Force:               false,            // Don't force delete standalone pods
-		GracePeriodSeconds:  -1,               // Use pod's own grace period
-		IgnoreAllDaemonSets: true,             // Skip DaemonSet pods
-		Timeout:             60 * time.Second, // Overall timeout
-		DeleteEmptyDirData:  true,             // Delete pods with emptyDir
-		DisableEviction:     false,            // Use eviction API when possible
+		Force:               opts.Force,
+		GracePeriodSeconds:  opts.GracePeriodSeconds,
+		IgnoreAllDaemonSets: opts.IgnoreAllDaemonSets,
+		Timeout:             opts.Timeout,
+		DeleteEmptyDirData:  opts.DeleteEmptyDirData,
+		DisableEviction:     opts.DisableEviction,
 		Out:                 buf,
 		ErrOut:              buf,
 		DryRunStrategy:      util.DryRunNone,
@@ -60,7 +103,7 @@ func (c *Client) DrainNode(nodeName string) (string, error) {
 		Msg("Configured drain helper")

 	// Get the node first
-	node, err := clientset.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
+	node, err := clientset.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
 	if err != nil {
 		log.Error().
 			Str("context", "libkubectl").
@@ -90,6 +133,13 @@ func (c *Client) DrainNode(nodeName string) (string, error) {
 		Str("node_name", nodeName).
 		Msg("Successfully cordoned node, proceeding to drain")

+	// Skip the Portainer agent pod in the normal drain pass and evict it
+	// separately, last (see evictAgentAndWaitForFailover).
+	agentPodsOnNode := findAgentPodsOnNode(ctx, clientset, nodeName)
+	if len(agentPodsOnNode) > 0 {
+		drainer.AdditionalFilters = append(drainer.AdditionalFilters, skipPodsFilter(agentPodsOnNode))
+	}
+
 	// Then drain it
 	if err := drain.RunNodeDrain(drainer, nodeName); err != nil {
 		log.Error().
@@ -100,6 +150,8 @@ func (c *Client) DrainNode(nodeName string) (string, error) {
 		return "", fmt.Errorf("failed to drain node %s: %w", nodeName, err)
 	}

+	evictAgentAndWaitForFailover(ctx, drainer, clientset, nodeName, agentPodsOnNode)
+
 	log.Debug().
 		Str("context", "libkubectl").
 		Str("node_name", nodeName).
@@ -0,0 +1,127 @@
+package libkubectl
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	kfake "k8s.io/client-go/kubernetes/fake"
+	kubetesting "k8s.io/client-go/testing"
+)
+
+func testNode(name string) *corev1.Node {
+	return &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: name},
+	}
+}
+
+// replicaSetOwnedPod returns a pod with an owner reference, so drain's
+// unreplicatedFilter allows it to be drained without requiring Force.
+func replicaSetOwnedPod(name, namespace, nodeName string) *corev1.Pod {
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: "apps/v1",
+					Kind:       "ReplicaSet",
+					Name:       "some-replicaset",
+					Controller: new(true),
+				},
+			},
+		},
+		Spec: corev1.PodSpec{NodeName: nodeName},
+		Status: corev1.PodStatus{
+			Phase: corev1.PodRunning,
+		},
+	}
+}
+
+// replicaSetOwnedAgentPod is agentPod (from agent_failover_test.go) with an
+// owner reference added, so it survives drain's unreplicatedFilter — needed
+// here since this pod goes through the real drain.RunNodeDrain filter chain.
+func replicaSetOwnedAgentPod(name, nodeName string, phase corev1.PodPhase) *corev1.Pod {
+	pod := agentPod(name, nodeName, phase)
+	pod.OwnerReferences = []metav1.OwnerReference{
+		{
+			APIVersion: "apps/v1",
+			Kind:       "ReplicaSet",
+			Name:       "portainer-agent-replicaset",
+			Controller: new(true),
+		},
+	}
+	return pod
+}
+
+func TestDrainNodeWithClient_NoAgent_DrainsRegularPods(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		testNode("node-1"),
+		replicaSetOwnedPod("app-pod", "default", "node-1"),
+	)
+	clientset.Resources = append(clientset.Resources, &metav1.APIResourceList{GroupVersion: "v1"})
+
+	opts := DefaultDrainOptions()
+	opts.Timeout = 10 * time.Second
+
+	_, err := drainNodeWithClient(context.Background(), clientset, "node-1", opts)
+	require.NoError(t, err)
+
+	_, err = clientset.CoreV1().Pods("default").Get(context.Background(), "app-pod", metav1.GetOptions{})
+	require.Error(t, err, "the regular pod on the drained node should have been deleted")
+
+	node, err := clientset.CoreV1().Nodes().Get(context.Background(), "node-1", metav1.GetOptions{})
+	require.NoError(t, err)
+	require.True(t, node.Spec.Unschedulable, "the node should have been cordoned")
+}
+
+// Proves drainNodeWithClient's wiring: the agent pod is skipped during the
+// normal drain pass, other pods still drain, and the agent is evicted last.
+func TestDrainNodeWithClient_AgentOnNode_SkipsAgentInNormalPassAndEvictsLast(t *testing.T) {
+	t.Parallel()
+
+	clientset := kfake.NewSimpleClientset(
+		testNode("node-1"),
+		agentDeployment(),
+		replicaSetOwnedAgentPod("portainer-agent-old", "node-1", corev1.PodRunning),
+		replicaSetOwnedPod("app-pod", "default", "node-1"),
+	)
+	clientset.Resources = append(clientset.Resources, &metav1.APIResourceList{GroupVersion: "v1"})
+
+	opts := DefaultDrainOptions()
+	opts.Timeout = 10 * time.Second
+
+	// Simulate the ReplicaSet controller creating a replacement once the old
+	// agent pod is deleted, so the poll succeeds quickly.
+	clientset.PrependReactor("delete", "pods", func(action kubetesting.Action) (bool, runtime.Object, error) {
+		deleteAction, ok := action.(kubetesting.DeleteAction)
+		if ok && deleteAction.GetName() == "portainer-agent-old" {
+			go func() {
+				_, _ = clientset.CoreV1().Pods(portainerAgentNamespace).Create(
+					context.Background(),
+					agentPod("portainer-agent-new", "node-2", corev1.PodRunning),
+					metav1.CreateOptions{},
+				)
+			}()
+		}
+		return false, nil, nil
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	_, err := drainNodeWithClient(ctx, clientset, "node-1", opts)
+	require.NoError(t, err, "drain must succeed and observe the agent's replacement coming up on another node")
+
+	_, err = clientset.CoreV1().Pods("default").Get(context.Background(), "app-pod", metav1.GetOptions{})
+	require.Error(t, err, "the regular pod must still be drained even though the node also runs the agent")
+
+	_, err = clientset.CoreV1().Pods(portainerAgentNamespace).Get(context.Background(), "portainer-agent-old", metav1.GetOptions{})
+	require.Error(t, err, "the agent pod must eventually be evicted too, just after everything else")
+}