0d0fc2d433
Co-authored-by: Anton Evangelatov <anton.evangelatov@gmail.com> Co-authored-by: Raúl Kripalani <raul@protocol.ai>
89 lines
2.2 KiB
Go
89 lines
2.2 KiB
Go
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)
|
|
}
|
|
}
|