From 2ea2fb3166f325ff3ad532b7ec538d5dd110e2a5 Mon Sep 17 00:00:00 2001
From: Ed Bartosh <eduard.bartosh@intel.com>
Date: Sat, 25 May 2024 00:22:05 +0300
Subject: [PATCH] e2e: test-driver: implement failure mode

---
 test/e2e/dra/test-driver/app/kubeletplugin.go | 62 +++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/test/e2e/dra/test-driver/app/kubeletplugin.go b/test/e2e/dra/test-driver/app/kubeletplugin.go
index ac1f2bfa69f..47d9dbfb00e 100644
--- a/test/e2e/dra/test-driver/app/kubeletplugin.go
+++ b/test/e2e/dra/test-driver/app/kubeletplugin.go
@@ -55,6 +55,12 @@ type ExamplePlugin struct {
 	gRPCCalls      []GRPCCall
 
 	block bool
+
+	prepareResourcesFailure   error
+	failPrepareResourcesMutex sync.Mutex
+
+	unprepareResourcesFailure   error
+	failUnprepareResourcesMutex sync.Mutex
 }
 
 type GRPCCall struct {
@@ -168,6 +174,52 @@ func (ex *ExamplePlugin) Block() {
 	ex.block = true
 }
 
+func (ex *ExamplePlugin) withLock(mutex *sync.Mutex, f func()) {
+	mutex.Lock()
+	f()
+	mutex.Unlock()
+}
+
+// SetNodePrepareResourcesFailureMode sets the failure mode for NodePrepareResources call
+// and returns a function to unset the failure mode
+func (ex *ExamplePlugin) SetNodePrepareResourcesFailureMode() func() {
+	ex.failPrepareResourcesMutex.Lock()
+	ex.prepareResourcesFailure = errors.New("simulated PrepareResources failure")
+	ex.failPrepareResourcesMutex.Unlock()
+
+	return func() {
+		ex.failPrepareResourcesMutex.Lock()
+		ex.prepareResourcesFailure = nil
+		ex.failPrepareResourcesMutex.Unlock()
+	}
+}
+
+func (ex *ExamplePlugin) getPrepareResourcesFailure() error {
+	ex.failPrepareResourcesMutex.Lock()
+	defer ex.failPrepareResourcesMutex.Unlock()
+	return ex.prepareResourcesFailure
+}
+
+// SetNodeUnprepareResourcesFailureMode sets the failure mode for NodeUnprepareResources call
+// and returns a function to unset the failure mode
+func (ex *ExamplePlugin) SetNodeUnprepareResourcesFailureMode() func() {
+	ex.failUnprepareResourcesMutex.Lock()
+	ex.unprepareResourcesFailure = errors.New("simulated UnprepareResources failure")
+	ex.failUnprepareResourcesMutex.Unlock()
+
+	return func() {
+		ex.failUnprepareResourcesMutex.Lock()
+		ex.unprepareResourcesFailure = nil
+		ex.failUnprepareResourcesMutex.Unlock()
+	}
+}
+
+func (ex *ExamplePlugin) getUnprepareResourcesFailure() error {
+	ex.failUnprepareResourcesMutex.Lock()
+	defer ex.failUnprepareResourcesMutex.Unlock()
+	return ex.unprepareResourcesFailure
+}
+
 // NodePrepareResource ensures that the CDI file for the claim exists. It uses
 // a deterministic name to simplify NodeUnprepareResource (no need to remember
 // or discover the name) and idempotency (when called again, the file simply
@@ -309,6 +361,11 @@ func (ex *ExamplePlugin) NodePrepareResources(ctx context.Context, req *drapbv1a
 	resp := &drapbv1alpha3.NodePrepareResourcesResponse{
 		Claims: make(map[string]*drapbv1alpha3.NodePrepareResourceResponse),
 	}
+
+	if failure := ex.getPrepareResourcesFailure(); failure != nil {
+		return resp, failure
+	}
+
 	for _, claimReq := range req.Claims {
 		cdiDevices, err := ex.nodePrepareResource(ctx, claimReq.Name, claimReq.Uid, claimReq.ResourceHandle, claimReq.StructuredResourceHandle)
 		if err != nil {
@@ -381,6 +438,11 @@ func (ex *ExamplePlugin) NodeUnprepareResources(ctx context.Context, req *drapbv
 	resp := &drapbv1alpha3.NodeUnprepareResourcesResponse{
 		Claims: make(map[string]*drapbv1alpha3.NodeUnprepareResourceResponse),
 	}
+
+	if failure := ex.getUnprepareResourcesFailure(); failure != nil {
+		return resp, failure
+	}
+
 	for _, claimReq := range req.Claims {
 		err := ex.nodeUnprepareResource(ctx, claimReq.Name, claimReq.Uid, claimReq.ResourceHandle, claimReq.StructuredResourceHandle)
 		if err != nil {