snapshot of https://github.com/mattermost/mattermost.git (master, 2026-09-16; blobs > 131072B omitted)

server/channels/jobs/batch_migration_worker_test.go

@@ -0,0 +1,182 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package jobs_test
+
+import (
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/shared/request"
+	"github.com/mattermost/mattermost/server/v8/channels/jobs"
+	"github.com/mattermost/mattermost/server/v8/channels/store"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type MockApp struct {
+	clusterInfo []*model.ClusterInfo
+}
+
+func (ma MockApp) GetClusterStatus(rctx request.CTX) ([]*model.ClusterInfo, error) {
+	return ma.clusterInfo, nil
+}
+
+func (ma *MockApp) SetInSync() {
+	ma.clusterInfo = nil
+	ma.clusterInfo = append(ma.clusterInfo, &model.ClusterInfo{
+		SchemaVersion: "a",
+	})
+	ma.clusterInfo = append(ma.clusterInfo, &model.ClusterInfo{
+		SchemaVersion: "a",
+	})
+}
+
+func (ma *MockApp) SetOutOfSync() {
+	ma.clusterInfo = nil
+	ma.clusterInfo = append(ma.clusterInfo, &model.ClusterInfo{
+		SchemaVersion: "a",
+	})
+	ma.clusterInfo = append(ma.clusterInfo, &model.ClusterInfo{
+		SchemaVersion: "b",
+	})
+}
+
+func TestBatchMigrationWorker(t *testing.T) {
+	mainHelper.Parallel(t)
+
+	setupBatchWorker := func(t *testing.T, th *TestHelper, mockApp *MockApp, doMigrationBatch func(model.StringMap, store.Store) (model.StringMap, bool, error)) (model.Worker, *model.Job) {
+		t.Helper()
+
+		worker := jobs.MakeBatchMigrationWorker(
+			th.Server.Jobs,
+			th.Server.Store(),
+			mockApp,
+			model.NewId(),
+			1*time.Second,
+			doMigrationBatch,
+		)
+		job := th.SetupBatchWorker(t, worker.BatchWorker)
+
+		return worker, job
+	}
+
+	stopWorker := func(t *testing.T, worker model.Worker) {
+		t.Helper()
+
+		stopped := make(chan bool, 1)
+		go func() {
+			worker.Stop()
+			close(stopped)
+		}()
+
+		waitDone(t, stopped, "worker did not stop")
+	}
+
+	assertJobResetEventually := func(t *testing.T, th *TestHelper, job *model.Job) {
+		t.Helper()
+		assert.EventuallyWithT(t, func(t *assert.CollectT) {
+			actualJob, appErr := th.Server.Jobs.GetJob(th.Context, job.Id)
+			require.Nil(t, appErr)
+			assert.Empty(t, actualJob.Progress, "expected no job progress")
+			assert.Empty(t, actualJob.Data, "expected no job data")
+		}, 5*time.Second, 250*time.Millisecond, "job never reset")
+	}
+
+	getBatchNumberFromData := func(t *testing.T, data model.StringMap) int {
+		t.Helper()
+
+		if data["batch_number"] == "" {
+			data["batch_number"] = "1"
+		}
+		batchNumber, err := strconv.Atoi(data["batch_number"])
+		require.NoError(t, err)
+
+		return batchNumber
+	}
+
+	getDataFromBatchNumber := func(batchNumber int) model.StringMap {
+		data := make(model.StringMap)
+		data["batch_number"] = strconv.Itoa(batchNumber)
+
+		return data
+	}
+
+	t.Run("done after three batches", func(t *testing.T) {
+		th := Setup(t).InitBasic(t)
+
+		mockApp := &MockApp{}
+
+		var worker model.Worker
+		var job *model.Job
+		worker, job = setupBatchWorker(t, th, mockApp, func(data model.StringMap, s store.Store) (model.StringMap, bool, error) {
+			batchNumber := getBatchNumberFromData(t, data)
+			require.LessOrEqual(t, batchNumber, 3, "only 3 batches should have run")
+
+			if batchNumber >= 3 {
+				go worker.Stop() // Shut down the worker when the job is done
+				return getDataFromBatchNumber(batchNumber), true, nil
+			}
+
+			batchNumber++
+			return getDataFromBatchNumber(batchNumber), false, nil
+		})
+
+		// Queue the work to be done
+		worker.JobChannel() <- *job
+
+		th.WaitForJobStatus(t, job, model.JobStatusSuccess)
+		th.WaitForBatchNumber(t, job, 3)
+	})
+
+	t.Run("clusters not in sync before first batch", func(t *testing.T) {
+		th := Setup(t).InitBasic(t)
+
+		mockApp := &MockApp{}
+		mockApp.SetOutOfSync()
+
+		var worker model.Worker
+		var job *model.Job
+		worker, job = setupBatchWorker(t, th, mockApp, func(model.StringMap, store.Store) (model.StringMap, bool, error) {
+			require.Fail(t, "migration batch should never run while clusters not in sync")
+
+			return nil, false, nil
+		})
+
+		// Queue the work to be done
+		worker.JobChannel() <- *job
+
+		th.WaitForJobStatus(t, job, model.JobStatusPending)
+		assertJobResetEventually(t, th, job)
+
+		stopWorker(t, worker)
+	})
+
+	t.Run("clusters not in sync after first batch", func(t *testing.T) {
+		th := Setup(t).InitBasic(t)
+
+		mockApp := &MockApp{}
+
+		var worker model.Worker
+		var job *model.Job
+		worker, job = setupBatchWorker(t, th, mockApp, func(data model.StringMap, s store.Store) (model.StringMap, bool, error) {
+			batchNumber := getBatchNumberFromData(t, data)
+			require.Equal(t, 1, batchNumber, "only batch 1 should have run")
+
+			mockApp.SetOutOfSync()
+			batchNumber++
+
+			return getDataFromBatchNumber(batchNumber), false, nil
+		})
+
+		// Queue the work to be done
+		worker.JobChannel() <- *job
+
+		th.WaitForJobStatus(t, job, model.JobStatusPending)
+		assertJobResetEventually(t, th, job)
+
+		stopWorker(t, worker)
+	})
+}