frostfs-node/pkg/local_object_storage/internal/testutil/generators.go

package testutil

import (
	"encoding/binary"
	"testing"

	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"
	"golang.org/x/exp/rand"
)

// AddressGenerator is the interface of types that generate object addresses.
type AddressGenerator interface {
	Next() oid.Address
}

// SeqAddrGenerator is an AddressGenerator that generates addresses sequentially and wraps around the given max ID.
type SeqAddrGenerator struct {
	cnt   atomic.Uint64
	MaxID uint64
}

var _ AddressGenerator = &SeqAddrGenerator{}

func (g *SeqAddrGenerator) Next() oid.Address {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], ((g.cnt.Inc()-1)%g.MaxID)+1)
	var addr oid.Address
	addr.SetContainer(cid.ID{})
	addr.SetObject(id)
	return addr
}

// RandAddrGenerator is an addressGenerator that generates random addresses in the given range.
type RandAddrGenerator uint64

func (g RandAddrGenerator) Next() oid.Address {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], uint64(1+int(rand.Int63n(int64(g)))))
	var addr oid.Address
	addr.SetContainer(cid.ID{})
	addr.SetObject(id)
	return addr
}

// ObjectGenerator is the interface of types that generate object entries.
type ObjectGenerator interface {
	Next() *object.Object
}

// SeqObjGenerator is an ObjectGenerator that generates entries with random payloads of size objSize and sequential IDs.
type SeqObjGenerator struct {
	cnt     atomic.Uint64
	ObjSize uint64
}

var _ ObjectGenerator = &SeqObjGenerator{}

func generateObjectWithOIDWithCIDWithSize(oid oid.ID, cid cid.ID, sz uint64) *object.Object {
	data := make([]byte, sz)
	_, _ = rand.Read(data)
	obj := GenerateObjectWithCIDWithPayload(cid, data)
	obj.SetID(oid)
	return obj
}

func (g *SeqObjGenerator) Next() *object.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], g.cnt.Inc())
	return generateObjectWithOIDWithCIDWithSize(id, cid.ID{}, g.ObjSize)
}

// RandObjGenerator is an ObjectGenerator that generates entries with random IDs and payloads of size objSize.
type RandObjGenerator struct {
	ObjSize uint64
}

var _ ObjectGenerator = &RandObjGenerator{}

func (g *RandObjGenerator) Next() *object.Object {
	var id oid.ID
	rand.Read(id[:])
	return generateObjectWithOIDWithCIDWithSize(id, cid.ID{}, g.ObjSize)
}

// OverwriteObjGenerator is an ObjectGenerator that generates entries with random payloads of size objSize and at most maxObjects distinct IDs.
type OverwriteObjGenerator struct {
	ObjSize    uint64
	MaxObjects uint64
}

func (g *OverwriteObjGenerator) Next() *object.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], uint64(1+rand.Int63n(int64(g.MaxObjects))))
	return generateObjectWithOIDWithCIDWithSize(id, cid.ID{}, g.ObjSize)
}

func AddressFromObject(t testing.TB, obj *object.Object) oid.Address {
	var addr oid.Address

	id, isSet := obj.ID()
	require.True(t, isSet, "object ID is not set")
	addr.SetObject(id)

	cid, isSet := obj.ContainerID()
	require.True(t, isSet, "container ID is not set")
	addr.SetContainer(cid)

	return addr
}