xk6-frostfs/internal/registry/obj_exporter_test.go

package registry

import (
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

type expectedResult struct {
	mode     string
	objects  []ObjectInfo
	dir      string
	dbName   string
	jsonName string
}

func TestObjectExporter(t *testing.T) {
	names := []string{"s3", "grpc"}
	for _, name := range names {
		t.Run(name, runExportTest)
		t.Run(name+"-changed", runExportChangedTest)
		t.Run(name+"-empty", runExportEmptyTest)
	}
}

func runExportTest(t *testing.T) {
	expected := getExpectedResult(t)
	objReg := getFilledRegistry(t, expected)
	objExp := NewObjExporter(NewObjSelector(objReg, 0, SelectorOneshot, &ObjFilter{Status: statusCreated}))

	require.NoError(t, objExp.ExportJSONPreGen(expected.jsonName))
	require.NoError(t, checkExported(expected.objects, expected.jsonName))
}

func runExportChangedTest(t *testing.T) {
	expected := getExpectedResult(t)
	objReg := getFilledRegistry(t, expected)

	newStatus := randString(10)
	num := randPositiveInt(1, len(expected.objects))
	changedObjects := make([]ObjectInfo, num)
	require.Equal(t, num, copy(changedObjects[:], expected.objects[:]))

	sel := NewObjSelector(objReg, 0, SelectorOneshot, &ObjFilter{Status: statusCreated})
	for i := range changedObjects {
		changedObjects[i].Status = newStatus
		require.NoError(t, objReg.SetObjectStatus(sel.NextObject().Id, statusCreated, newStatus))
	}

	objExp := NewObjExporter(NewObjSelector(objReg, 0, SelectorOneshot, &ObjFilter{Status: newStatus}))
	require.NoError(t, objExp.ExportJSONPreGen(expected.jsonName))
	require.NoError(t, checkExported(changedObjects, expected.jsonName))
}

func runExportEmptyTest(t *testing.T) {
	expected := getExpectedResult(t)
	expected.objects = make([]ObjectInfo, 0)
	objReg := getFilledRegistry(t, expected)
	objExp := NewObjExporter(NewObjSelector(objReg, 0, SelectorOneshot, &ObjFilter{Status: statusCreated}))

	require.NoError(t, objExp.ExportJSONPreGen(expected.jsonName))
	require.NoError(t, checkExported(expected.objects, expected.jsonName))
}

func getExpectedResult(t *testing.T) expectedResult {
	num := randPositiveInt(2, 100)
	mode := getMode(t.Name())
	require.NotEqual(t, "", mode, "test mode should contain either \"s3\" or\"grpc\"")
	dir := t.TempDir()
	res := expectedResult{
		mode:     mode,
		objects:  generateObjectInfo(num, t.Name()),
		dir:      dir,
		dbName:   filepath.Join(dir, "registry-"+mode+".db"),
		jsonName: filepath.Join(dir, "registry-"+mode+".json"),
	}
	return res
}

func randPositiveInt(min, max int) int {
	return rand.Intn(max-min) + min
}

func getMode(name string) (res string) {
	if strings.Contains(name, "s3") {
		res = filepath.Base(name)
	}
	if strings.Contains(name, "grpc") {
		res = filepath.Base(name)
	}
	return res
}

func generateObjectInfo(num int, mode string) []ObjectInfo {
	res := make([]ObjectInfo, num)
	for i := range res {
		res[i] = randomObjectInfo()
		if !strings.Contains(mode, "s3") {
			res[i].S3Bucket = ""
			res[i].S3Key = ""
		}
		if !strings.Contains(mode, "grpc") {
			res[i].CID = ""
			res[i].OID = ""
		}
	}
	return res
}

func getFilledRegistry(t *testing.T, expected expectedResult) *ObjRegistry {
	objReg := NewObjRegistry(context.Background(), expected.dbName)
	for i := range expected.objects {
		require.NoError(t, objReg.AddObject(expected.objects[i].CID, expected.objects[i].OID, expected.objects[i].S3Bucket, expected.objects[i].S3Key, expected.objects[i].PayloadHash))
	}
	return objReg
}

func checkExported(expected []ObjectInfo, fileName string) error {
	file, err := os.ReadFile(fileName)
	if err != nil {
		return err
	}

	if !json.Valid(file) {
		return fmt.Errorf("exported json file %s is invalid", fileName)
	}

	var actual PreGenerateInfo
	if json.Unmarshal(file, &actual) != nil {
		return err
	}

	if len(expected) != len(actual.Objects) {
		return fmt.Errorf("expected len(): %v, got len(): %v", len(expected), len(actual.Objects))
	}

	for i := range expected {
		if !slices.ContainsFunc(actual.Objects, func(oi ObjInfo) bool {
			compareS3 := oi.Bucket == expected[i].S3Bucket && oi.Object == expected[i].S3Key
			comparegRPC := oi.CID == expected[i].CID && oi.OID == expected[i].OID
			return compareS3 && comparegRPC
		}) {
			return fmt.Errorf("object %v not found in exported json file %s", expected[i], fileName)
		}
	}

	return nil
}