2 files changed, 10 insertions, 8 deletions
diff --git a/internal/command/shared/customaction/customaction.go b/internal/command/shared/customaction/customaction.go
index 2ba1091..0675d36 100644
--- a/internal/command/shared/customaction/customaction.go
+++ b/internal/command/shared/customaction/customaction.go
@@ -2,6 +2,7 @@ package customaction
 
 import (
 	"bytes"
+	"context"
 	"errors"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client"
@@ -34,7 +35,7 @@ type Command struct {
 	EOFSent    bool
 }
 
-func (c *Command) Execute(response *accessverifier.Response) error {
+func (c *Command) Execute(ctx context.Context, response *accessverifier.Response) error {
 	data := response.Payload.Data
 	apiEndpoints := data.ApiEndpoints
 
@@ -42,10 +43,10 @@ func (c *Command) Execute(response *accessverifier.Response) error {
 		return errors.New("Custom action error: Empty API endpoints")
 	}
 
-	return c.processApiEndpoints(response)
+	return c.processApiEndpoints(ctx, response)
 }
 
-func (c *Command) processApiEndpoints(response *accessverifier.Response) error {
+func (c *Command) processApiEndpoints(ctx context.Context, response *accessverifier.Response) error {
 	client, err := gitlabnet.GetClient(c.Config)
 
 	if err != nil {
@@ -64,7 +65,7 @@ func (c *Command) processApiEndpoints(response *accessverifier.Response) error {
 
 		log.WithFields(fields).Info("Performing custom action")
 
-		response, err := c.performRequest(client, endpoint, request)
+		response, err := c.performRequest(ctx, client, endpoint, request)
 		if err != nil {
 			return err
 		}
@@ -95,8 +96,8 @@ func (c *Command) processApiEndpoints(response *accessverifier.Response) error {
 	return nil
 }
 
-func (c *Command) performRequest(client *client.GitlabNetClient, endpoint string, request *Request) (*Response, error) {
-	response, err := client.DoRequest(http.MethodPost, endpoint, request)
+func (c *Command) performRequest(ctx context.Context, client *client.GitlabNetClient, endpoint string, request *Request) (*Response, error) {
+	response, err := client.DoRequest(ctx, http.MethodPost, endpoint, request)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/command/shared/customaction/customaction_test.go b/internal/command/shared/customaction/customaction_test.go
index 46c5f32..119da5b 100644
--- a/internal/command/shared/customaction/customaction_test.go
+++ b/internal/command/shared/customaction/customaction_test.go
@@ -2,6 +2,7 @@ package customaction
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
@@ -78,7 +79,7 @@ func TestExecuteEOFSent(t *testing.T) {
 		EOFSent:    true,
 	}
 
-	require.NoError(t, cmd.Execute(response))
+	require.NoError(t, cmd.Execute(context.Background(), response))
 
 	// expect printing of info message, "custom" string from the first request
 	// and "output" string from the second request
@@ -148,7 +149,7 @@ func TestExecuteNoEOFSent(t *testing.T) {
 		EOFSent:    false,
 	}
 
-	require.NoError(t, cmd.Execute(response))
+	require.NoError(t, cmd.Execute(context.Background(), response))
 
 	// expect printing of info message, "custom" string from the first request
 	// and "output" string from the second request