frostfs-node/scripts/populate-metabase/internal/generate.go

package internal

import (
	"crypto/sha256"
	"fmt"

	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/checksum"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id/test"
	objecttest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/test"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/user"
	usertest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/user/test"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/version"
	"git.frostfs.info/TrueCloudLab/tzhash/tz"
	"golang.org/x/exp/rand"
)

func GeneratePayloadPool(count uint, size uint) [][]byte {
	var pool [][]byte
	for range count {
		payload := make([]byte, size)
		_, _ = rand.Read(payload)

		pool = append(pool, payload)
	}
	return pool
}

func GenerateAttributePool(count uint) []objectSDK.Attribute {
	var pool []objectSDK.Attribute
	for i := range count {
		for j := range count {
			attr := *objectSDK.NewAttribute()
			attr.SetKey(fmt.Sprintf("key%d", i))
			attr.SetValue(fmt.Sprintf("value%d", j))
			pool = append(pool, attr)
		}
	}
	return pool
}

func GenerateOwnerPool(count uint) []user.ID {
	var pool []user.ID
	for range count {
		pool = append(pool, usertest.ID())
	}
	return pool
}

type ObjectOption func(obj *objectSDK.Object)

func GenerateObject(options ...ObjectOption) *objectSDK.Object {
	var ver version.Version
	ver.SetMajor(2)
	ver.SetMinor(1)

	payload := make([]byte, 0)

	var csum checksum.Checksum
	csum.SetSHA256(sha256.Sum256(payload))

	var csumTZ checksum.Checksum
	csumTZ.SetTillichZemor(tz.Sum(csum.Value()))

	obj := objectSDK.New()
	obj.SetID(oidtest.ID())
	obj.SetOwnerID(usertest.ID())
	obj.SetContainerID(cidtest.ID())

	header := objecttest.Object().GetECHeader()
	header.SetParent(oidtest.ID())
	obj.SetECHeader(header)

	obj.SetVersion(&ver)
	obj.SetPayload(payload)
	obj.SetPayloadSize(uint64(len(payload)))
	obj.SetPayloadChecksum(csum)
	obj.SetPayloadHomomorphicHash(csumTZ)

	for _, option := range options {
		option(obj)
	}

	return obj
}

func WithContainerID(cid cid.ID) ObjectOption {
	return func(obj *objectSDK.Object) {
		obj.SetContainerID(cid)
	}
}

func WithType(typ objectSDK.Type) ObjectOption {
	return func(obj *objectSDK.Object) {
		obj.SetType(typ)
	}
}

func WithPayloadFromPool(pool [][]byte) ObjectOption {
	payload := pool[rand.Intn(len(pool))]

	var csum checksum.Checksum
	csum.SetSHA256(sha256.Sum256(payload))

	var csumTZ checksum.Checksum
	csumTZ.SetTillichZemor(tz.Sum(csum.Value()))

	return func(obj *objectSDK.Object) {
		obj.SetPayload(payload)
		obj.SetPayloadSize(uint64(len(payload)))
		obj.SetPayloadChecksum(csum)
		obj.SetPayloadHomomorphicHash(csumTZ)
	}
}

func WithAttributesFromPool(pool []objectSDK.Attribute, count uint) ObjectOption {
	return func(obj *objectSDK.Object) {
		var attrs []objectSDK.Attribute
		for range count {
			attrs = append(attrs, pool[rand.Intn(len(pool))])
		}
		obj.SetAttributes(attrs...)
	}
}

func WithOwnerIDFromPool(pool []user.ID) ObjectOption {
	return func(obj *objectSDK.Object) {
		obj.SetOwnerID(pool[rand.Intn(len(pool))])
	}
}