frostfs-sdk-go/pool/mock_test.go

package pool

import (
	"context"
	"crypto/ecdsa"
	"errors"

	"go.uber.org/zap"

	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/accounting"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/ape"
	sessionv2 "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/session"
	apistatus "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client/status"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	frostfsecdsa "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/crypto/ecdsa"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/netmap"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/session"
	"github.com/google/uuid"
)

type mockClient struct {
	key ecdsa.PrivateKey
	clientStatusMonitor

	errorOnDial          bool
	errorOnCreateSession bool
	errorOnEndpointInfo  error
	resOnEndpointInfo    netmap.NodeInfo
	healthcheckFn        func()
	errorOnNetworkInfo   bool
	stOnGetObject        apistatus.Status
}

var _ client = (*mockClient)(nil)

func newMockClient(addr string, key ecdsa.PrivateKey) *mockClient {
	return &mockClient{
		key:                 key,
		clientStatusMonitor: newClientStatusMonitor(zap.NewExample(), addr, 10),
	}
}

func newMockClientHealthy(addr string, key ecdsa.PrivateKey, healthy bool) *mockClient {
	m := newMockClient(addr, key)
	if healthy {
		m.setHealthy()
	} else {
		m.setUnhealthy()
	}
	return m
}

func (m *mockClient) setThreshold(threshold uint32) {
	m.errorThreshold = threshold
}

func (m *mockClient) errOnCreateSession() {
	m.errorOnCreateSession = true
}

func (m *mockClient) errOnEndpointInfo() {
	m.errorOnEndpointInfo = errors.New("error")
}

func (m *mockClient) errOnNetworkInfo() {
	m.errorOnEndpointInfo = errors.New("error")
}

func (m *mockClient) errOnDial() {
	m.errorOnDial = true
	m.errOnCreateSession()
	m.errOnEndpointInfo()
	m.errOnNetworkInfo()
}

func (m *mockClient) statusOnGetObject(st apistatus.Status) {
	m.stOnGetObject = st
}

func newToken(key ecdsa.PrivateKey) *session.Object {
	var tok session.Object
	tok.SetID(uuid.New())
	pk := frostfsecdsa.PublicKey(key.PublicKey)
	tok.SetAuthKey(&pk)

	return &tok
}

func (m *mockClient) balanceGet(context.Context, PrmBalanceGet) (accounting.Decimal, error) {
	return accounting.Decimal{}, nil
}

func (m *mockClient) containerPut(context.Context, PrmContainerPut) (cid.ID, error) {
	return cid.ID{}, nil
}

func (m *mockClient) containerGet(context.Context, PrmContainerGet) (container.Container, error) {
	return container.Container{}, nil
}

func (m *mockClient) containerList(context.Context, PrmContainerList) ([]cid.ID, error) {
	return nil, nil
}

func (m *mockClient) containerListStream(context.Context, PrmListStream) (ResListStream, error) {
	return ResListStream{}, nil
}

func (m *mockClient) containerDelete(context.Context, PrmContainerDelete) error {
	return nil
}

func (m *mockClient) apeManagerAddChain(context.Context, PrmAddAPEChain) error {
	return nil
}

func (m *mockClient) apeManagerRemoveChain(context.Context, PrmRemoveAPEChain) error {
	return nil
}

func (m *mockClient) apeManagerListChains(context.Context, PrmListAPEChains) ([]ape.Chain, error) {
	return []ape.Chain{}, nil
}

func (m *mockClient) endpointInfo(ctx context.Context, _ prmEndpointInfo) (netmap.NodeInfo, error) {
	if m.errorOnEndpointInfo != nil {
		return netmap.NodeInfo{}, m.handleError(ctx, nil, m.errorOnEndpointInfo)
	}

	m.resOnEndpointInfo.SetNetworkEndpoints(m.addr)
	return m.resOnEndpointInfo, nil
}

func (m *mockClient) healthcheck(ctx context.Context) (netmap.NodeInfo, error) {
	if m.healthcheckFn != nil {
		m.healthcheckFn()
	}
	return m.endpointInfo(ctx, prmEndpointInfo{})
}

func (m *mockClient) networkInfo(ctx context.Context, _ prmNetworkInfo) (netmap.NetworkInfo, error) {
	var ni netmap.NetworkInfo

	if m.errorOnNetworkInfo {
		return ni, m.handleError(ctx, nil, errors.New("error"))
	}

	return ni, nil
}

func (m *mockClient) netMapSnapshot(context.Context, prmNetMapSnapshot) (netmap.NetMap, error) {
	var nm netmap.NetMap
	return nm, nil
}

func (m *mockClient) objectPut(context.Context, PrmObjectPut) (ResPutObject, error) {
	return ResPutObject{}, nil
}

func (m *mockClient) objectPatch(context.Context, PrmObjectPatch) (ResPatchObject, error) {
	return ResPatchObject{}, nil
}

func (m *mockClient) objectDelete(context.Context, PrmObjectDelete) error {
	return nil
}

func (m *mockClient) objectGet(ctx context.Context, _ PrmObjectGet) (ResGetObject, error) {
	var res ResGetObject

	if m.stOnGetObject == nil {
		return res, nil
	}

	status := apistatus.ErrFromStatus(m.stOnGetObject)
	return res, m.handleError(ctx, status, nil)
}

func (m *mockClient) objectHead(context.Context, PrmObjectHead) (object.Object, error) {
	return object.Object{}, nil
}

func (m *mockClient) objectRange(context.Context, PrmObjectRange) (ResObjectRange, error) {
	return ResObjectRange{}, nil
}

func (m *mockClient) objectSearch(context.Context, PrmObjectSearch) (ResObjectSearch, error) {
	return ResObjectSearch{}, nil
}

func (m *mockClient) sessionCreate(ctx context.Context, _ prmCreateSession) (resCreateSession, error) {
	if m.errorOnCreateSession {
		return resCreateSession{}, m.handleError(ctx, nil, errors.New("error"))
	}

	tok := newToken(m.key)

	var v2tok sessionv2.Token
	tok.WriteToV2(&v2tok)

	return resCreateSession{
		id:         v2tok.GetBody().GetID(),
		sessionKey: v2tok.GetBody().GetSessionKey(),
	}, nil
}

func (m *mockClient) dial(context.Context) error {
	if m.errorOnDial {
		return errors.New("dial error")
	}
	return nil
}

func (m *mockClient) restart(context.Context) error {
	if m.errorOnDial {
		return errors.New("restart dial error")
	}

	return nil
}

func (m *mockClient) close() error {
	return nil
}