fix(e2e): resolve agent pods by selector in clique-identity assertion (#437)

assertNodeCliqueIdentities was hardwired to the node-wide scenario's
app=gpu-agent selector, so when the IMEX scenario reused it the per-node pod
lookup polled for gpu-agent pods that never exist in the imex cluster and timed
out ("no running gpu-agent pod on node ..."). Parameterize the assertion with
namespace + selector (via the shared agentPodOnNode helper) and update both call
sites; the now-redundant nriAgentPodOnNode is removed. Surfaced by e2e-nri once
the imex scenario joined the nri label group.

Signed-off-by: Giulio Calzolari <gcalzolari@nvidia.com>

Author: giuliocalzo

tests/e2e/go/scenario_imex_test.go

@@ -84,7 +84,7 @@ var _ = Describe("nvml-mock NRI IMEX channel injection", Label("imex", "nri"), O
 				if !p.HasFabric() {
 					Skip("profile " + name + " declares no device_defaults.fabric block")
 				}
-				assertNodeCliqueIdentities(ctx, h, workers)
+				assertNodeCliqueIdentities(ctx, h, imexWorkloadNS, imexAgentSelector, workers)
 			})
 		})
 	}
@@ -146,8 +146,8 @@ func assertWorkloadSeesImexChannels(ctx context.Context, h *harness.Harness, wor
 	}
 }
 
-// agentPodOnNode generalizes nriAgentPodOnNode with namespace + selector
-// parameters so multiple scenarios can locate their per-node agent pod.
+// agentPodOnNode locates a running agent pod on the given node by namespace and
+// label selector, so multiple scenarios (gpu-agent, imex-agent) can share it.
 func agentPodOnNode(ctx context.Context, h *harness.Harness, ns, selector, node string) kube.PodRef {
 	GinkgoHelper()
 	var name string

tests/e2e/go/scenario_nri_test.go

@@ -100,7 +100,7 @@ var _ = Describe("nvml-mock node-wide NRI injection", Label("nri"), Ordered, fun
 				if !computeDomain {
 					Skip("profile " + name + " declares no device_defaults.fabric block; ComputeDomain identity is unsupported")
 				}
-				assertNodeCliqueIdentities(ctx, h, workers)
+				assertNodeCliqueIdentities(ctx, h, nriWorkloadNS, nriAgentSelector, workers)
 			})
 		})
 	}
@@ -171,15 +171,17 @@ func assertAgentSeesGPUs(ctx context.Context, h *harness.Harness, expectedGPUs i
 }
 
 // assertNodeCliqueIdentities runs the staged `check-fabric` consumer inside the
-// gpu-agent pod on every worker and asserts each node reports the clique /
-// cluster UUID the topology overlay assigned to it — with no nvidia.com/gpu
-// request and no MOCK_* env in the pod spec (identity comes only from NRI).
-func assertNodeCliqueIdentities(ctx context.Context, h *harness.Harness, workers []cluster.Node) {
+// agent pod (identified by ns/selector) on every worker and asserts each node
+// reports the clique / cluster UUID the topology overlay assigned to it — with
+// no nvidia.com/gpu request and no MOCK_* env in the pod spec (identity comes
+// only from NRI). ns/selector are parameterized so both the node-wide (gpu-agent)
+// and IMEX (imex-agent) scenarios can share this assertion.
+func assertNodeCliqueIdentities(ctx context.Context, h *harness.Harness, ns, selector string, workers []cluster.Node) {
 	GinkgoHelper()
 	cliqueByNode := nriCliqueByNode(workers)
 	for _, w := range workers {
 		expected := cliqueByNode[w.Name]
-		pod := nriAgentPodOnNode(ctx, h, w.Name)
+		pod := agentPodOnNode(ctx, h, ns, selector, w.Name)
 		res, err := h.Kube.ExecSh(ctx, pod, "check-fabric 2>&1")
 		Expect(err).NotTo(HaveOccurred(), "check-fabric on %s: %s", w.Name, res.Combined())
 		out := res.Combined()
@@ -205,30 +207,6 @@ func firstNRIAgentPod(ctx context.Context, h *harness.Harness) kube.PodRef {
 	return kube.PodRef{Namespace: nriWorkloadNS, Pod: name}
 }
 
-func nriAgentPodOnNode(ctx context.Context, h *harness.Harness, node string) kube.PodRef {
-	GinkgoHelper()
-	var name string
-	Eventually(func() (string, error) {
-		pods, err := h.Kube.RunningPodNames(ctx, nriWorkloadNS, nriAgentSelector)
-		if err != nil {
-			return "", err
-		}
-		for _, pod := range pods {
-			podNode, err := h.Kube.PodNode(ctx, nriWorkloadNS, pod)
-			if err != nil {
-				return "", err
-			}
-			if podNode == node {
-				name = pod
-				return name, nil
-			}
-		}
-		return "", nil
-	}).WithContext(ctx).WithTimeout(config.ReadyTimeout()).WithPolling(config.PollInterval()).
-		ShouldNot(BeEmpty(), "no running gpu-agent pod on node %s", node)
-	return kube.PodRef{Namespace: nriWorkloadNS, Pod: name}
-}
-
 func writeNRITopologyValues(workers []cluster.Node) string {
 	GinkgoHelper()
 	path, err := assets.WriteTemp("nri-topology-*.yaml", nriTopologyValues(workers))