Update tests to handle push notification implementation

3 files changed, 61 insertions, 10 deletions

diff --git a/internal/command/twofactorverify/twofactorverify_test.go b/internal/command/twofactorverify/twofactorverify_test.go
index 9d5f54d..2ee0226 100644
--- a/internal/command/twofactorverify/twofactorverify_test.go
+++ b/internal/command/twofactorverify/twofactorverify_test.go
@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -30,14 +31,31 @@ func setup(t *testing.T) []testserver.TestRequestHandler {
 				var requestBody *twofactorverify.RequestBody
 				require.NoError(t, json.Unmarshal(b, &requestBody))
 
+				var body map[string]interface{}
 				switch requestBody.KeyId {
 				case "1":
-					body := map[string]interface{}{
-						"success": true,
+					if requestBody.PushAuth {
+						body = map[string]interface{}{
+							"success": false,
+						}
+					} else {
+						body = map[string]interface{}{
+							"success": true,
+						}
+					}
+					json.NewEncoder(w).Encode(body)
+				case "2":
+					if requestBody.PushAuth {
+						body = map[string]interface{}{
+							"success": true,
+						}
+					} else {
+						// Stall verifyOTP long enough for pushAuth to complete
+						time.Sleep(10 * time.Second)
 					}
 					json.NewEncoder(w).Encode(body)
 				case "error":
-					body := map[string]interface{}{
+					body = map[string]interface{}{
 						"success": false,
 						"message": "error message",
 					}
@@ -70,6 +88,13 @@ func TestExecute(t *testing.T) {
 		expectedOutput string
 	}{
 		{
+			desc:      "When push is provided",
+			arguments: &commandargs.Shell{GitlabKeyId: "2"},
+			answer:    "",
+			expectedOutput: question +
+				"Push OTP validation successful. Git operations are now allowed.\n",
+		},
+		{
 			desc:      "With a known key id",
 			arguments: &commandargs.Shell{GitlabKeyId: "1"},
 			answer:    "123456\n",
@@ -105,8 +130,8 @@ func TestExecute(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			output := &bytes.Buffer{}
-			input := bytes.NewBufferString(tc.answer)
 
+			input := bytes.NewBufferString(tc.answer)
 			cmd := &Command{
 				Config:     &config.Config{GitlabUrl: url},
 				Args:       tc.arguments,
diff --git a/internal/gitlabnet/twofactorverify/client_test.go b/internal/gitlabnet/twofactorverify/client_test.go
index 7bb037e..22f5989 100644
--- a/internal/gitlabnet/twofactorverify/client_test.go
+++ b/internal/gitlabnet/twofactorverify/client_test.go
@@ -85,7 +85,7 @@ func TestVerifyOTPByKeyId(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabKeyId: "0"}
-	err := client.VerifyOTP(context.Background(), args, otpAttempt)
+	_, _, err := client.VerifyOTP(context.Background(), args, otpAttempt)
 	require.NoError(t, err)
 }
 
@@ -94,7 +94,7 @@ func TestVerifyOTPByUsername(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabUsername: "jane-doe"}
-	err := client.VerifyOTP(context.Background(), args, otpAttempt)
+	_, _, err := client.VerifyOTP(context.Background(), args, otpAttempt)
 	require.NoError(t, err)
 }
 
@@ -103,8 +103,8 @@ func TestErrorMessage(t *testing.T) {
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabKeyId: "1"}
-	err := client.VerifyOTP(context.Background(), args, otpAttempt)
-	require.Equal(t, "error message", err.Error())
+	_, reason, _ := client.VerifyOTP(context.Background(), args, otpAttempt)
+	require.Equal(t, "error message", reason)
 }
 
 func TestErrorResponses(t *testing.T) {
@@ -136,7 +136,7 @@ func TestErrorResponses(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			args := &commandargs.Shell{GitlabKeyId: tc.fakeId}
-			err := client.VerifyOTP(context.Background(), args, otpAttempt)
+			_, _, err := client.VerifyOTP(context.Background(), args, otpAttempt)
 
 			require.EqualError(t, err, tc.expectedError)
 		})
diff --git a/spec/gitlab_shell_two_factor_verify_spec.rb b/spec/gitlab_shell_two_factor_verify_spec.rb
index 25d8869..244d7d8 100644
--- a/spec/gitlab_shell_two_factor_verify_spec.rb
+++ b/spec/gitlab_shell_two_factor_verify_spec.rb
@@ -24,7 +24,15 @@ describe 'bin/gitlab-shell 2fa_verify' do
       key_id = params['key_id'] || params['user_id'].to_s
 
       if key_id == '100'
-        res.body = { success: true }.to_json
+        if params['push_auth']
+          res.body = { success: false }.to_json
+        else
+          res.body = { success: true }.to_json
+        end
+      elsif key_id == '102'
+        if params['push_auth']
+          res.body = { success: true }.to_json
+        end
       else
         res.body = { success: false, message: 'boom!' }.to_json
       end
@@ -38,6 +46,14 @@ describe 'bin/gitlab-shell 2fa_verify' do
   end
 
   describe 'command' do
+    context 'when push is provided' do
+      let(:cmd) { "#{gitlab_shell_path} key-102" }
+
+      it 'prints a successful push verification message' do
+        verify_successful_verification_push!(cmd)
+      end
+    end
+
     context 'when key is provided' do
       let(:cmd) { "#{gitlab_shell_path} key-100" }
 
@@ -78,4 +94,14 @@ describe 'bin/gitlab-shell 2fa_verify' do
       expect(stdout.flush.read).to eq("\nOTP validation successful. Git operations are now allowed.\n")
     end
   end
+
+  def verify_successful_verification_push!(cmd)
+    Open3.popen2(env, cmd) do |stdin, stdout|
+      expect(stdout.gets(5)).to eq('OTP: ')
+
+      stdin.puts('123456')
+
+      expect(stdout.flush.read).to eq("\nPush OTP validation successful. Git operations are now allowed.\n")
+    end
+  end
 end