2 files changed, 11 insertions, 9 deletions

diff --git a/internal/gitlabnet/twofactorrecover/client.go b/internal/gitlabnet/twofactorrecover/client.go
index d22daca..456f892 100644
--- a/internal/gitlabnet/twofactorrecover/client.go
+++ b/internal/gitlabnet/twofactorrecover/client.go
@@ -1,6 +1,7 @@
 package twofactorrecover
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -37,14 +38,14 @@ func NewClient(config *config.Config) (*Client, error) {
 	return &Client{config: config, client: client}, nil
 }
 
-func (c *Client) GetRecoveryCodes(args *commandargs.Shell) ([]string, error) {
-	requestBody, err := c.getRequestBody(args)
+func (c *Client) GetRecoveryCodes(ctx context.Context, args *commandargs.Shell) ([]string, error) {
+	requestBody, err := c.getRequestBody(ctx, args)
 
 	if err != nil {
 		return nil, err
 	}
 
-	response, err := c.client.Post("/two_factor_recovery_codes", requestBody)
+	response, err := c.client.Post(ctx, "/two_factor_recovery_codes", requestBody)
 	if err != nil {
 		return nil, err
 	}
@@ -66,7 +67,7 @@ func parse(hr *http.Response) ([]string, error) {
 	return response.RecoveryCodes, nil
 }
 
-func (c *Client) getRequestBody(args *commandargs.Shell) (*RequestBody, error) {
+func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell) (*RequestBody, error) {
 	client, err := discover.NewClient(c.config)
 
 	if err != nil {
@@ -77,7 +78,7 @@ func (c *Client) getRequestBody(args *commandargs.Shell) (*RequestBody, error) {
 	if args.GitlabKeyId != "" {
 		requestBody = &RequestBody{KeyId: args.GitlabKeyId}
 	} else {
-		userInfo, err := client.GetByCommandArgs(args)
+		userInfo, err := client.GetByCommandArgs(ctx, args)
 
 		if err != nil {
 			return nil, err
diff --git a/internal/gitlabnet/twofactorrecover/client_test.go b/internal/gitlabnet/twofactorrecover/client_test.go
index 372afec..46291aa 100644
--- a/internal/gitlabnet/twofactorrecover/client_test.go
+++ b/internal/gitlabnet/twofactorrecover/client_test.go
@@ -1,6 +1,7 @@
 package twofactorrecover
 
 import (
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
@@ -85,7 +86,7 @@ func TestGetRecoveryCodesByKeyId(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabKeyId: "0"}
-	result, err := client.GetRecoveryCodes(args)
+	result, err := client.GetRecoveryCodes(context.Background(), args)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"recovery 1", "codes 1"}, result)
 }
@@ -95,7 +96,7 @@ func TestGetRecoveryCodesByUsername(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabUsername: "jane-doe"}
-	result, err := client.GetRecoveryCodes(args)
+	result, err := client.GetRecoveryCodes(context.Background(), args)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"recovery 2", "codes 2"}, result)
 }
@@ -105,7 +106,7 @@ func TestMissingUser(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabKeyId: "1"}
-	_, err := client.GetRecoveryCodes(args)
+	_, err := client.GetRecoveryCodes(context.Background(), args)
 	assert.Equal(t, "missing user", err.Error())
 }
 
@@ -138,7 +139,7 @@ func TestErrorResponses(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			args := &commandargs.Shell{GitlabKeyId: tc.fakeId}
-			resp, err := client.GetRecoveryCodes(args)
+			resp, err := client.GetRecoveryCodes(context.Background(), args)
 
 			assert.EqualError(t, err, tc.expectedError)
 			assert.Nil(t, resp)