deploy/test/testutil.go

package test

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/client-go/tools/cache"
	"sigs.k8s.io/yaml"

	"github.com/hanzoai/cd/pkg/apis/application/v1alpha1"
)

// StartInformer is a helper to start an informer, wait for its cache to sync and return a cancel func
func StartInformer(informer cache.SharedIndexInformer) context.CancelFunc {
	ctx, cancel := context.WithCancel(context.Background())
	go informer.Run(ctx.Done())
	if !cache.WaitForCacheSync(ctx.Done(), informer.HasSynced) {
		log.Fatal("Timed out waiting for informer cache to sync")
	}
	return cancel
}

// GetFreePort finds an available free port on the OS
func GetFreePort() (int, error) {
	lc := &net.ListenConfig{}
	ln, err := lc.Listen(context.Background(), "tcp", "[::]:0")
	if err != nil {
		return 0, err
	}
	return ln.Addr().(*net.TCPAddr).Port, ln.Close()
}

// WaitForPortListen waits until the given address is listening on the port
func WaitForPortListen(addr string, timeout time.Duration) error {
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()
	timer := time.NewTimer(timeout)
	if timeout == 0 {
		timer.Stop()
	} else {
		defer timer.Stop()
	}
	for {
		select {
		case <-ticker.C:
			if portIsOpen(addr) {
				return nil
			}
		case <-timer.C:
			return fmt.Errorf("timeout after %s", timeout.String())
		}
	}
}

func portIsOpen(addr string) bool {
	d := &net.Dialer{}
	conn, err := d.DialContext(context.Background(), "tcp", addr)
	if err != nil {
		return false
	}
	_ = conn.Close()
	return true
}

// Read the contents of a file and returns it as string. Panics on error.
func MustLoadFileToString(path string) string {
	o, err := os.ReadFile(path)
	if err != nil {
		panic(err.Error())
	}
	return string(o)
}

func YamlToUnstructured(yamlStr string) *unstructured.Unstructured {
	obj := make(map[string]any)
	err := yaml.Unmarshal([]byte(yamlStr), &obj)
	if err != nil {
		panic(err)
	}
	return &unstructured.Unstructured{Object: obj}
}

func YamlToApplication(yamlStr string) *v1alpha1.Application {
	app := v1alpha1.Application{}
	err := yaml.Unmarshal([]byte(yamlStr), &app)
	if err != nil {
		panic(err)
	}
	return &app
}

// ToMap converts any object to a map[string]any
func ToMap(obj any) map[string]any {
	data, err := json.Marshal(obj)
	if err != nil {
		panic(err)
	}
	var res map[string]any
	err = json.Unmarshal(data, &res)
	if err != nil {
		panic(err)
	}
	return res
}

// GetTestDir will return the full directory path of the
// calling test file.
func GetTestDir(t *testing.T) string {
	t.Helper()
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	return cwd
}

var logCaptureMutex sync.Mutex

type threadSafeBuffer struct {
	buf bytes.Buffer
	mu  sync.RWMutex
}

func (tsb *threadSafeBuffer) Write(p []byte) (n int, err error) {
	tsb.mu.Lock()
	defer tsb.mu.Unlock()
	return tsb.buf.Write(p)
}

func (tsb *threadSafeBuffer) String() string {
	tsb.mu.RLock()
	defer tsb.mu.RUnlock()
	return tsb.buf.String()
}

// CaptureLogEntries captures log entries generated by the logger and return it as string
func CaptureLogEntries(run func()) string {
	logCaptureMutex.Lock()
	defer logCaptureMutex.Unlock()

	f := log.StandardLogger().Formatter
	log.SetFormatter(&log.TextFormatter{DisableColors: true})
	defer log.SetFormatter(f)
	output := &threadSafeBuffer{}
	original := log.StandardLogger().Out
	log.SetOutput(output)
	log.SetLevel(log.DebugLevel)
	defer log.SetOutput(original)

	run()

	return output.String()
}