twofactorrecover_test.go 3.75 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
package twofactorrecover

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/twofactorrecover"
)

var (
21
	requests []testserver.TestRequestHandler
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
)

func setup(t *testing.T) {
	requests = []testserver.TestRequestHandler{
		{
			Path: "/api/v4/internal/two_factor_recovery_codes",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				b, err := ioutil.ReadAll(r.Body)
				defer r.Body.Close()

				require.NoError(t, err)

				var requestBody *twofactorrecover.RequestBody
				json.Unmarshal(b, &requestBody)

				switch requestBody.KeyId {
				case "1":
					body := map[string]interface{}{
						"success":        true,
						"recovery_codes": [2]string{"recovery", "codes"},
					}
					json.NewEncoder(w).Encode(body)
				case "forbidden":
					body := map[string]interface{}{
						"success": false,
						"message": "Forbidden!",
					}
					json.NewEncoder(w).Encode(body)
				case "broken":
					w.WriteHeader(http.StatusInternalServerError)
				}
			},
		},
	}
}

const (
	question = "Are you sure you want to generate new two-factor recovery codes?\n" +
		"Any existing recovery codes you saved will be invalidated. (yes/no)\n\n"
	errorHeader = "An error occurred while trying to generate new recovery codes.\n"
)

func TestExecute(t *testing.T) {
	setup(t)

67
	url, cleanup := testserver.StartSocketHttpServer(t, requests)
68 69 70 71
	defer cleanup()

	testCases := []struct {
		desc           string
72
		arguments      *commandargs.Shell
73 74 75 76 77
		answer         string
		expectedOutput string
	}{
		{
			desc:      "With a known key id",
78
			arguments: &commandargs.Shell{GitlabKeyId: "1"},
79 80 81 82 83 84 85 86 87
			answer:    "yes\n",
			expectedOutput: question +
				"Your two-factor authentication recovery codes are:\n\nrecovery\ncodes\n\n" +
				"During sign in, use one of the codes above when prompted for\n" +
				"your two-factor code. Then, visit your Profile Settings and add\n" +
				"a new device so you do not lose access to your account again.\n",
		},
		{
			desc:           "With bad response",
88
			arguments:      &commandargs.Shell{GitlabKeyId: "-1"},
89 90 91 92 93
			answer:         "yes\n",
			expectedOutput: question + errorHeader + "Parsing failed\n",
		},
		{
			desc:           "With API returns an error",
94
			arguments:      &commandargs.Shell{GitlabKeyId: "forbidden"},
95 96 97 98 99
			answer:         "yes\n",
			expectedOutput: question + errorHeader + "Forbidden!\n",
		},
		{
			desc:           "With API fails",
100
			arguments:      &commandargs.Shell{GitlabKeyId: "broken"},
101 102 103 104 105
			answer:         "yes\n",
			expectedOutput: question + errorHeader + "Internal API error (500)\n",
		},
		{
			desc:           "With missing arguments",
106
			arguments:      &commandargs.Shell{},
107 108 109 110 111
			answer:         "yes\n",
			expectedOutput: question + errorHeader + "who='' is invalid\n",
		},
		{
			desc:      "With negative answer",
112
			arguments: &commandargs.Shell{},
113 114 115 116 117 118 119 120 121 122 123
			answer:    "no\n",
			expectedOutput: question +
				"New recovery codes have *not* been generated. Existing codes will remain valid.\n",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			output := &bytes.Buffer{}
			input := bytes.NewBufferString(tc.answer)

124 125 126 127 128
			cmd := &Command{
				Config:     &config.Config{GitlabUrl: url},
				Args:       tc.arguments,
				ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
			}
129

130
			err := cmd.Execute()
131 132 133 134 135 136

			assert.NoError(t, err)
			assert.Equal(t, tc.expectedOutput, output.String())
		})
	}
}