upstream-mattermost/server/channels/jobs/batch_report_worker_test.go

// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.txt for license information.

package jobs_test

import (
	"errors"
	"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/stretchr/testify/require"
)

type ReportMockApp struct{}

func (rma *ReportMockApp) SaveReportChunk(format string, prefix string, count int, reportData []model.ReportableObject) *model.AppError {
	return nil
}

func (rma *ReportMockApp) CompileReportChunks(format string, prefix string, numberOfChunks int, headers []string) *model.AppError {
	return nil
}

func (rma *ReportMockApp) SendReportToUser(rctx request.CTX, job *model.Job, format string) *model.AppError {
	return nil
}

func (rma *ReportMockApp) CleanupReportChunks(format string, prefix string, numberOfChunks int) *model.AppError {
	return nil
}

func TestBatchReportWorker(t *testing.T) {
	mainHelper.Parallel(t)

	setupBatchWorker := func(
		t *testing.T,
		th *TestHelper,
		getData func(jobData model.StringMap) ([]model.ReportableObject, model.StringMap, bool, error),
	) (*jobs.BatchReportWorker, *model.Job) {
		t.Helper()

		worker := jobs.MakeBatchReportWorker(
			th.Server.Jobs,
			th.Server.Store(),
			&ReportMockApp{},
			1*time.Second,
			"csv",
			[]string{},
			getData)
		job := th.SetupBatchWorker(t, worker.BatchWorker)
		return worker, job
	}

	waitForFileCount := func(t *testing.T, th *TestHelper, job *model.Job, fileCount int) {
		t.Helper()

		require.Eventuallyf(t, func() bool {
			actualJob, appErr := th.Server.Jobs.GetJob(th.Context, job.Id)
			require.Nil(t, appErr)
			require.Equal(t, job.Id, actualJob.Id)

			finalFileCount, err := strconv.Atoi(actualJob.Data["file_count"])
			require.NoError(t, err)
			return finalFileCount == fileCount
		}, 5*time.Second, 250*time.Millisecond, "job did not stop at batch %d", fileCount)
	}

	getFileCountFromData := func(t *testing.T, data model.StringMap) int {
		t.Helper()

		if data["file_count"] == "" {
			return 0
		}

		fileCount, err := strconv.Atoi(data["file_count"])
		require.NoError(t, err)

		return fileCount
	}

	createData := func(th *TestHelper, data model.StringMap) model.StringMap {
		data["requesting_user_id"] = th.SystemAdminUser.Id
		return data
	}

	t.Run("should finish when the report is done, incrementing file count along the way", func(t *testing.T) {
		mainHelper.Parallel(t)

		th := Setup(t).InitBasic(t)

		var worker model.Worker
		var job *model.Job

		iterations := 0

		worker, job = setupBatchWorker(t, th, func(data model.StringMap) ([]model.ReportableObject, model.StringMap, bool, error) {
			fileCount := getFileCountFromData(t, data)
			require.Equal(t, iterations, fileCount)
			require.LessOrEqual(t, fileCount, 3, "only 3 batches should have run")

			iterations++

			if fileCount >= 3 {
				go worker.Stop() // Shut down the worker when the job is done
				return []model.ReportableObject{}, createData(th, data), true, nil
			}

			return []model.ReportableObject{}, createData(th, data), false, nil
		})

		// Queue the work to be done
		worker.JobChannel() <- *job

		th.WaitForJobStatus(t, job, model.JobStatusSuccess)
		waitForFileCount(t, th, job, 3)
	})

	t.Run("should fail job when get data throws an error", func(t *testing.T) {
		mainHelper.Parallel(t)
		th := Setup(t).InitBasic(t)

		var worker model.Worker
		var job *model.Job
		worker, job = setupBatchWorker(t, th, func(data model.StringMap) ([]model.ReportableObject, model.StringMap, bool, error) {
			go worker.Stop() // Shut down the worker right after this
			return []model.ReportableObject{}, createData(th, data), false, errors.New("failed to fetch data")
		})

		// Queue the work to be done
		worker.JobChannel() <- *job

		th.WaitForJobStatus(t, job, model.JobStatusError)
	})
}