package testkit import ( "context" "fmt" "strings" "time" "github.com/davecgh/go-spew/spew" "github.com/testground/sdk-go/run" "github.com/testground/sdk-go/runtime" ) type TestEnvironment struct { *runtime.RunEnv *run.InitContext Role string } // workaround for default params being wrapped in quote chars func (t *TestEnvironment) StringParam(name string) string { return strings.Trim(t.RunEnv.StringParam(name), "\"") } func (t *TestEnvironment) DurationParam(name string) time.Duration { d, err := time.ParseDuration(t.StringParam(name)) if err != nil { panic(fmt.Errorf("invalid duration value for param '%s': %w", name, err)) } return d } func (t *TestEnvironment) DebugSpew(format string, args ...interface{}) { t.RecordMessage(spew.Sprintf(format, args...)) } // WaitUntilAllDone waits until all instances in the test case are done. func (t *TestEnvironment) WaitUntilAllDone() { ctx := context.Background() t.SyncClient.MustSignalAndWait(ctx, StateDone, t.TestInstanceCount) } // WrapTestEnvironment takes a test case function that accepts a // *TestEnvironment, and adapts it to the original unwrapped SDK style // (run.InitializedTestCaseFn). func WrapTestEnvironment(f func(t *TestEnvironment) error) run.InitializedTestCaseFn { return func(runenv *runtime.RunEnv, initCtx *run.InitContext) error { t := &TestEnvironment{RunEnv: runenv, InitContext: initCtx} t.Role = t.StringParam("role") return f(t) } }