package testkit import ( "context" "encoding/json" "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) DurationRangeParam(name string) DurationRange { var r DurationRange t.JSONParam(name, &r) return r } func (t *TestEnvironment) FloatRangeParam(name string) FloatRange { r := FloatRange{} t.JSONParam(name, &r) return r } func (t *TestEnvironment) DebugSpew(format string, args ...interface{}) { t.RecordMessage(spew.Sprintf(format, args...)) } func (t *TestEnvironment) DumpJSON(filename string, v interface{}) { b, err := json.Marshal(v) if err != nil { t.RecordMessage("unable to marshal object to JSON: %s", err) return } f, err := t.CreateRawAsset(filename) if err != nil { t.RecordMessage("unable to create asset file: %s", err) return } defer f.Close() _, err = f.Write(b) if err != nil { t.RecordMessage("error writing json object dump: %s", err) } } // 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") t.DumpJSON("test-parameters.json", t.TestInstanceParams) return f(t) } }