1 files changed, 69 insertions, 33 deletions

diff --git a/api4/apitestlib.go b/api4/apitestlib.go
index f50095f79..640f38fe6 100644
--- a/api4/apitestlib.go
+++ b/api4/apitestlib.go
@@ -22,6 +22,8 @@ import (
 	"github.com/mattermost/mattermost-server/app"
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/store"
+	"github.com/mattermost/mattermost-server/store/sqlstore"
+	"github.com/mattermost/mattermost-server/store/storetest"
 	"github.com/mattermost/mattermost-server/utils"
 	"github.com/mattermost/mattermost-server/wsapi"
 
@@ -46,39 +48,66 @@ type TestHelper struct {
 	SystemAdminUser   *model.User
 }
 
-func setupTestHelper(enterprise bool) *TestHelper {
-	th := &TestHelper{
-		App: app.New(),
+type persistentTestStore struct {
+	store.Store
+}
+
+func (*persistentTestStore) Close() {}
+
+var testStoreContainer *storetest.RunningContainer
+var testStore *persistentTestStore
+
+// UseTestStore sets the container and corresponding settings to use for tests. Once the tests are
+// complete (e.g. at the end of your TestMain implementation), you should call StopTestStore.
+func UseTestStore(container *storetest.RunningContainer, settings *model.SqlSettings) {
+	testStoreContainer = container
+	testStore = &persistentTestStore{store.NewLayeredStore(sqlstore.NewSqlSupplier(*settings, nil), nil, nil)}
+}
+
+func StopTestStore() {
+	if testStoreContainer != nil {
+		testStoreContainer.Stop()
+		testStoreContainer = nil
 	}
+}
 
-	if th.App.Srv == nil {
+func setupTestHelper(enterprise bool) *TestHelper {
+	if utils.T == nil {
 		utils.TranslationsPreInit()
-		utils.LoadConfig("config.json")
-		utils.InitTranslations(utils.Cfg.LocalizationSettings)
-		*utils.Cfg.TeamSettings.MaxUsersPerTeam = 50
-		*utils.Cfg.RateLimitSettings.Enable = false
-		utils.Cfg.EmailSettings.SendEmailNotifications = true
-		utils.DisableDebugLogForTest()
-		th.App.NewServer()
-		th.App.InitStores()
-		th.App.Srv.Router = NewRouter()
-		th.App.Srv.WebSocketRouter = th.App.NewWebSocketRouter()
-		th.App.StartServer()
-		Init(th.App, th.App.Srv.Router, true)
-		wsapi.Init(th.App, th.App.Srv.WebSocketRouter)
-		utils.EnableDebugLogForTest()
-		th.App.Srv.Store.MarkSystemRanUnitTests()
+	}
+	utils.LoadConfig("config.json")
+	utils.InitTranslations(utils.Cfg.LocalizationSettings)
 
-		*utils.Cfg.TeamSettings.EnableOpenServer = true
+	var options []app.Option
+	if testStore != nil {
+		options = append(options, app.StoreOverride(testStore))
+		options = append(options, app.ConfigOverride(func(cfg *model.Config) {
+			cfg.ServiceSettings.ListenAddress = new(string)
+			*cfg.ServiceSettings.ListenAddress = ":0"
+		}))
+	}
+
+	th := &TestHelper{
+		App: app.New(options...),
 	}
 
+	*utils.Cfg.TeamSettings.MaxUsersPerTeam = 50
+	*utils.Cfg.RateLimitSettings.Enable = false
+	utils.Cfg.EmailSettings.SendEmailNotifications = true
+	utils.DisableDebugLogForTest()
+	th.App.StartServer()
+	Init(th.App, th.App.Srv.Router, true)
+	wsapi.Init(th.App, th.App.Srv.WebSocketRouter)
+	utils.EnableDebugLogForTest()
+	th.App.Srv.Store.MarkSystemRanUnitTests()
+
+	*utils.Cfg.TeamSettings.EnableOpenServer = true
+
 	utils.SetIsLicensed(enterprise)
 	if enterprise {
 		utils.License().Features.SetDefaults()
 	}
 
-	th.App.Jobs.Store = th.App.Srv.Store
-
 	th.Client = th.CreateClient()
 	th.SystemAdminClient = th.CreateClient()
 	return th
@@ -150,23 +179,18 @@ func (me *TestHelper) TearDown() {
 	me.App.Shutdown()
 
 	utils.EnableDebugLogForTest()
-}
 
-func (me *TestHelper) waitForConnectivity() {
-	for i := 0; i < 1000; i++ {
-		_, err := net.Dial("tcp", "localhost"+*utils.Cfg.ServiceSettings.ListenAddress)
-		if err == nil {
-			return
-		}
-		time.Sleep(time.Millisecond * 20)
+	if err := recover(); err != nil {
+		StopTestStore()
+		panic(err)
 	}
-	panic("unable to connect")
 }
 
 func (me *TestHelper) InitBasic() *TestHelper {
 	me.waitForConnectivity()
 
 	me.TeamAdminUser = me.CreateUser()
+	me.App.UpdateUserRoles(me.TeamAdminUser.Id, model.ROLE_SYSTEM_USER.Id)
 	me.LoginTeamAdmin()
 	me.BasicTeam = me.CreateTeam()
 	me.BasicChannel = me.CreatePublicChannel()
@@ -199,12 +223,24 @@ func (me *TestHelper) InitSystemAdmin() *TestHelper {
 	return me
 }
 
+func (me *TestHelper) waitForConnectivity() {
+	for i := 0; i < 1000; i++ {
+		conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%v", me.App.Srv.ListenAddr.Port))
+		if err == nil {
+			conn.Close()
+			return
+		}
+		time.Sleep(time.Millisecond * 20)
+	}
+	panic("unable to connect")
+}
+
 func (me *TestHelper) CreateClient() *model.Client4 {
-	return model.NewAPIv4Client("http://localhost" + *utils.Cfg.ServiceSettings.ListenAddress)
+	return model.NewAPIv4Client(fmt.Sprintf("http://localhost:%v", me.App.Srv.ListenAddr.Port))
 }
 
 func (me *TestHelper) CreateWebSocketClient() (*model.WebSocketClient, *model.AppError) {
-	return model.NewWebSocketClient4("ws://localhost"+*utils.Cfg.ServiceSettings.ListenAddress, me.Client.AuthToken)
+	return model.NewWebSocketClient4(fmt.Sprintf("ws://localhost:%v", me.App.Srv.ListenAddr.Port), me.Client.AuthToken)
 }
 
 func (me *TestHelper) CreateUser() *model.User {