76 lines
1.5 KiB
Go
76 lines
1.5 KiB
Go
package testsupport
|
|
|
|
import (
|
|
"context"
|
|
"git.wamblee.org/converge/pkg/support/pprof"
|
|
"github.com/stretchr/testify/suite"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
_ "runtime/pprof"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type TestFunction func() any
|
|
|
|
func RunAndWait(suite *suite.Suite, functions ...TestFunction) []any {
|
|
wg := sync.WaitGroup{}
|
|
wg.Add(len(functions))
|
|
res := make([]any, len(functions))
|
|
for i, function := range functions {
|
|
go func() {
|
|
defer func() {
|
|
wg.Done()
|
|
}()
|
|
res[i] = function()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
return res
|
|
}
|
|
|
|
func StartPprof(port string) *http.Server {
|
|
if os.Getenv("PPROF") == "" {
|
|
return nil
|
|
}
|
|
if port == "" {
|
|
port = ":9000"
|
|
}
|
|
mux := http.NewServeMux()
|
|
pprof.RegisterPprof(mux, "/debug/pprof")
|
|
srv := http.Server{
|
|
Addr: port,
|
|
Handler: mux,
|
|
}
|
|
go func() {
|
|
if err := srv.ListenAndServe(); err != http.ErrServerClosed {
|
|
log.Fatalf("Could not start pprof listener: %v", err)
|
|
}
|
|
log.Println("Test pprof server started: " + port)
|
|
}()
|
|
return &srv
|
|
}
|
|
|
|
func StopPprof(ctx context.Context, server *http.Server) {
|
|
if os.Getenv("PPROF") == "" {
|
|
return
|
|
}
|
|
err := server.Shutdown(ctx)
|
|
if err != nil {
|
|
log.Println("Error shutting down test pprof server")
|
|
return
|
|
}
|
|
log.Println("Test pprof server stopped")
|
|
}
|
|
|
|
func CreateTestContext(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
|
|
ctx, cancelFunc := context.WithCancel(ctx)
|
|
ctx, timeoutCancelFunc := context.WithTimeout(ctx, timeout)
|
|
compositeCancelFunc := func() {
|
|
timeoutCancelFunc()
|
|
cancelFunc()
|
|
}
|
|
return ctx, compositeCancelFunc
|
|
}
|