[MM-17297] Extend EnsureBot helper to include bot images (#12153)

plugin/api.go

@@ -664,7 +664,7 @@ type API interface {
 	// Minimum server version: 5.6
 	GetFileLink(fileId string) (string, *model.AppError)
 
-	// ReadFileAtPath reads the file from the backend for a specific path
+	// ReadFile reads the file from the backend for a specific path
 	//
 	// @tag File
 	// Minimum server version: 5.3

plugin/helpers.go

@@ -14,10 +14,11 @@ import (
 // Plugins obtain access to the Helpers by embedding MattermostPlugin.
 type Helpers interface {
 	// EnsureBot either returns an existing bot user matching the given bot, or creates a bot user from the given bot.
+	// A profile image or icon image may be optionally passed in to be set for the existing or newly created bot.
 	// Returns the id of the resulting bot.
 	//
 	// Minimum server version: 5.10
-	EnsureBot(bot *model.Bot) (string, error)
+	EnsureBot(bot *model.Bot, options ...EnsureBotOption) (string, error)
 
 	// KVSetJSON stores a key-value pair, unique per plugin, marshalling the given value as a JSON string.
 	//

plugin/helpers_bots.go

@@ -4,19 +4,77 @@
 package plugin
 
 import (
+	"io/ioutil"
+	"path/filepath"
+
 	"github.com/pkg/errors"
 
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/utils"
 )
 
+type ensureBotOptions struct {
+	ProfileImagePath string
+	IconImagePath    string
+}
+
+type EnsureBotOption func(*ensureBotOptions)
+
+func ProfileImagePath(path string) EnsureBotOption {
+	return func(args *ensureBotOptions) {
+		args.ProfileImagePath = path
+	}
+}
+
+func IconImagePath(path string) EnsureBotOption {
+	return func(args *ensureBotOptions) {
+		args.IconImagePath = path
+	}
+}
+
+func (p *HelpersImpl) readFile(path string) ([]byte, error) {
+	bundlePath, err := p.API.GetBundlePath()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get bundle path")
+	}
+
+	imageBytes, err := ioutil.ReadFile(filepath.Join(bundlePath, path))
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read image")
+	}
+	return imageBytes, nil
+}
+
 // EnsureBot implements Helpers.EnsureBot
-func (p *HelpersImpl) EnsureBot(bot *model.Bot) (retBotID string, retErr error) {
+func (p *HelpersImpl) EnsureBot(bot *model.Bot, options ...EnsureBotOption) (retBotID string, retErr error) {
 	err := p.ensureServerVersion("5.10.0")
 	if err != nil {
 		return "", errors.Wrap(err, "failed to ensure bot")
 	}
 
+	// Default options
+	o := &ensureBotOptions{
+		ProfileImagePath: "",
+		IconImagePath:    "",
+	}
+
+	for _, setter := range options {
+		setter(o)
+	}
+
+	botID, err := p.ensureBot(bot)
+	if err != nil {
+		return "", err
+	}
+
+	err = p.setBotImages(botID, o.ProfileImagePath, o.IconImagePath)
+	if err != nil {
+		return "", err
+	}
+	return botID, nil
+}
+
+func (p *HelpersImpl) ensureBot(bot *model.Bot) (retBotID string, retErr error) {
 	// Must provide a bot with a username
 	if bot == nil || len(bot.Username) < 1 {
 		return "", errors.New("passed a bad bot, nil or no username")
@@ -79,3 +137,27 @@ func (p *HelpersImpl) EnsureBot(bot *model.Bot) (retBotID string, retErr error)
 
 	return createdBot.UserId, nil
 }
+
+func (p *HelpersImpl) setBotImages(botID, profileImagePath, iconImagePath string) error {
+	if profileImagePath != "" {
+		imageBytes, err := p.readFile(profileImagePath)
+		if err != nil {
+			return errors.Wrap(err, "failed to read profile image")
+		}
+		appErr := p.API.SetProfileImage(botID, imageBytes)
+		if appErr != nil {
+			return errors.Wrap(appErr, "failed to set profile image")
+		}
+	}
+	if iconImagePath != "" {
+		imageBytes, err := p.readFile(iconImagePath)
+		if err != nil {
+			return errors.Wrap(err, "failed to read icon image")
+		}
+		appErr := p.API.SetBotIconImage(botID, imageBytes)
+		if appErr != nil {
+			return errors.Wrap(appErr, "failed to set icon image")
+		}
+	}
+	return nil
+}

plugin/helpers_bots_test.go

@@ -4,12 +4,15 @@
 package plugin_test
 
 import (
+	"io/ioutil"
+	"path/filepath"
 	"testing"
 
 	"github.com/mattermost/mattermost-server/model"
 	"github.com/mattermost/mattermost-server/plugin"
 	"github.com/mattermost/mattermost-server/plugin/plugintest"
 	"github.com/mattermost/mattermost-server/plugin/plugintest/mock"
+	"github.com/mattermost/mattermost-server/utils/fileutils"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -95,6 +98,77 @@ func TestEnsureBot(t *testing.T) {
 			assert.Equal(t, "", botId)
 			assert.NotNil(t, err)
 		})
+
+		t.Run("should set the bot profile image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return([]byte(expectedBotId), nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetProfileImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			assert.Nil(t, err)
+
+			botId, err := p.EnsureBot(testbot, plugin.ProfileImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
+
+		t.Run("should set the bot icon image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+			assert.Nil(t, err)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return([]byte(expectedBotId), nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetBotIconImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			botId, err := p.EnsureBot(testbot, plugin.IconImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
+
+		t.Run("should set both the profile image and bot icon image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+			assert.Nil(t, err)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return([]byte(expectedBotId), nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetProfileImage", expectedBotId, imageBytes).Return(nil)
+			api.On("SetBotIconImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			botId, err := p.EnsureBot(testbot, plugin.ProfileImagePath(testImage), plugin.IconImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
 	})
 
 	t.Run("if bot doesn't exist", func(t *testing.T) {
@@ -179,5 +253,90 @@ func TestEnsureBot(t *testing.T) {
 			assert.Equal(t, "", botId)
 			assert.NotNil(t, err)
 		})
+
+		t.Run("should create bot and set the bot profile image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+			assert.Nil(t, err)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return(nil, nil)
+			api.On("GetUserByUsername", testbot.Username).Return(nil, nil)
+			api.On("CreateBot", testbot).Return(&model.Bot{
+				UserId: expectedBotId,
+			}, nil)
+			api.On("KVSet", plugin.BOT_USER_KEY, []byte(expectedBotId)).Return(nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetProfileImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			botId, err := p.EnsureBot(testbot, plugin.ProfileImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
+
+		t.Run("should create bot and set the bot icon image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+			assert.Nil(t, err)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return(nil, nil)
+			api.On("GetUserByUsername", testbot.Username).Return(nil, nil)
+			api.On("CreateBot", testbot).Return(&model.Bot{
+				UserId: expectedBotId,
+			}, nil)
+			api.On("KVSet", plugin.BOT_USER_KEY, []byte(expectedBotId)).Return(nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetBotIconImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			botId, err := p.EnsureBot(testbot, plugin.IconImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
+
+		t.Run("should create bot and set both the profile image and bot icon image when specified", func(t *testing.T) {
+			expectedBotId := model.NewId()
+			api := setupAPI()
+
+			testsDir, _ := fileutils.FindDir("tests")
+			testImage := filepath.Join(testsDir, "test.png")
+			imageBytes, err := ioutil.ReadFile(testImage)
+			assert.Nil(t, err)
+
+			api.On("KVGet", plugin.BOT_USER_KEY).Return(nil, nil)
+			api.On("GetUserByUsername", testbot.Username).Return(nil, nil)
+			api.On("CreateBot", testbot).Return(&model.Bot{
+				UserId: expectedBotId,
+			}, nil)
+			api.On("KVSet", plugin.BOT_USER_KEY, []byte(expectedBotId)).Return(nil)
+			api.On("GetBundlePath").Return("", nil)
+			api.On("SetProfileImage", expectedBotId, imageBytes).Return(nil)
+			api.On("SetBotIconImage", expectedBotId, imageBytes).Return(nil)
+			api.On("GetServerVersion").Return("5.10.0")
+			defer api.AssertExpectations(t)
+
+			p := &plugin.HelpersImpl{}
+			p.API = api
+
+			botId, err := p.EnsureBot(testbot, plugin.ProfileImagePath(testImage), plugin.IconImagePath(testImage))
+			assert.Equal(t, expectedBotId, botId)
+			assert.Nil(t, err)
+		})
 	})
 }

plugin/plugintest/helpers.go

@@ -6,6 +6,7 @@ package plugintest
 
 import (
 	model "github.com/mattermost/mattermost-server/model"
+	plugin "github.com/mattermost/mattermost-server/plugin"
 	mock "github.com/stretchr/testify/mock"
 )
 
@@ -14,20 +15,27 @@ type Helpers struct {
 	mock.Mock
 }
 
-// EnsureBot provides a mock function with given fields: bot
+// EnsureBot provides a mock function with given fields: bot, options
-func (_m *Helpers) EnsureBot(bot *model.Bot) (string, error) {
+func (_m *Helpers) EnsureBot(bot *model.Bot, options ...plugin.EnsureBotOption) (string, error) {
-	ret := _m.Called(bot)
+	_va := make([]interface{}, len(options))
+	for _i := range options {
+		_va[_i] = options[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, bot)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
 
 	var r0 string
-	if rf, ok := ret.Get(0).(func(*model.Bot) string); ok {
+	if rf, ok := ret.Get(0).(func(*model.Bot, ...plugin.EnsureBotOption) string); ok {
-		r0 = rf(bot)
+		r0 = rf(bot, options...)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(*model.Bot) error); ok {
+	if rf, ok := ret.Get(1).(func(*model.Bot, ...plugin.EnsureBotOption) error); ok {
-		r1 = rf(bot)
+		r1 = rf(bot, options...)
 	} else {
 		r1 = ret.Error(1)
 	}