Fix unit test

Emyrk · Emyrk · commit a343da91b762 · 2022-09-01T10:51:15.000-04:00
diff --git a/coderd/userauth_test.go b/coderd/userauth_test.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"strings"
 	"testing"
 	"time"
 
@@ -244,7 +245,7 @@ func TestUserOAuth2Github(t *testing.T) {
 		resp := oauth2Callback(t, client)
 		require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode)
 
-		client.SessionToken = resp.Cookies()[0].Value
+		client.SessionToken = authCookieValue(resp.Cookies())
 		user, err := client.User(context.Background(), "me")
 		require.NoError(t, err)
 		require.Equal(t, "kyle@coder.com", user.Email)
@@ -374,7 +375,7 @@ func TestUserOIDC(t *testing.T) {
 			defer cancel()
 
 			if tc.Username != "" {
-				client.SessionToken = resp.Cookies()[0].Value
+				client.SessionToken = authCookieValue(resp.Cookies())
 				user, err := client.User(ctx, "me")
 				require.NoError(t, err)
 				require.Equal(t, tc.Username, user.Username)
@@ -503,3 +504,12 @@ func oidcCallback(t *testing.T, client *codersdk.Client) *http.Response {
 func i64ptr(i int64) *int64 {
 	return &i
 }
+
+func authCookieValue(cookies []*http.Cookie) string {
+	for _, cookie := range cookies {
+		if strings.Contains(cookie.Name, "token") {
+			return cookie.Value
+		}
+	}
+	return ""
+}
diff --git a/coderd/users_test.go b/coderd/users_test.go
@@ -281,10 +281,14 @@ func TestPostLogout(t *testing.T) {
 		require.Equal(t, http.StatusOK, res.StatusCode)
 
 		cookies := res.Cookies()
-		require.Len(t, cookies, 1, "Exactly one cookie should be returned")
+		require.Len(t, cookies, 2, "Exactly one cookie should be returned")
 
-		require.Equal(t, codersdk.SessionTokenKey, cookies[0].Name, "Cookie should be the auth cookie")
-		require.Equal(t, -1, cookies[0].MaxAge, "Cookie should be set to delete")
+		for _, cookie := range cookies {
+			if cookie.Name == codersdk.SessionTokenKey {
+				require.Equal(t, codersdk.SessionTokenKey, cookies[0].Name, "Cookie should be the auth cookie")
+				require.Equal(t, -1, cookies[0].MaxAge, "Cookie should be set to delete")
+			}
+		}
 
 		_, err = client.GetAPIKey(ctx, admin.UserID.String(), keyID)
 		sdkErr := &codersdk.Error{}
