egress-proxies-operator/internal/provider/kvm/fake_test.go

package kvm

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"slices"
	"strings"
	"sync"

	libvirt "github.com/digitalocean/go-libvirt"
)

// fakeAPI is an in-memory libvirtAPI: pools of volumes and a set of
// domains, with per-method error injection. It models just enough libvirt
// semantics for the provider's decisions to be observable — ErrStorageVolExist
// on duplicate create, ErrNoDomain/ErrNoStorageVol on missing lookups,
// ErrOperationInvalid on destroying a stopped domain, and so on.
type fakeAPI struct {
	mu       sync.Mutex
	poolPath string
	pools    map[string]map[string]*fakeVol
	domains  map[string]*fakeDomain
	calls    []string
	fail     map[string]error
	failOnce map[string]error
	// uploadHook, when set, sees every StorageVolUpload before the fake
	// reads the stream; it can return an error to simulate a broken upload.
	uploadHook func(vol string, r io.Reader) error
	// hostname, libVersion, poolCapacity and poolStates answer the health
	// check's RPCs. A pool missing from poolStates is running.
	hostname     string
	libVersion   uint64
	poolCapacity uint64
	poolStates   map[string]libvirt.StoragePoolState
}

type fakeVol struct {
	format   string
	capacity int64
	backing  string
	data     []byte
}

type fakeDomain struct {
	name     string
	xml      string
	metadata string
	running  bool
	// addrs is what DomainInterfaceAddresses returns per source.
	addrs map[uint32][]libvirt.DomainInterface
}

func newFakeAPI(pools ...string) *fakeAPI {
	f := &fakeAPI{
		poolPath: "/var/lib/libvirt/images",
		pools:    map[string]map[string]*fakeVol{},
		domains:  map[string]*fakeDomain{},
		fail:     map[string]error{},
		failOnce: map[string]error{},

		hostname:     "fake-host",
		libVersion:   10002000,
		poolCapacity: 100 << 30,
		poolStates:   map[string]libvirt.StoragePoolState{},
	}
	for _, p := range pools {
		f.pools[p] = map[string]*fakeVol{}
	}
	return f
}

// enter records the call and returns any injected error for method.
func (f *fakeAPI) enter(method string) error {
	f.calls = append(f.calls, method)
	if err, ok := f.failOnce[method]; ok {
		delete(f.failOnce, method)
		return err
	}
	return f.fail[method]
}

func (f *fakeAPI) callCount(method string) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	n := 0
	for _, c := range f.calls {
		if c == method {
			n++
		}
	}
	return n
}

func (f *fakeAPI) hasCall(method string) bool { return f.callCount(method) > 0 }

func (f *fakeAPI) vol(pool, name string) *fakeVol {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.pools[pool][name]
}

func (f *fakeAPI) putVol(pool, name, format string, data []byte) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.pools[pool][name] = &fakeVol{format: format, capacity: int64(len(data)), data: data}
}

func (f *fakeAPI) volNames(pool string) []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []string
	for n := range f.pools[pool] {
		out = append(out, n)
	}
	slices.Sort(out)
	return out
}

func (f *fakeAPI) domain(name string) *fakeDomain {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.domains[name]
}

func (f *fakeAPI) putDomain(d *fakeDomain) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.domains[d.name] = d
}

func lvError(code libvirt.ErrorNumber, format string, args ...any) error {
	return libvirt.Error{Code: uint32(code), Message: fmt.Sprintf(format, args...)}
}

// --- domains ---

func (f *fakeAPI) DomainLookupByName(name string) (libvirt.Domain, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainLookupByName"); err != nil {
		return libvirt.Domain{}, err
	}
	if _, ok := f.domains[name]; !ok {
		return libvirt.Domain{}, lvError(libvirt.ErrNoDomain, "Domain not found: no domain with matching name '%s'", name)
	}
	return libvirt.Domain{Name: name}, nil
}

func (f *fakeAPI) DomainDefineXML(raw string) (libvirt.Domain, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainDefineXML"); err != nil {
		return libvirt.Domain{}, err
	}
	var d domainXML
	if err := xml.Unmarshal([]byte(raw), &d); err != nil {
		return libvirt.Domain{}, lvError(libvirt.ErrXMLError, "%v", err)
	}
	meta, err := xml.Marshal(d.Metadata.Instance)
	if err != nil {
		return libvirt.Domain{}, lvError(libvirt.ErrXMLError, "%v", err)
	}
	// libvirt hands the element back without its namespace declaration
	// (seen on 11.6); mimic that so Get/ListByTag parse the real shape.
	stored := strings.Replace(string(meta), ` xmlns="`+metadataURI+`"`, "", 1)
	existing := f.domains[d.Name]
	nd := &fakeDomain{name: d.Name, xml: raw, metadata: stored}
	if existing != nil {
		nd.running = existing.running
		nd.addrs = existing.addrs
	}
	f.domains[d.Name] = nd
	return libvirt.Domain{Name: d.Name}, nil
}

func (f *fakeAPI) DomainCreate(d libvirt.Domain) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainCreate"); err != nil {
		return err
	}
	dom, ok := f.domains[d.Name]
	if !ok {
		return lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	if dom.running {
		return lvError(libvirt.ErrOperationInvalid, "domain is already running")
	}
	dom.running = true
	return nil
}

func (f *fakeAPI) DomainDestroy(d libvirt.Domain) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainDestroy"); err != nil {
		return err
	}
	dom, ok := f.domains[d.Name]
	if !ok {
		return lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	if !dom.running {
		return lvError(libvirt.ErrOperationInvalid, "domain is not running")
	}
	dom.running = false
	return nil
}

func (f *fakeAPI) DomainUndefineFlags(d libvirt.Domain, _ libvirt.DomainUndefineFlagsValues) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainUndefineFlags"); err != nil {
		return err
	}
	if _, ok := f.domains[d.Name]; !ok {
		return lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	delete(f.domains, d.Name)
	return nil
}

func (f *fakeAPI) DomainGetState(d libvirt.Domain, _ uint32) (int32, int32, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainGetState"); err != nil {
		return 0, 0, err
	}
	dom, ok := f.domains[d.Name]
	if !ok {
		return 0, 0, lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	if dom.running {
		return int32(libvirt.DomainRunning), 0, nil
	}
	return int32(libvirt.DomainShutoff), 0, nil
}

func (f *fakeAPI) DomainGetMetadata(d libvirt.Domain, typ int32, uri libvirt.OptString, _ libvirt.DomainModificationImpact) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainGetMetadata"); err != nil {
		return "", err
	}
	dom, ok := f.domains[d.Name]
	if !ok {
		return "", lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	if typ != int32(libvirt.DomainMetadataElement) || len(uri) != 1 || uri[0] != metadataURI || dom.metadata == "" {
		return "", lvError(libvirt.ErrNoDomainMetadata, "Requested metadata element is not present")
	}
	return dom.metadata, nil
}

func (f *fakeAPI) DomainInterfaceAddresses(d libvirt.Domain, source uint32, _ uint32) ([]libvirt.DomainInterface, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("DomainInterfaceAddresses"); err != nil {
		return nil, err
	}
	dom, ok := f.domains[d.Name]
	if !ok {
		return nil, lvError(libvirt.ErrNoDomain, "no domain %s", d.Name)
	}
	if !dom.running {
		return nil, lvError(libvirt.ErrOperationInvalid, "domain is not running")
	}
	if source == uint32(libvirt.DomainInterfaceAddressesSrcAgent) && dom.addrs[source] == nil {
		return nil, lvError(libvirt.ErrAgentUnresponsive, "Guest agent is not responding")
	}
	return dom.addrs[source], nil
}

func (f *fakeAPI) ConnectListAllDomains(_ int32, _ libvirt.ConnectListAllDomainsFlags) ([]libvirt.Domain, uint32, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("ConnectListAllDomains"); err != nil {
		return nil, 0, err
	}
	names := make([]string, 0, len(f.domains))
	for n := range f.domains {
		names = append(names, n)
	}
	slices.Sort(names)
	out := make([]libvirt.Domain, 0, len(names))
	for _, n := range names {
		out = append(out, libvirt.Domain{Name: n})
	}
	return out, uint32(len(out)), nil
}

// --- storage ---

func (f *fakeAPI) StoragePoolLookupByName(name string) (libvirt.StoragePool, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StoragePoolLookupByName"); err != nil {
		return libvirt.StoragePool{}, err
	}
	if _, ok := f.pools[name]; !ok {
		return libvirt.StoragePool{}, lvError(libvirt.ErrNoStoragePool, "Storage pool not found: no storage pool with matching name '%s'", name)
	}
	return libvirt.StoragePool{Name: name}, nil
}

func (f *fakeAPI) StoragePoolRefresh(pool libvirt.StoragePool, _ uint32) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StoragePoolRefresh"); err != nil {
		return err
	}
	if _, ok := f.pools[pool.Name]; !ok {
		return lvError(libvirt.ErrNoStoragePool, "no pool %s", pool.Name)
	}
	return nil
}

func (f *fakeAPI) StorageVolLookupByName(pool libvirt.StoragePool, name string) (libvirt.StorageVol, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StorageVolLookupByName"); err != nil {
		return libvirt.StorageVol{}, err
	}
	if _, ok := f.pools[pool.Name][name]; !ok {
		return libvirt.StorageVol{}, lvError(libvirt.ErrNoStorageVol, "Storage volume not found: no storage vol with matching name '%s'", name)
	}
	return libvirt.StorageVol{Pool: pool.Name, Name: name, Key: f.poolPath + "/" + name}, nil
}

func (f *fakeAPI) StorageVolGetPath(vol libvirt.StorageVol) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StorageVolGetPath"); err != nil {
		return "", err
	}
	if _, ok := f.pools[vol.Pool][vol.Name]; !ok {
		return "", lvError(libvirt.ErrNoStorageVol, "no vol %s", vol.Name)
	}
	return f.poolPath + "/" + vol.Name, nil
}

func (f *fakeAPI) StorageVolCreateXML(pool libvirt.StoragePool, raw string, _ libvirt.StorageVolCreateFlags) (libvirt.StorageVol, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StorageVolCreateXML"); err != nil {
		return libvirt.StorageVol{}, err
	}
	vols, ok := f.pools[pool.Name]
	if !ok {
		return libvirt.StorageVol{}, lvError(libvirt.ErrNoStoragePool, "no pool %s", pool.Name)
	}
	var v volumeXML
	if err := xml.Unmarshal([]byte(raw), &v); err != nil {
		return libvirt.StorageVol{}, lvError(libvirt.ErrXMLError, "%v", err)
	}
	if _, exists := vols[v.Name]; exists {
		return libvirt.StorageVol{}, lvError(libvirt.ErrStorageVolExist, "storage volume '%s' exists already", v.Name)
	}
	fv := &fakeVol{format: v.Target.Format.Type, capacity: v.Capacity.Value}
	if v.BackingStore != nil {
		fv.backing = v.BackingStore.Path
	}
	vols[v.Name] = fv
	return libvirt.StorageVol{Pool: pool.Name, Name: v.Name, Key: f.poolPath + "/" + v.Name}, nil
}

func (f *fakeAPI) StorageVolUpload(vol libvirt.StorageVol, r io.Reader, _, _ uint64, _ libvirt.StorageVolUploadFlags) error {
	f.mu.Lock()
	if err := f.enter("StorageVolUpload"); err != nil {
		f.mu.Unlock()
		return err
	}
	fv, ok := f.pools[vol.Pool][vol.Name]
	hook := f.uploadHook
	f.mu.Unlock()
	if !ok {
		return lvError(libvirt.ErrNoStorageVol, "no vol %s", vol.Name)
	}
	if hook != nil {
		if err := hook(vol.Name, r); err != nil {
			return err
		}
	}
	// Read outside the lock: the stream may be a slow HTTP body.
	data, err := io.ReadAll(r)
	if err != nil {
		return fmt.Errorf("reading upload stream: %w", err)
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	fv.data = data
	return nil
}

func (f *fakeAPI) StorageVolDownload(vol libvirt.StorageVol, w io.Writer, _, _ uint64, _ libvirt.StorageVolDownloadFlags) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StorageVolDownload"); err != nil {
		return err
	}
	fv, ok := f.pools[vol.Pool][vol.Name]
	if !ok {
		return lvError(libvirt.ErrNoStorageVol, "no vol %s", vol.Name)
	}
	_, err := io.Copy(w, bytes.NewReader(fv.data))
	return err
}

func (f *fakeAPI) StorageVolDelete(vol libvirt.StorageVol, _ libvirt.StorageVolDeleteFlags) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StorageVolDelete"); err != nil {
		return err
	}
	if _, ok := f.pools[vol.Pool][vol.Name]; !ok {
		return lvError(libvirt.ErrNoStorageVol, "no vol %s", vol.Name)
	}
	delete(f.pools[vol.Pool], vol.Name)
	return nil
}

// --- host ---

func (f *fakeAPI) ConnectGetLibVersion() (uint64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("ConnectGetLibVersion"); err != nil {
		return 0, err
	}
	return f.libVersion, nil
}

func (f *fakeAPI) ConnectGetHostname() (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("ConnectGetHostname"); err != nil {
		return "", err
	}
	return f.hostname, nil
}

// StoragePoolGetInfo reports the pool's state and poolCapacity minus the
// volumes' capacities as available space.
func (f *fakeAPI) StoragePoolGetInfo(pool libvirt.StoragePool) (uint8, uint64, uint64, uint64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.enter("StoragePoolGetInfo"); err != nil {
		return 0, 0, 0, 0, err
	}
	vols, ok := f.pools[pool.Name]
	if !ok {
		return 0, 0, 0, 0, lvError(libvirt.ErrNoStoragePool, "no pool %s", pool.Name)
	}
	var used uint64
	for _, v := range vols {
		used += uint64(v.capacity)
	}
	state := libvirt.StoragePoolRunning
	if s, ok := f.poolStates[pool.Name]; ok {
		state = s
	}
	return uint8(state), f.poolCapacity, used, f.poolCapacity - used, nil
}