From 9b8d4c9c45846f440fcf92be5093dc45c0982a93 Mon Sep 17 00:00:00 2001
From: Lorna Barber <SimbaGithub@magnitude.com>
Date: Thu, 6 Oct 2022 13:16:50 -0700
Subject: [PATCH] add unit test

---
 auth_test.go | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/auth_test.go b/auth_test.go
index 08df368e9..8ea27c7e1 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -226,6 +226,25 @@ func postAuthCheckJWTToken(_ context.Context, _ *snowflakeRestful, _ *url.Values
 	}, nil
 }
 
+func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+	var ar authRequest
+	if err := json.Unmarshal(jsonBody, &ar); err != nil {
+		return nil, err
+	}
+
+	return &authResponse{
+		Success: true,
+		Data: authResponseMain{
+			Token:       "t",
+			MasterToken: "m",
+			MfaToken:    "mockedMfaToken",
+			SessionInfo: authResponseSessionInfo{
+				DatabaseName: "dbn",
+			},
+		},
+	}, nil
+}
+
 func getDefaultSnowflakeConn() *snowflakeConn {
 	cfg := Config{
 		Account:            "a",
@@ -469,3 +488,20 @@ func TestUnitAuthenticateJWT(t *testing.T) {
 		t.Fatalf("invalid token passed")
 	}
 }
+
+func TestUnitAuthenticateUsernamePasswordMfa(t *testing.T) {
+	var err error
+	sr := &snowflakeRestful{
+		FuncPostAuth:  postAuthCheckUsernamePasswordMfa,
+		TokenAccessor: getSimpleTokenAccessor(),
+	}
+	sc := getDefaultSnowflakeConn()
+	sc.cfg.Authenticator = AuthTypeUsernamePasswordMFA
+	requestMfaToken := "true"
+	sc.cfg.Params[clientRequestMfaToken] = &requestMfaToken
+	sc.rest = sr
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
+	if err != nil {
+		t.Fatalf("failed to run. err: %v", err)
+	}
+}