Improve BaseTestSuite: move global variable ctx into suite

Use BaseTestSuite in TestRunner tests

Signed-off-by: Ilya <[email protected]>

pkg/runner/base_test_suite_test.go

@@ -1,38 +1,81 @@
 package runner
 
 import (
+	"context"
 	"encoding/json"
 
 	"github.com/benbjohnson/clock"
 	"github.com/linuxboot/contest/pkg/event/testevent"
 	"github.com/linuxboot/contest/pkg/pluginregistry"
+	"github.com/linuxboot/contest/pkg/storage"
 	"github.com/linuxboot/contest/pkg/target"
 	"github.com/linuxboot/contest/pkg/test"
 	"github.com/linuxboot/contest/pkg/xcontext"
+	"github.com/linuxboot/contest/pkg/xcontext/bundles/logrusctx"
+	"github.com/linuxboot/contest/pkg/xcontext/logger"
+	"github.com/linuxboot/contest/plugins/storage/memory"
 	"github.com/linuxboot/contest/plugins/targetlocker/inmemory"
+	"github.com/linuxboot/contest/tests/common"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
 
+type MemoryStorageEngine struct {
+	Ctx                xcontext.Context
+	Storage            storage.ResettableStorage
+	StorageEngineVault *storage.SimpleEngineVault
+}
+
+func NewMemoryStorageEngine() (*MemoryStorageEngine, error) {
+	ms, err := memory.New()
+	if err != nil {
+		return nil, err
+	}
+
+	storageEngineVault := storage.NewSimpleEngineVault()
+	if err := storageEngineVault.StoreEngine(ms, storage.SyncEngine); err != nil {
+		return nil, err
+	}
+
+	return &MemoryStorageEngine{
+		Storage:            ms,
+		StorageEngineVault: storageEngineVault,
+	}, nil
+}
+
+func (mse *MemoryStorageEngine) GetStepEvents(ctx xcontext.Context, testName string, stepLabel string) string {
+	return common.GetTestEventsAsString(mse.Ctx, mse.Storage, testName, nil, &stepLabel)
+}
+
+func (mse *MemoryStorageEngine) GetTargetEvents(ctx xcontext.Context, testName string, targetID string) string {
+	return common.GetTestEventsAsString(mse.Ctx, mse.Storage, testName, &targetID, nil)
+}
+
 type BaseTestSuite struct {
 	suite.Suite
 
-	pluginRegistry  *pluginregistry.PluginRegistry
-	internalStorage *MemoryStorageEngine
+	Ctx    xcontext.Context
+	Cancel context.CancelFunc
+
+	PluginRegistry *pluginregistry.PluginRegistry
+	MemoryStorage  *MemoryStorageEngine
 }
 
 func (s *BaseTestSuite) SetupTest() {
+	s.Ctx, s.Cancel = xcontext.WithCancel(logrusctx.NewContext(logger.LevelDebug))
+
 	storageEngine, err := NewMemoryStorageEngine()
 	require.NoError(s.T(), err)
-	s.internalStorage = storageEngine
+	s.MemoryStorage = storageEngine
 
 	target.SetLocker(inmemory.New(clock.New()))
 
-	s.pluginRegistry = pluginregistry.NewPluginRegistry(xcontext.Background())
+	s.PluginRegistry = pluginregistry.NewPluginRegistry(s.Ctx)
 }
 
 func (s *BaseTestSuite) TearDownTest() {
 	target.SetLocker(nil)
+	s.Cancel()
 }
 
 func (s *BaseTestSuite) RegisterStateFullStep(
@@ -41,7 +84,7 @@ func (s *BaseTestSuite) RegisterStateFullStep(
 		ev testevent.Emitter, resumeState json.RawMessage) (json.RawMessage, error),
 	validateFunction func(ctx xcontext.Context, params test.TestStepParameters) error) error {
 
-	return s.pluginRegistry.RegisterTestStep(stateFullStepName, func() test.TestStep {
+	return s.PluginRegistry.RegisterTestStep(stateFullStepName, func() test.TestStep {
 		return &stateFullStep{
 			runFunction:      runFunction,
 			validateFunction: validateFunction,
@@ -55,7 +98,7 @@ func (s *BaseTestSuite) NewStep(label, name string, params test.TestStepParamete
 		Label:      label,
 		Parameters: params,
 	}
-	sb, err := s.pluginRegistry.NewTestStepBundle(ctx, td)
+	sb, err := s.PluginRegistry.NewTestStepBundle(s.Ctx, td)
 	require.NoError(s.T(), err)
 	return *sb
 }

pkg/runner/job_runner_test.go

@@ -99,7 +99,7 @@ func (s *JobRunnerSuite) SetupTest() {
 	}{
 		{echo.Name, echo.New, echo.Events},
 	} {
-		require.NoError(s.T(), s.pluginRegistry.RegisterTestStep(e.name, e.factory, e.events))
+		require.NoError(s.T(), s.PluginRegistry.RegisterTestStep(e.name, e.factory, e.events))
 	}
 }
 
@@ -147,11 +147,11 @@ func (s *JobRunnerSuite) TestSimpleJobStartFinish() {
 		},
 	}
 
-	jsm := storage.NewJobStorageManager(s.internalStorage.StorageEngineVault)
-	jr := NewJobRunner(jsm, s.internalStorage.StorageEngineVault, clock.New(), time.Second)
+	jsm := storage.NewJobStorageManager(s.MemoryStorage.StorageEngineVault)
+	jr := NewJobRunner(jsm, s.MemoryStorage.StorageEngineVault, clock.New(), time.Second)
 	require.NotNil(s.T(), jr)
 
-	resumeState, err := jr.Run(ctx, &j, nil)
+	resumeState, err := jr.Run(s.Ctx, &j, nil)
 	require.NoError(s.T(), err)
 	require.Nil(s.T(), resumeState)
 
@@ -162,7 +162,7 @@ func (s *JobRunnerSuite) TestSimpleJobStartFinish() {
 {[1 1 SimpleTest 0 test_step_label][Target{ID: "T1"} TargetIn]}
 {[1 1 SimpleTest 0 test_step_label][Target{ID: "T1"} TargetOut]}
 {[1 1 SimpleTest 0 ][Target{ID: "T1"} TargetReleased]}
-`, s.internalStorage.GetTargetEvents(testName, "T1"))
+`, s.MemoryStorage.GetTargetEvents(s.Ctx, testName, "T1"))
 }
 
 func (s *JobRunnerSuite) TestJobWithTestRetry() {
@@ -232,11 +232,11 @@ func (s *JobRunnerSuite) TestJobWithTestRetry() {
 		},
 	}
 
-	jsm := storage.NewJobStorageManager(s.internalStorage.StorageEngineVault)
-	jr := NewJobRunner(jsm, s.internalStorage.StorageEngineVault, clock.New(), time.Second)
+	jsm := storage.NewJobStorageManager(s.MemoryStorage.StorageEngineVault)
+	jr := NewJobRunner(jsm, s.MemoryStorage.StorageEngineVault, clock.New(), time.Second)
 	require.NotNil(s.T(), jr)
 
-	resumeState, err := jr.Run(ctx, &j, nil)
+	resumeState, err := jr.Run(s.Ctx, &j, nil)
 	require.NoError(s.T(), err)
 	require.Nil(s.T(), resumeState)
 
@@ -255,7 +255,7 @@ func (s *JobRunnerSuite) TestJobWithTestRetry() {
 {[1 1 SimpleTest 1 echo2_step_label][Target{ID: "T1"} TargetIn]}
 {[1 1 SimpleTest 1 echo2_step_label][Target{ID: "T1"} TargetOut]}
 {[1 1 SimpleTest 1 ][Target{ID: "T1"} TargetReleased]}
-`, s.internalStorage.GetTargetEvents(testName, "T1"))
+`, s.MemoryStorage.GetTargetEvents(s.Ctx, testName, "T1"))
 
 	require.Len(s.T(), reporter.runStatuses, 1)
 	require.Len(s.T(), reporter.runStatuses[0].TestStatuses, 1)
@@ -311,8 +311,8 @@ func (s *JobRunnerSuite) TestResumeStateBadJobId() {
 		},
 	}
 
-	jsm := storage.NewJobStorageManager(s.internalStorage.StorageEngineVault)
-	jr := NewJobRunner(jsm, s.internalStorage.StorageEngineVault, clock.New(), time.Second)
+	jsm := storage.NewJobStorageManager(s.MemoryStorage.StorageEngineVault)
+	jr := NewJobRunner(jsm, s.MemoryStorage.StorageEngineVault, clock.New(), time.Second)
 	require.NotNil(s.T(), jr)
 
 	inputResumeState := job.PauseEventPayload{
@@ -322,7 +322,7 @@ func (s *JobRunnerSuite) TestResumeStateBadJobId() {
 		TestID:  1,
 	}
 
-	resumeState, err := jr.Run(ctx, &j, &inputResumeState)
+	resumeState, err := jr.Run(s.Ctx, &j, &inputResumeState)
 	require.Error(s.T(), err)
 	require.Nil(s.T(), resumeState)
 }

pkg/runner/step_runner_test.go

@@ -56,21 +56,21 @@ func (s *StepRunnerSuite) TestRunningStep() {
 	stepRunner := NewStepRunner()
 	require.NotNil(s.T(), stepRunner)
 
-	emitterFactory := NewTestStepEventsEmitterFactory(s.internalStorage.StorageEngineVault, 1, 1, testName, 0)
+	emitterFactory := NewTestStepEventsEmitterFactory(s.MemoryStorage.StorageEngineVault, 1, 1, testName, 0)
 	emitter := emitterFactory.New("test_step_label")
 
 	inputResumeState := json.RawMessage("{\"some_input\": 42}")
-	resultChan, err := stepRunner.Run(ctx, s.NewStep("test_step_label", stateFullStepName, nil), emitter, inputResumeState)
+	resultChan, err := stepRunner.Run(s.Ctx, s.NewStep("test_step_label", stateFullStepName, nil), emitter, inputResumeState)
 	require.NoError(s.T(), err)
 	require.NotNil(s.T(), resultChan)
 
-	require.NoError(s.T(), stepRunner.AddTarget(ctx, tgt("TSucc")))
+	require.NoError(s.T(), stepRunner.AddTarget(s.Ctx, tgt("TSucc")))
 	ev, ok := <-resultChan
 	require.True(s.T(), ok)
 	require.Equal(s.T(), tgt("TSucc"), ev.Target)
 	require.NoError(s.T(), ev.Err)
 
-	require.NoError(s.T(), stepRunner.AddTarget(ctx, tgt("TFail")))
+	require.NoError(s.T(), stepRunner.AddTarget(s.Ctx, tgt("TFail")))
 	ev, ok = <-resultChan
 	require.True(s.T(), ok)
 	require.Equal(s.T(), tgt("TFail"), ev.Target)
@@ -86,7 +86,7 @@ func (s *StepRunnerSuite) TestRunningStep() {
 	ev, ok = <-resultChan
 	require.False(s.T(), ok)
 
-	closedCtx, cancel := xcontext.WithCancel(ctx)
+	closedCtx, cancel := xcontext.WithCancel(s.Ctx)
 	cancel()
 
 	// if step runner has results, it should return them even if input context is closed