GODT-1166: Reduce the number of auth for live test

- Changed: Do not reauth controller clients.
- Changed: Verbosisty is set only once before run
- Changed: AddUser takes TestAccount as argument
- Added: Setup/clean up before/after test run
- Added: Access to the current refresh token from pmapi.Client interface.
- Added: Context function to add test a user to bridge without login, just call users.FinishLogin.
- Added: PMAPIController.GetAuthClient returns authenticated client for username.
- Added: Persistent clients does not loggout after every scenario.
- Changed: Disabled no-internet tests.

pkg/pmapi/auth.go

@@ -217,3 +217,13 @@ func randomString(length int) string {
 
 	return base64.StdEncoding.EncodeToString(noise)[:length]
 }
+
+func (c *client) GetCurrentAuth() *Auth {
+	return &Auth{
+		UserID: c.user.ID,
+		AuthRefresh: AuthRefresh{
+			UID:          c.uid,
+			RefreshToken: c.ref,
+		},
+	}
+}

pkg/pmapi/client_types.go

@@ -73,6 +73,8 @@ type Client interface {
 
 	KeyRingForAddressID(string) (kr *crypto.KeyRing, err error)
 	GetPublicKeysForEmail(context.Context, string) ([]PublicKey, bool, error)
+
+	GetCurrentAuth() *Auth
 }
 
 type AuthRefreshHandler func(*AuthRefresh)

pkg/pmapi/mocks/mocks.go

@@ -301,6 +301,20 @@ func (mr *MockClientMockRecorder) GetContactEmailByEmail(arg0, arg1, arg2, arg3
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetContactEmailByEmail", reflect.TypeOf((*MockClient)(nil).GetContactEmailByEmail), arg0, arg1, arg2, arg3)
 }
 
+// GetCurrentAuth mocks base method
+func (m *MockClient) GetCurrentAuth() *pmapi.Auth {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetCurrentAuth")
+	ret0, _ := ret[0].(*pmapi.Auth)
+	return ret0
+}
+
+// GetCurrentAuth indicates an expected call of GetCurrentAuth
+func (mr *MockClientMockRecorder) GetCurrentAuth() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetCurrentAuth", reflect.TypeOf((*MockClient)(nil).GetCurrentAuth))
+}
+
 // GetEvent mocks base method
 func (m *MockClient) GetEvent(arg0 context.Context, arg1 string) (*pmapi.Event, error) {
 	m.ctrl.T.Helper()

test/Makefile

@@ -1,7 +1,7 @@
 .PHONY: check-go check-godog install-godog test test-bridge test-ie test-live test-live-bridge test-live-ie test-stage test-debug test-live-debug bench
 
 export GO111MODULE=on
-export BRIDGE_VERSION:=1.5.5+integrationtests
+export BRIDGE_VERSION:=1.8.2+integrationtests
 export VERBOSITY?=fatal
 export TEST_DATA=testdata
 export TEST_APP?=bridge

test/bdd_test.go

@@ -29,6 +29,9 @@ const (
 )
 
 func FeatureContext(s *godog.Suite) {
+	s.BeforeSuite(context.BeforeRun)
+	s.AfterSuite(context.AfterRun)
+
 	s.BeforeScenario(beforeScenario)
 	s.AfterScenario(afterScenario)
 

test/benchmarks/bench_test.go

@@ -35,7 +35,7 @@ func benchTestContext() (*context.TestContext, *mocks.IMAPClient) {
 		panic("account " + username + " does not exist")
 	}
 
-	_ = ctx.GetPMAPIController().AddUser(account.User(), account.Addresses(), account.Password(), account.IsTwoFAEnabled())
+	_ = ctx.GetPMAPIController().AddUser(account)
 	if err := ctx.LoginUser(account.Username(), account.Password(), account.MailboxPassword()); err != nil {
 		panic(err)
 	}

test/context/context.go

@@ -94,8 +94,6 @@ type TestContext struct {
 
 // New returns a new test TestContext.
 func New(app string) *TestContext {
-	setLogrusVerbosityFromEnv()
-
 	listener := listener.New()
 	pmapiController, clientManager := newPMAPIController(app, listener)
 

test/context/globals.go

@@ -0,0 +1,40 @@
+// Copyright (c) 2021 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package context
+
+import (
+	"os"
+
+	"github.com/ProtonMail/proton-bridge/test/liveapi"
+)
+
+// BeforeRun does necessary setup.
+func BeforeRun() {
+	setLogrusVerbosityFromEnv()
+
+	if os.Getenv(EnvName) == EnvLive {
+		liveapi.SetupPersistentClients()
+	}
+}
+
+// AfterRun does necessary cleanup.
+func AfterRun() {
+	if os.Getenv(EnvName) == EnvLive {
+		liveapi.CleanupPersistentClients()
+	}
+}

test/context/pmapi_controller.go

@@ -23,6 +23,7 @@ import (
 	"github.com/ProtonMail/proton-bridge/internal/events"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/ProtonMail/proton-bridge/test/accounts"
 	"github.com/ProtonMail/proton-bridge/test/fakeapi"
 	"github.com/ProtonMail/proton-bridge/test/liveapi"
 )
@@ -30,7 +31,8 @@ import (
 type PMAPIController interface {
 	TurnInternetConnectionOff()
 	TurnInternetConnectionOn()
-	AddUser(user *pmapi.User, addresses *pmapi.AddressList, password []byte, twoFAEnabled bool) error
+	GetAuthClient(username string) pmapi.Client
+	AddUser(account *accounts.TestAccount) error
 	AddUserLabel(username string, label *pmapi.Label) error
 	GetLabelIDs(username string, labelNames []string) ([]string, error)
 	AddUserMessage(username string, message *pmapi.Message) (string, error)

test/context/users.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ProtonMail/go-srp"
 	"github.com/ProtonMail/proton-bridge/internal/store"
 	"github.com/ProtonMail/proton-bridge/internal/users"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 )
@@ -61,6 +62,18 @@ func (ctx *TestContext) LoginUser(username string, password, mailboxPassword []b
 	return nil
 }
 
+// FinishLogin prevents authentication if not necessary.
+func (ctx *TestContext) FinishLogin(client pmapi.Client, mailboxPassword string) error {
+	user, err := ctx.users.FinishLogin(client, client.GetCurrentAuth(), mailboxPassword)
+	if err != nil {
+		return errors.Wrap(err, "failed to finish login")
+	}
+
+	ctx.addCleanupChecked(user.Logout, "Logging out user")
+
+	return nil
+}
+
 // GetUser retrieves the bridge user matching the given query string.
 func (ctx *TestContext) GetUser(username string) (*users.User, error) {
 	return ctx.users.GetUser(username)

test/fakeapi/auth.go

@@ -63,3 +63,13 @@ func (api *FakePMAPI) AuthDelete(_ context.Context) error {
 
 	return nil
 }
+
+func (api *FakePMAPI) GetCurrentAuth() *pmapi.Auth {
+	return &pmapi.Auth{
+		UserID: api.userID,
+		AuthRefresh: pmapi.AuthRefresh{
+			UID:          api.uid,
+			RefreshToken: api.ref,
+		},
+	}
+}

test/fakeapi/controller.go

@@ -24,6 +24,8 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+// Controller implements dummy PMAPIController interface without actual
+// endpoint.
 type Controller struct {
 	// Internal states.
 	lock               *sync.RWMutex

test/fakeapi/controller_control.go

@@ -22,8 +22,10 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/ProtonMail/proton-bridge/test/accounts"
 )
 
 var systemLabelNameToID = map[string]string{ //nolint[gochecknoglobals]
@@ -61,13 +63,15 @@ func (ctl *Controller) ReorderAddresses(user *pmapi.User, addressIDs []string) e
 	return api.ReorderAddresses(context.Background(), addressIDs)
 }
 
-func (ctl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList, password []byte, twoFAEnabled bool) error {
-	ctl.usersByUsername[user.Name] = &fakeUser{
-		user:     user,
-		password: password,
-		has2FA:   twoFAEnabled,
+func (ctl *Controller) AddUser(account *accounts.TestAccount) error {
+	ctl.usersByUsername[account.User().Name] = &fakeUser{
+		user:     account.User(),
+		password: account.Password(),
+		has2FA:   account.IsTwoFAEnabled(),
 	}
-	ctl.addressesByUsername[user.Name] = addresses
+	ctl.addressesByUsername[account.User().Name] = account.Addresses()
+	ctl.createSession(account.User().Name, true)
+
 	return nil
 }
 
@@ -181,3 +185,15 @@ func (ctl *Controller) GetMessages(username, labelID string) ([]*pmapi.Message,
 	}
 	return messages, nil
 }
+
+func (ctl *Controller) GetAuthClient(username string) pmapi.Client {
+	for uid, session := range ctl.sessionsByUID {
+		if session.username == username {
+			return ctl.clientManager.NewClient(uid, session.acc, session.ref, time.Now())
+		}
+	}
+
+	ctl.log.WithField("username", username).Fatal("Cannot get authenticated client.")
+
+	return nil
+}

test/fakeapi/controller_session.go

@@ -51,24 +51,25 @@ func (ctl *Controller) checkScope(uid string) bool {
 }
 
 func (ctl *Controller) createSessionIfAuthorized(username string, password []byte) (*fakeSession, error) {
-	// get user
 	user, ok := ctl.usersByUsername[username]
 	if !ok || !bytes.Equal(user.password, password) {
 		return nil, errWrongNameOrPassword
 	}
 
-	// create session
+	return ctl.createSession(username, !user.has2FA), nil
+}
+
+func (ctl *Controller) createSession(username string, hasFullScope bool) *fakeSession {
 	session := &fakeSession{
 		username:     username,
 		uid:          ctl.tokenGenerator.next("uid"),
 		acc:          ctl.tokenGenerator.next("acc"),
 		ref:          ctl.tokenGenerator.next("ref"),
-		hasFullScope: !user.has2FA,
+		hasFullScope: hasFullScope,
 	}
 
 	ctl.sessionsByUID[session.uid] = session
-
-	return session, nil
+	return session
 }
 
 func (ctl *Controller) refreshSessionIfAuthorized(uid, ref string) (*fakeSession, error) {

test/features/bridge/imap/auth.feature

@@ -15,7 +15,7 @@ Feature: IMAP auth
     Then IMAP response is "IMAP error: NO account is logged out, use the app to login again"
 
   Scenario: Authenticates with connected user that was loaded without internet
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is no internet connection
     When bridge starts
     And the internet connection is restored
@@ -28,13 +28,13 @@ Feature: IMAP auth
     Then "user" is connected
 
   Scenario: Authenticates with freshly logged-out user
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     When "user" logs out
     And IMAP client authenticates "user"
     Then IMAP response is "IMAP error: NO account is logged out, use the app to login again"
 
   Scenario: Authenticates user which was re-logged in
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     When "user" logs out
     And IMAP client authenticates "user"
     Then IMAP response is "IMAP error: NO account is logged out, use the app to login again"

test/features/bridge/no_internet.feature

@@ -1,15 +1,17 @@
 Feature: Servers are closed when no internet
 
+  # FIXME: Locally works, has lags on CI. Looks like it breaks other tests as well.
+  @ignore
   Scenario: All connection are closed and then restored multiple times
     Given there is connected user "user"
     And there is IMAP client "i1" logged in as "user"
     And there is SMTP client "s1" logged in as "user"
     When there is no internet connection
-    And 1 second pass
+    And 3 seconds pass
     Then IMAP client "i1" is logged out
     And SMTP client "s1" is logged out
     Given the internet connection is restored
-    And 1 second pass
+    And 3 seconds pass
     And there is IMAP client "i2" logged in as "user"
     And there is SMTP client "s2" logged in as "user"
     When IMAP client "i2" gets info of "INBOX"
@@ -17,11 +19,11 @@ Feature: Servers are closed when no internet
     Then IMAP response to "i2" is "OK"
     Then SMTP response to "s2" is "OK"
     When there is no internet connection
-    And 1 second pass
+    And 3 seconds pass
     Then IMAP client "i2" is logged out
     And SMTP client "s2" is logged out
     Given the internet connection is restored
-    And 1 second pass
+    And 3 seconds pass
     And there is IMAP client "i3" logged in as "user"
     And there is SMTP client "s3" logged in as "user"
     When IMAP client "i3" gets info of "INBOX"

test/features/bridge/start.feature

@@ -1,6 +1,6 @@
 Feature: Start bridge
   Scenario: Start with connected user, database file and internet connection
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is database file for "user"
     When bridge starts
     Then "user" is connected
@@ -8,7 +8,7 @@ Feature: Start bridge
     And "user" has running event loop
 
   Scenario: Start with connected user, database file and no internet connection
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is database file for "user"
     And there is no internet connection
     When bridge starts
@@ -17,7 +17,7 @@ Feature: Start bridge
     And "user" has running event loop
 
   Scenario: Start with connected user, no database file and internet connection
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is no database file for "user"
     When bridge starts
     Then "user" is connected
@@ -25,7 +25,7 @@ Feature: Start bridge
     And "user" has running event loop
 
   Scenario: Start with connected user, no database file and no internet connection
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is no database file for "user"
     And there is no internet connection
     When bridge starts

test/features/bridge/users/delete.feature

@@ -1,18 +1,18 @@
 Feature: Delete user
   Scenario: Deleting connected user
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     When user deletes "user"
     Then last response is "OK"
     And "user" has database file
 
   Scenario: Deleting connected user with cache
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     When user deletes "user" with cache
     Then last response is "OK"
     And "user" does not have database file
 
   Scenario: Deleting connected user without database file
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is no database file for "user"
     When user deletes "user" with cache
     Then last response is "OK"

test/features/bridge/users/relogin.feature

@@ -1,6 +1,6 @@
 Feature: Re-login
   Scenario: Re-login with connected user and database file
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is database file for "user"
     When "user" logs in
     Then last response is "failed to finish login: user is already connected"
@@ -9,7 +9,7 @@ Feature: Re-login
 
   @ignore
   Scenario: Re-login with connected user and no database file
-    Given there is connected user "user"
+    Given there is user "user" which just logged in
     And there is no database file for "user"
     When "user" logs in
     Then last response is "failed to finish login: user is already connected"

test/liveapi/controller.go

@@ -21,45 +21,36 @@ import (
 	"net/http"
 	"sync"
 
-	"github.com/ProtonMail/proton-bridge/internal/constants"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/sirupsen/logrus"
 )
 
+// Controller implements PMAPIController interface for specified endpoint.
 type Controller struct {
+	log *logrus.Entry
 	// Internal states.
 	lock                 *sync.RWMutex
 	calls                []*fakeCall
-	pmapiByUsername      map[string]pmapi.Client
 	messageIDsByUsername map[string][]string
-	clientManager        pmapi.Manager
 
 	// State controlled by test.
 	noInternetConnection bool
 }
 
-func NewController(app string) (*Controller, pmapi.Manager) {
-	cm := pmapi.New(pmapi.NewConfig(getAppVersionName(app), constants.Version))
+func NewController(_ string) (*Controller, pmapi.Manager) {
 	controller := &Controller{
+		log:                  logrus.WithField("pkg", "live-controller"),
 		lock:                 &sync.RWMutex{},
 		calls:                []*fakeCall{},
-		pmapiByUsername:      map[string]pmapi.Client{},
 		messageIDsByUsername: map[string][]string{},
-		clientManager:        cm,
 
 		noInternetConnection: false,
 	}
 
-	cm.SetTransport(&fakeTransport{
+	persistentClients.manager.SetTransport(&fakeTransport{
 		ctl:       controller,
 		transport: http.DefaultTransport,
 	})
 
-	return controller, cm
-}
-
-func getAppVersionName(app string) string {
-	if app == "ie" {
-		return "importExport"
-	}
-	return app
+	return controller, persistentClients.manager
 }

test/liveapi/labels.go

@@ -37,9 +37,9 @@ var systemLabelNameToID = map[string]string{ //nolint[gochecknoglobals]
 }
 
 func (ctl *Controller) AddUserLabel(username string, label *pmapi.Label) error {
-	client, ok := ctl.pmapiByUsername[username]
-	if !ok {
-		return fmt.Errorf("user %s does not exist", username)
+	client, err := getPersistentClient(username)
+	if err != nil {
+		return err
 	}
 
 	label.Exclusive = getLabelExclusive(label.Name)
@@ -68,9 +68,9 @@ func (ctl *Controller) getLabelID(username, labelName string) (string, error) {
 		return labelID, nil
 	}
 
-	client, ok := ctl.pmapiByUsername[username]
-	if !ok {
-		return "", fmt.Errorf("user %s does not exist", username)
+	client, err := getPersistentClient(username)
+	if err != nil {
+		return "", err
 	}
 
 	labels, err := client.ListLabels(context.Background())

test/liveapi/messages.go

@@ -19,7 +19,6 @@ package liveapi
 
 import (
 	"context"
-	"fmt"
 
 	messageUtils "github.com/ProtonMail/proton-bridge/pkg/message"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
@@ -31,9 +30,9 @@ func (ctl *Controller) AddUserMessage(username string, message *pmapi.Message) (
 		return "", errors.New("add user messages with attachments is not implemented for live")
 	}
 
-	client, ok := ctl.pmapiByUsername[username]
-	if !ok {
-		return "", fmt.Errorf("user %s does not exist", username)
+	client, err := getPersistentClient(username)
+	if err != nil {
+		return "", err
 	}
 
 	if message.Flags == 0 {
@@ -75,9 +74,9 @@ func (ctl *Controller) AddUserMessage(username string, message *pmapi.Message) (
 }
 
 func (ctl *Controller) GetMessages(username, labelID string) ([]*pmapi.Message, error) {
-	client, ok := ctl.pmapiByUsername[username]
-	if !ok {
-		return nil, fmt.Errorf("user %s does not exist", username)
+	client, err := getPersistentClient(username)
+	if err != nil {
+		return nil, err
 	}
 
 	page := 0

test/liveapi/persistent_clients.go

@@ -0,0 +1,131 @@
+// Copyright (c) 2021 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package liveapi
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"os"
+
+	"github.com/ProtonMail/proton-bridge/internal/constants"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/ProtonMail/proton-bridge/pkg/srp"
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+)
+
+// persistentClients keeps authenticated clients for tests.
+//
+// We need to reduce the number of authentication done by live tests.
+// Before every *scenario* we are creating and authenticating new client.
+// This is not necessary for controller purposes. We can reuse the same clients
+// for all tests.
+//
+//nolint[gochecknoglobals]
+var persistentClients = struct {
+	manager    pmapi.Manager
+	byName     map[string]pmapi.Client
+	saltByName map[string]string
+}{}
+
+type persistentClient struct {
+	pmapi.Client
+	username string
+}
+
+// AuthDelete is noop. All sessions will be closed in CleanupPersistentClients.
+func (pc *persistentClient) AuthDelete(_ context.Context) error {
+	return nil
+}
+
+// AuthSalt returns cached string. Otherwise after some time there is an error:
+//
+//     Access token does not have sufficient scope
+//
+// while all other routes works normally. Need to confirm with Aron that this
+// is expected behaviour.
+func (pc *persistentClient) AuthSalt(_ context.Context) (string, error) {
+	return persistentClients.saltByName[pc.username], nil
+}
+
+func SetupPersistentClients() {
+	app := os.Getenv("TEST_APP")
+
+	persistentClients.manager = pmapi.New(pmapi.NewConfig(getAppVersionName(app), constants.Version))
+	persistentClients.manager.SetLogging(logrus.WithField("pkg", "liveapi"), logrus.GetLevel() == logrus.TraceLevel)
+
+	persistentClients.byName = map[string]pmapi.Client{}
+	persistentClients.saltByName = map[string]string{}
+}
+
+func getAppVersionName(app string) string {
+	if app == "ie" {
+		return "importExport"
+	}
+	return app
+}
+
+func CleanupPersistentClients() {
+	for username, client := range persistentClients.byName {
+		if err := client.AuthDelete(context.Background()); err != nil {
+			logrus.WithError(err).
+				WithField("username", username).
+				Error("Failed to logout persistent client")
+		}
+	}
+}
+
+func addPersistentClient(username string, password, mailboxPassword []byte) (pmapi.Client, error) {
+	if cl, ok := persistentClients.byName[username]; ok {
+		return cl, nil
+	}
+
+	srp.RandReader = rand.New(rand.NewSource(42)) //nolint[gosec] It is OK to use weaker random number generator here
+
+	client, _, err := persistentClients.manager.NewClientWithLogin(context.Background(), username, password)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to create new persistent client")
+	}
+
+	salt, err := client.AuthSalt(context.Background())
+	if err != nil {
+		return nil, errors.Wrap(err, "persistent client: failed to get salt")
+	}
+
+	hashedMboxPass, err := pmapi.HashMailboxPassword(mailboxPassword, salt)
+	if err != nil {
+		return nil, errors.Wrap(err, "persistent client: failed to hash mailbox password")
+	}
+
+	if err := client.Unlock(context.Background(), hashedMboxPass); err != nil {
+		return nil, errors.Wrap(err, "persistent client: failed to unlock user")
+	}
+
+	persistentClients.byName[username] = client
+	persistentClients.saltByName[username] = salt
+	return client, nil
+}
+
+func getPersistentClient(username string) (pmapi.Client, error) {
+	v, ok := persistentClients.byName[username]
+	if !ok {
+		return nil, fmt.Errorf("user %s does not exist", username)
+	}
+	return &persistentClient{v, username}, nil
+}

test/liveapi/users.go

@@ -21,43 +21,42 @@ import (
 	"context"
 
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/ProtonMail/proton-bridge/test/accounts"
 	"github.com/cucumber/godog"
 	"github.com/pkg/errors"
 )
 
-func (ctl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList, password []byte, twoFAEnabled bool) error {
-	if twoFAEnabled {
+func (ctl *Controller) AddUser(account *accounts.TestAccount) error {
+	if account.IsTwoFAEnabled() {
 		return godog.ErrPending
 	}
 
-	client, _, err := ctl.clientManager.NewClientWithLogin(context.Background(), user.Name, password)
+	client, err := addPersistentClient(account.User().Name, account.Password(), account.MailboxPassword())
 	if err != nil {
-		return errors.Wrap(err, "failed to create new client")
+		return errors.Wrap(err, "failed to add persistent client")
 	}
 
-	salt, err := client.AuthSalt(context.Background())
-	if err != nil {
-		return errors.Wrap(err, "failed to get salt")
-	}
-
-	mailboxPassword, err := pmapi.HashMailboxPassword(password, salt)
-	if err != nil {
-		return errors.Wrap(err, "failed to hash mailbox password")
-	}
-
-	if err := client.Unlock(context.Background(), mailboxPassword); err != nil {
-		return errors.Wrap(err, "failed to unlock user")
-	}
-
-	if err := cleanup(client, addresses); err != nil {
+	if err := cleanup(client, account.Addresses()); err != nil {
 		return errors.Wrap(err, "failed to clean user")
 	}
 
-	ctl.pmapiByUsername[user.Name] = client
-
 	return nil
 }
 
 func (ctl *Controller) ReorderAddresses(user *pmapi.User, addressIDs []string) error {
-	return ctl.pmapiByUsername[user.Name].ReorderAddresses(context.Background(), addressIDs)
+	client, err := getPersistentClient(user.Name)
+	if err != nil {
+		return err
+	}
+	return client.ReorderAddresses(context.Background(), addressIDs)
+}
+
+func (ctl *Controller) GetAuthClient(username string) pmapi.Client {
+	client, err := getPersistentClient(username)
+	if err != nil {
+		ctl.log.WithError(err).
+			WithField("username", username).
+			Fatal("Cannot get authenticated client")
+	}
+	return client
 }

test/users_setup_test.go

@@ -28,6 +28,7 @@ import (
 func UsersSetupFeatureContext(s *godog.Suite) {
 	s.Step(`^there is user "([^"]*)"$`, thereIsUser)
 	s.Step(`^there is connected user "([^"]*)"$`, thereIsConnectedUser)
+	s.Step(`^there is user "([^"]*)" which just logged in$`, thereIsUserWhichJustLoggedIn)
 	s.Step(`^there is disconnected user "([^"]*)"$`, thereIsDisconnectedUser)
 	s.Step(`^there is database file for "([^"]*)"$`, thereIsDatabaseFileForUser)
 	s.Step(`^there is no database file for "([^"]*)"$`, thereIsNoDatabaseFileForUser)
@@ -39,7 +40,7 @@ func thereIsUser(bddUserID string) error {
 	if account == nil {
 		return godog.ErrPending
 	}
-	err := ctx.GetPMAPIController().AddUser(account.User(), account.Addresses(), account.Password(), account.IsTwoFAEnabled())
+	err := ctx.GetPMAPIController().AddUser(account)
 	return internalError(err, "adding user %s", account.Username())
 }
 
@@ -48,7 +49,21 @@ func thereIsConnectedUser(bddUserID string) error {
 	if account == nil {
 		return godog.ErrPending
 	}
-	err := ctx.GetPMAPIController().AddUser(account.User(), account.Addresses(), account.Password(), account.IsTwoFAEnabled())
+	username := account.Username()
+	ctl := ctx.GetPMAPIController()
+	err := ctl.AddUser(account)
+	if err != nil {
+		return internalError(err, "adding user %s", username)
+	}
+	return ctx.FinishLogin(ctx.GetPMAPIController().GetAuthClient(username), account.MailboxPassword())
+}
+
+func thereIsUserWhichJustLoggedIn(bddUserID string) error {
+	account := ctx.GetTestAccount(bddUserID)
+	if account == nil {
+		return godog.ErrPending
+	}
+	err := ctx.GetPMAPIController().AddUser(account)
 	if err != nil {
 		return internalError(err, "adding user %s", account.Username())
 	}
@@ -60,7 +75,7 @@ func thereIsDisconnectedUser(bddUserID string) error {
 	if account == nil {
 		return godog.ErrPending
 	}
-	err := ctx.GetPMAPIController().AddUser(account.User(), account.Addresses(), account.Password(), account.IsTwoFAEnabled())
+	err := ctx.GetPMAPIController().AddUser(account)
 	if err != nil {
 		return internalError(err, "adding user %s", account.Username())
 	}