Repository: kubernetes/kubernetes
File: test/e2e_node/dra_test.go (Go)

Code before:
		err: err,
	}
}

// Accept blocks until Close is called, then returns the configured error.
func (l *errorOnCloseListener) Accept() (net.Conn, error) {
	<-l.ch
	return nil, l.err
}

// Close unblocks Accept and causes it to return the configured error.
// It is safe to call multiple times.
func (l *errorOnCloseListener) Close() error {
	if l.closed.Swap(true) {
		return nil // already closed
	}
	close(l.ch)
	return nil
}

// Addr returns a dummy Unix address. Required to satisfy net.Listener.
func (*errorOnCloseListener) Addr() net.Addr {
	return &net.UnixAddr{Name: "errorOnCloseListener", Net: "unix"}
}

// driverNameLabelKey maps metric samples to their driver_name label value.
// Used with gstruct.MatchAllElements.
func driverNameLabelKey(element any) string {
	el := element.(*testutil.Sample)
	return string(el.Metric[testutil.LabelName("driver_name")])
}

// countingListener is a wrapper around net.Listener that counts the number of times
// Accept method is called. It embeds a net.Listener and uses atomic counters
// to safely track the number of accept operations.
type countingListener struct {
	net.Listener
	acceptCalls atomic.Int32
}

// newCountingListener creates a new Unix domain socket listener
// wrapped in a countingListener to track accepted connections.
func newCountingListener(socketPath string) (*countingListener, error) {
	addr, err := net.ResolveUnixAddr("unix", socketPath)
	if err != nil {
		return nil, err
	}

	l, err := net.ListenUnix("unix", addr)
	if err != nil {
		return nil, err

Review comment (refactor):
This looks good, thank you!

BTW, would it be possible to simply include this Status into the pod , so we can create pod with the required status in one step? That would simplify the code even more and eliminate any possibility for the race.

Code after:
					Command: []string{"/bin/sh", "-c", "sleep 600"},
					Resources: v1.ResourceRequirements{
						Claims: []v1.ResourceClaim{{Name: claimName, Request: "my-request"}},
					},
				},
			},
		},
	}
	// Create the pod on the API server to assign the real UID.
	createdPod, err := f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(ctx, pod, metav1.CreateOptions{})
	framework.ExpectNoError(err, "failed to create Pod "+podName)
	ginkgo.DeferCleanup(func(ctx context.Context) {
		e2epod.DeletePodOrFail(ctx, f.ClientSet, createdPod.Namespace, createdPod.Name)
	})

	// Patch the Pod's status to include the ResourceClaimStatuses, mimicking the scheduler.
	patch, err := json.Marshal(v1.Pod{
		Status: v1.PodStatus{
			ResourceClaimStatuses: []v1.PodResourceClaimStatus{
				{Name: claimName, ResourceClaimName: &claimName},
			},
		},
	})
	framework.ExpectNoError(err, "failed to marshal patch for Pod status")
	_, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Patch(ctx, createdPod.Name, apitypes.StrategicMergePatchType, patch, metav1.PatchOptions{}, "status")
	framework.ExpectNoError(err, "failed to patch Pod status with ResourceClaimStatuses")

	ginkgo.By(fmt.Sprintf("Allocating claim %q to pod %q with its real UID", claimName, podName))
	// Get the created claim to ensure the latest version before updating.
	claimToUpdate, err := f.ClientSet.ResourceV1().ResourceClaims(f.Namespace.Name).Get(ctx, claimName, metav1.GetOptions{})
	framework.ExpectNoError(err, "failed to get latest version of ResourceClaim "+claimName)

	// Update the claims status to reserve it for the *real* pod UID.
	claimToUpdate.Status = resourceapi.ResourceClaimStatus{
		ReservedFor: []resourceapi.ResourceClaimConsumerReference{
			{Resource: "pods", Name: createdPod.Name, UID: createdPod.UID},
		},
		Allocation: &resourceapi.AllocationResult{
			Devices: resourceapi.DeviceAllocationResult{
				Results: []resourceapi.DeviceRequestAllocationResult{
					{
						Driver:  driverName,
						Pool:    poolName,
						Device:  deviceName,
						Request: "my-request",
					},
				},
			},
		},
	}
	_, err = f.ClientSet.ResourceV1().ResourceClaims(f.Namespace.Name).UpdateStatus(ctx, claimToUpdate, metav1.UpdateOptions{})