feat: make store use ClientManager

test/context/bridge.go

@@ -21,9 +21,9 @@ import (
 	"os"
 	"runtime"
 
+	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/internal/preferences"
 	"github.com/ProtonMail/proton-bridge/pkg/listener"
-	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 )
 
 // GetBridge returns bridge instance.
@@ -34,10 +34,7 @@ func (ctx *TestContext) GetBridge() *bridge.Bridge {
 // withBridgeInstance creates a bridge instance for use in the test.
 // Every TestContext has this by default and thus this doesn't need to be exported.
 func (ctx *TestContext) withBridgeInstance() {
-	pmapiFactory := func(userID string) pmapi.Client {
-		return ctx.pmapiController.GetClient(userID)
-	}
-	ctx.bridge = newBridgeInstance(ctx.t, ctx.cfg, ctx.credStore, ctx.listener, pmapiFactory)
+	ctx.bridge = newBridgeInstance(ctx.t, ctx.cfg, ctx.credStore, ctx.listener, ctx.clientManager)
 	ctx.addCleanupChecked(ctx.bridge.ClearData, "Cleaning bridge data")
 }
 
@@ -62,7 +59,7 @@ func newBridgeInstance(
 	cfg *fakeConfig,
 	credStore bridge.CredentialsStorer,
 	eventListener listener.Listener,
-	pmapiFactory bridge.PMAPIProviderFactory,
+	clientManager bridge.ClientManager,
 ) *bridge.Bridge {
 	version := os.Getenv("VERSION")
 	bridge.UpdateCurrentUserAgent(version, runtime.GOOS, "", "")
@@ -70,7 +67,7 @@ func newBridgeInstance(
 	panicHandler := &panicHandler{t: t}
 	pref := preferences.New(cfg)
 
-	return bridge.New(cfg, pref, panicHandler, eventListener, version, pmapiFactory, credStore)
+	return bridge.New(cfg, pref, panicHandler, eventListener, version, clientManager, credStore)
 }
 
 // SetLastBridgeError sets the last error that occurred while executing a bridge action.

test/context/config.go

@@ -28,7 +28,6 @@ import (
 
 type fakeConfig struct {
 	dir string
-	tm  *pmapi.TokenManager
 }
 
 // newFakeConfig creates a temporary folder for files.
@@ -41,7 +40,6 @@ func newFakeConfig() *fakeConfig {
 
 	return &fakeConfig{
 		dir: dir,
-		tm:  pmapi.NewTokenManager(),
 	}
 }
 
@@ -53,8 +51,6 @@ func (c *fakeConfig) GetAPIConfig() *pmapi.ClientConfig {
 		AppVersion: "Bridge_" + os.Getenv("VERSION"),
 		ClientID:   "bridge",
 		SentryDSN:  "",
-		// TokenManager should not be required, but PMAPI still doesn't handle not-set cases everywhere.
-		TokenManager: c.tm,
 	}
 }
 func (c *fakeConfig) GetDBDir() string {

test/context/context.go

@@ -21,6 +21,7 @@ package context
 import (
 	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"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/mocks"
 	"github.com/sirupsen/logrus"
@@ -44,6 +45,7 @@ type TestContext struct {
 	bridge          *bridge.Bridge
 	bridgeLastError error
 	credStore       bridge.CredentialsStorer
+	clientManager   *pmapi.ClientManager
 
 	// IMAP related variables.
 	imapAddr          string
@@ -70,11 +72,14 @@ func New() *TestContext {
 
 	cfg := newFakeConfig()
 
+	cm := pmapi.NewClientManager(cfg.GetAPIConfig())
+
 	ctx := &TestContext{
 		t:                 &bddT{},
 		cfg:               cfg,
 		listener:          listener.New(),
-		pmapiController:   newPMAPIController(),
+		pmapiController:   newPMAPIController(cm),
+		clientManager:     cm,
 		testAccounts:      newTestAccounts(),
 		credStore:         newFakeCredStore(),
 		imapClients:       make(map[string]*mocks.IMAPClient),

test/context/pmapi_controller.go

@@ -20,14 +20,12 @@ package context
 import (
 	"os"
 
-	"github.com/ProtonMail/proton-bridge/internal/bridge"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	"github.com/ProtonMail/proton-bridge/test/fakeapi"
 	"github.com/ProtonMail/proton-bridge/test/liveapi"
 )
 
 type PMAPIController interface {
-	GetClient(userID string) bridge.PMAPIProvider
 	TurnInternetConnectionOff()
 	TurnInternetConnectionOn()
 	AddUser(user *pmapi.User, addresses *pmapi.AddressList, password string, twoFAEnabled bool) error
@@ -40,19 +38,19 @@ type PMAPIController interface {
 	GetCalls(method, path string) [][]byte
 }
 
-func newPMAPIController() PMAPIController {
+func newPMAPIController(cm *pmapi.ClientManager) PMAPIController {
 	switch os.Getenv(EnvName) {
 	case EnvFake:
-		return newFakePMAPIController()
+		return newFakePMAPIController(cm)
 	case EnvLive:
-		return newLivePMAPIController()
+		return newLivePMAPIController(cm)
 	default:
 		panic("unknown env")
 	}
 }
 
-func newFakePMAPIController() PMAPIController {
-	return newFakePMAPIControllerWrap(fakeapi.NewController())
+func newFakePMAPIController(cm *pmapi.ClientManager) PMAPIController {
+	return newFakePMAPIControllerWrap(fakeapi.NewController(cm))
 }
 
 type fakePMAPIControllerWrap struct {
@@ -63,12 +61,8 @@ func newFakePMAPIControllerWrap(controller *fakeapi.Controller) PMAPIController
 	return &fakePMAPIControllerWrap{Controller: controller}
 }
 
-func (s *fakePMAPIControllerWrap) GetClient(userID string) bridge.PMAPIProvider {
-	return s.Controller.GetClient(userID)
-}
-
-func newLivePMAPIController() PMAPIController {
-	return newLiveAPIControllerWrap(liveapi.NewController())
+func newLivePMAPIController(cm *pmapi.ClientManager) PMAPIController {
+	return newLiveAPIControllerWrap(liveapi.NewController(cm))
 }
 
 type liveAPIControllerWrap struct {
@@ -78,7 +72,3 @@ type liveAPIControllerWrap struct {
 func newLiveAPIControllerWrap(controller *liveapi.Controller) PMAPIController {
 	return &liveAPIControllerWrap{Controller: controller}
 }
-
-func (s *liveAPIControllerWrap) GetClient(userID string) bridge.PMAPIProvider {
-	return s.Controller.GetClient(userID)
-}

test/fakeapi/attachments.go

@@ -45,3 +45,11 @@ func (api *FakePMAPI) CreateAttachment(attachment *pmapi.Attachment, data io.Rea
 	attachment.KeyPackets = base64.StdEncoding.EncodeToString(bytes)
 	return attachment, nil
 }
+
+func (api *FakePMAPI) DeleteAttachment(attachmentID string) error {
+	if err := api.checkAndRecordCall(GET, "/attachments/"+attachmentID, nil); err != nil {
+		return err
+	}
+
+	return nil
+}

test/fakeapi/auth.go

@@ -141,13 +141,23 @@ func (api *FakePMAPI) AuthRefresh(token string) (*pmapi.Auth, error) {
 	return auth, nil
 }
 
-func (api *FakePMAPI) Logout() error {
+func (api *FakePMAPI) Logout() {
+	_ = api.DeleteAuth()
+	api.ClearData()
+}
+
+func (api *FakePMAPI) DeleteAuth() error {
 	if err := api.checkAndRecordCall(DELETE, "/auth", nil); err != nil {
 		return err
 	}
+
 	// Logout will also emit change to auth channel
 	api.sendAuth(nil)
-	api.controller.deleteSession(api.uid)
-	api.unsetUser()
+
 	return nil
 }
+
+func (api *FakePMAPI) ClearData() {
+	api.controller.deleteSession(api.uid)
+	api.unsetUser()
+}

test/fakeapi/controller.go

@@ -44,8 +44,8 @@ type Controller struct {
 	log *logrus.Entry
 }
 
-func NewController() *Controller {
-	return &Controller{
+func NewController(cm *pmapi.ClientManager) (cntrl *Controller) {
+	cntrl = &Controller{
 		lock:               &sync.RWMutex{},
 		fakeAPIs:           []*FakePMAPI{},
 		calls:              []*fakeCall{},
@@ -62,10 +62,12 @@ func NewController() *Controller {
 
 		log: logrus.WithField("pkg", "fakeapi-controller"),
 	}
-}
 
-func (cntrl *Controller) GetClient(userID string) *FakePMAPI {
-	fakeAPI := New(cntrl)
-	cntrl.fakeAPIs = append(cntrl.fakeAPIs, fakeAPI)
-	return fakeAPI
+	cm.SetClientConstructor(func(userID string) pmapi.Client {
+		fakeAPI := New(cntrl)
+		cntrl.fakeAPIs = append(cntrl.fakeAPIs, fakeAPI)
+		return fakeAPI
+	})
+
+	return
 }

test/fakeapi/reports.go

@@ -42,3 +42,7 @@ func (api *FakePMAPI) SendSimpleMetric(category, action, label string) error {
 	v.Set("Label", label)
 	return api.checkInternetAndRecordCall(GET, "/metrics?"+v.Encode(), nil)
 }
+
+func (api *FakePMAPI) ReportSentryCrash(reportErr error) (err error) {
+	return nil
+}

test/fakeapi/user.go

@@ -50,6 +50,10 @@ func (api *FakePMAPI) UpdateUser() (*pmapi.User, error) {
 	return api.user, nil
 }
 
+func (api *FakePMAPI) GetAddresses() (pmapi.AddressList, error) {
+	return *api.addresses, nil
+}
+
 func (api *FakePMAPI) Addresses() pmapi.AddressList {
 	return *api.addresses
 }

test/liveapi/cleanup.go

@@ -24,7 +24,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-func cleanup(client *pmapi.Client) error {
+func cleanup(client pmapi.Client) error {
 	if err := cleanSystemFolders(client); err != nil {
 		return errors.Wrap(err, "failed to clean system folders")
 	}
@@ -37,7 +37,7 @@ func cleanup(client *pmapi.Client) error {
 	return nil
 }
 
-func cleanSystemFolders(client *pmapi.Client) error {
+func cleanSystemFolders(client pmapi.Client) error {
 	for _, labelID := range []string{pmapi.InboxLabel, pmapi.SentLabel, pmapi.ArchiveLabel, pmapi.AllMailLabel, pmapi.DraftLabel} {
 		for {
 			messages, total, err := client.ListMessages(&pmapi.MessagesFilter{
@@ -69,7 +69,7 @@ func cleanSystemFolders(client *pmapi.Client) error {
 	return nil
 }
 
-func cleanCustomLables(client *pmapi.Client) error {
+func cleanCustomLables(client pmapi.Client) error {
 	labels, err := client.ListLabels()
 	if err != nil {
 		return errors.Wrap(err, "failed to list labels")
@@ -87,7 +87,7 @@ func cleanCustomLables(client *pmapi.Client) error {
 	return nil
 }
 
-func cleanTrash(client *pmapi.Client) error {
+func cleanTrash(client pmapi.Client) error {
 	for {
 		_, total, err := client.ListMessages(&pmapi.MessagesFilter{
 			PageSize: 1,
@@ -110,7 +110,7 @@ func cleanTrash(client *pmapi.Client) error {
 	return nil
 }
 
-func emptyFolder(client *pmapi.Client, labelID string) error {
+func emptyFolder(client pmapi.Client, labelID string) error {
 	err := client.EmptyFolder(labelID, "")
 	if err != nil {
 		return err

test/liveapi/controller.go

@@ -18,9 +18,7 @@
 package liveapi
 
 import (
-	"fmt"
 	"net/http"
-	"os"
 	"sync"
 
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
@@ -30,33 +28,27 @@ type Controller struct {
 	// Internal states.
 	lock                 *sync.RWMutex
 	calls                []*fakeCall
-	pmapiByUsername      map[string]*pmapi.Client
 	messageIDsByUsername map[string][]string
+	clientManager        *pmapi.ClientManager
 
 	// State controlled by test.
 	noInternetConnection bool
 }
 
-func NewController() *Controller {
-	return &Controller{
+func NewController(cm *pmapi.ClientManager) (cntrl *Controller) {
+	cntrl = &Controller{
 		lock:                 &sync.RWMutex{},
 		calls:                []*fakeCall{},
-		pmapiByUsername:      map[string]*pmapi.Client{},
 		messageIDsByUsername: map[string][]string{},
+		clientManager:        cm,
 
 		noInternetConnection: false,
 	}
-}
 
-func (cntrl *Controller) GetClient(userID string) *pmapi.Client {
-	cfg := &pmapi.ClientConfig{
-		AppVersion: fmt.Sprintf("Bridge_%s", os.Getenv("VERSION")),
-		ClientID:   "bridge-test",
-		Transport: &fakeTransport{
-			cntrl:     cntrl,
-			transport: http.DefaultTransport,
-		},
-		TokenManager: pmapi.NewTokenManager(),
-	}
-	return pmapi.NewClient(cfg, userID)
+	cm.SetRoundTripper(&fakeTransport{
+		cntrl:     cntrl,
+		transport: http.DefaultTransport,
+	})
+
+	return
 }

test/liveapi/labels.go

@@ -36,11 +36,7 @@ var systemLabelNameToID = map[string]string{ //nolint[gochecknoglobals]
 }
 
 func (cntrl *Controller) AddUserLabel(username string, label *pmapi.Label) error {
-	client, ok := cntrl.pmapiByUsername[username]
-	if !ok {
-		return fmt.Errorf("user %s does not exist", username)
-	}
-
+	client := cntrl.clientManager.GetClient(username)
 	label.Exclusive = getLabelExclusive(label.Name)
 	label.Name = getLabelNameWithoutPrefix(label.Name)
 	label.Color = pmapi.LabelColors[0]
@@ -67,11 +63,7 @@ func (cntrl *Controller) getLabelID(username, labelName string) (string, error)
 		return labelID, nil
 	}
 
-	client, ok := cntrl.pmapiByUsername[username]
-	if !ok {
-		return "", fmt.Errorf("user %s does not exist", username)
-	}
-
+	client := cntrl.clientManager.GetClient(username)
 	labels, err := client.ListLabels()
 	if err != nil {
 		return "", errors.Wrap(err, "failed to list labels")

test/liveapi/messages.go

@@ -31,10 +31,7 @@ import (
 )
 
 func (cntrl *Controller) AddUserMessage(username string, message *pmapi.Message) error {
-	client, ok := cntrl.pmapiByUsername[username]
-	if !ok {
-		return fmt.Errorf("user %s does not exist", username)
-	}
+	client := cntrl.clientManager.GetClient(username)
 
 	body, err := buildMessage(client, message)
 	if err != nil {
@@ -64,7 +61,7 @@ func (cntrl *Controller) AddUserMessage(username string, message *pmapi.Message)
 	return nil
 }
 
-func buildMessage(client *pmapi.Client, message *pmapi.Message) (*bytes.Buffer, error) {
+func buildMessage(client pmapi.Client, message *pmapi.Message) (*bytes.Buffer, error) {
 	if err := encryptMessage(client, message); err != nil {
 		return nil, errors.Wrap(err, "failed to encrypt message")
 	}
@@ -79,7 +76,7 @@ func buildMessage(client *pmapi.Client, message *pmapi.Message) (*bytes.Buffer,
 	return body, nil
 }
 
-func encryptMessage(client *pmapi.Client, message *pmapi.Message) error {
+func encryptMessage(client pmapi.Client, message *pmapi.Message) error {
 	addresses, err := client.GetAddresses()
 	if err != nil {
 		return errors.Wrap(err, "failed to get address")

test/liveapi/users.go

@@ -18,9 +18,7 @@
 package liveapi
 
 import (
-	"fmt"
-	"os"
-
+	"github.com/ProtonMail/bridge/pkg/pmapi"
 	"github.com/cucumber/godog"
 	"github.com/pkg/errors"
 )
@@ -30,11 +28,7 @@ func (cntrl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList,
 		return godog.ErrPending
 	}
 
-	client := pmapi.NewClient(&pmapi.ClientConfig{
-		AppVersion:   fmt.Sprintf("Bridge_%s", os.Getenv("VERSION")),
-		ClientID:     "bridge-cntrl",
-		TokenManager: pmapi.NewTokenManager(),
-	}, user.ID)
+	client := cntrl.clientManager.GetClient(user.ID)
 
 	authInfo, err := client.AuthInfo(user.Name)
 	if err != nil {
@@ -60,6 +54,5 @@ func (cntrl *Controller) AddUser(user *pmapi.User, addresses *pmapi.AddressList,
 		return errors.Wrap(err, "failed to clean user")
 	}
 
-	cntrl.pmapiByUsername[user.Name] = client
 	return nil
 }