based on feedback, make methods consistant and enable context based a…

…ccess.

pkg/environment/magic.go

@@ -23,7 +23,9 @@ import (
 	"testing"
 
 	corev1 "k8s.io/api/core/v1"
+
 	"knative.dev/reconciler-test/pkg/feature"
+	"knative.dev/reconciler-test/pkg/state"
 )
 
 func NewGlobalEnvironment(ctx context.Context) GlobalEnvironment {
@@ -133,6 +135,11 @@ func (mr *MagicEnvironment) Prerequisite(ctx context.Context, t *testing.T, f *f
 func (mr *MagicEnvironment) Test(ctx context.Context, t *testing.T, f *feature.Feature) {
 	t.Helper() // Helper marks the calling function as a test helper function.
 
+	if f.State == nil {
+		f.State = &state.KVStore{}
+	}
+	ctx = state.ContextWith(ctx, f.State)
+
 	steps := feature.CollapseSteps(f.Steps)
 
 	for _, timing := range feature.Timings() {

pkg/feature/feature.go

@@ -55,22 +55,6 @@ func (s *Step) TestName() string {
 	}
 }
 
-// Save a value for use in a later feature step, or error on issue.
-func (f *Feature) Save(ctx context.Context, t *testing.T, key string, value interface{}) {
-	t.Helper()
-	if err := f.State.Set(ctx, key, value); err != nil {
-		t.Error(err)
-	}
-}
-
-// Load a value stored in a previous feature step, or fail if not found.
-func (f *Feature) Load(ctx context.Context, t *testing.T, key string, value interface{}) {
-	t.Helper()
-	if err := f.State.Get(ctx, key, &value); err != nil {
-		t.Error(err)
-	}
-}
-
 // Setup adds a step function to the feature set at the Setup timing phase.
 func (f *Feature) Setup(name string, fn StepFn) {
 	f.AddStep(Step{
@@ -117,9 +101,6 @@ func (f *Feature) Teardown(name string, fn StepFn) {
 
 // AddStep appends one or more steps to the Feature.
 func (f *Feature) AddStep(step ...Step) {
-	if f.State == nil {
-		f.State = &state.KVStore{}
-	}
 	f.Steps = append(f.Steps, step...)
 }
 

pkg/state/context.go

@@ -0,0 +1,62 @@
+/*
+Copyright 2021 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package state
+
+import (
+	"context"
+	"testing"
+)
+
+type envKey struct{}
+
+func ContextWith(ctx context.Context, s Store) context.Context {
+	return context.WithValue(ctx, envKey{}, s)
+}
+
+func FromContext(ctx context.Context) Store {
+	if e, ok := ctx.Value(envKey{}).(Store); ok {
+		return e
+	}
+	panic("no Store found in context")
+}
+
+// Get the string value from the kvstore from key.
+func GetStringOrFail(ctx context.Context, t *testing.T, key string) string {
+	t.Helper()
+	value := ""
+	state := FromContext(ctx)
+	if err := state.Get(ctx, key, &value); err != nil {
+		t.Error(err)
+	}
+	return value
+}
+
+// Get gets the key from the Store into the provided value
+func GetOrFail(ctx context.Context, t *testing.T, key string, value interface{}) {
+	state := FromContext(ctx)
+	if err := state.Get(ctx, key, value); err != nil {
+		t.Error(err)
+	}
+}
+
+// Set sets the key into the Store from the provided value
+func SetOrFail(ctx context.Context, t *testing.T, key string, value interface{}) {
+	state := FromContext(ctx)
+	if err := state.Set(ctx, key, value); err != nil {
+		t.Error(err)
+	}
+}

test/example/recorder_feature.go

@@ -27,6 +27,7 @@ import (
 	"knative.dev/reconciler-test/pkg/eventshub"
 	"knative.dev/reconciler-test/pkg/feature"
 	"knative.dev/reconciler-test/pkg/k8s"
+	"knative.dev/reconciler-test/pkg/state"
 
 	// Dot import the eventshub asserts and sdk-go test packages to include all the assert utilities
 	. "github.com/cloudevents/sdk-go/v2/test"
@@ -40,41 +41,38 @@ func RecorderFeature() *feature.Feature {
 
 	f.Setup("create an event", func(ctx context.Context, t *testing.T) {
 		t.Helper()
-		f.Save(ctx, t, "event", FullEvent())
+		state.SetOrFail(ctx, t, "event", FullEvent())
 	})
 
 	f.Setup("install recorder", func(ctx context.Context, t *testing.T) {
 		t.Helper()
 		to := feature.MakeRandomK8sName("recorder")
 		eventshub.Install(to, eventshub.StartReceiver)
-		f.Save(ctx, t, "to", to)
+		state.SetOrFail(ctx, t, "to", to)
 	})
 
 	f.Setup("install sender", func(ctx context.Context, t *testing.T) {
 		t.Helper()
-		var to string
+		to := state.GetStringOrFail(ctx, t, "to")
 		var event cloudevents.Event
-		f.Load(ctx, t, "to", &to)
-		f.Load(ctx, t, "event", &event)
+		state.GetOrFail(ctx, t, "event", &event)
 		from := feature.MakeRandomK8sName("sender")
 		eventshub.Install(from, eventshub.StartSender(to), eventshub.InputEvent(event))
 	})
 
 	f.Requirement("recorder is addressable", func(ctx context.Context, t *testing.T) {
 		t.Helper()
-		var to string
-		f.Load(ctx, t, "to", &to)
+		to := state.GetStringOrFail(ctx, t, "to")
 		k8s.IsAddressable(svc, to, time.Second, 30*time.Second)
 	})
 
 	f.Alpha("direct sending between a producer and a recorder").
 		Must("the recorder received all sent events within the time",
 			func(ctx context.Context, t *testing.T) {
 				t.Helper()
-				var to string
+				to := state.GetStringOrFail(ctx, t, "to")
 				var event cloudevents.Event
-				f.Load(ctx, t, "to", &to)
-				f.Load(ctx, t, "event", &event)
+				state.GetOrFail(ctx, t, "event", &event)
 				OnStore(to).MatchEvent(HasId(event.ID())).Exact(1)
 			})