From e54d3cd79e4f7dfbecaf9e19d0717cd60e20d39c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Mon, 30 Oct 2023 16:44:19 +0100
Subject: [PATCH] Add unit tests for runner recovery.

---
 .../environment/nomad_environment_test.go     | 27 +++++++
 internal/environment/nomad_manager_test.go    | 81 +++++++++++++++++++
 internal/runner/nomad_manager_test.go         |  2 +-
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/internal/environment/nomad_environment_test.go b/internal/environment/nomad_environment_test.go
index d3b61774..ae43d080 100644
--- a/internal/environment/nomad_environment_test.go
+++ b/internal/environment/nomad_environment_test.go
@@ -235,3 +235,30 @@ func (s *MainTestSuite) TestNomadEnvironment_DeleteLocally() {
 	s.NoError(err)
 	apiMock.AssertExpectations(s.T())
 }
+
+func (s *MainTestSuite) TestNomadEnvironment_AddRunner() {
+	s.Run("Destroys runner before replacing it", func() {
+		apiMock := &nomad.ExecutorAPIMock{}
+		environment, err := NewNomadEnvironment(tests.DefaultEnvironmentIDAsInteger, apiMock, templateEnvironmentJobHCL)
+		s.Require().NoError(err)
+		r := &runner.RunnerMock{}
+		r.On("ID").Return(tests.DefaultRunnerID)
+		r.On("Destroy", mock.Anything).Run(func(args mock.Arguments) {
+			err, ok := args[0].(error)
+			s.Require().True(ok)
+			s.ErrorIs(err, runner.ErrLocalDestruction)
+		}).Return(nil).Once()
+		r2 := &runner.RunnerMock{}
+		r2.On("ID").Return(tests.DefaultRunnerID)
+
+		environment.AddRunner(r)
+		environment.AddRunner(r2)
+		r.AssertExpectations(s.T())
+
+		// Teardown test case
+		r2.On("Destroy", mock.Anything).Return(nil)
+		apiMock.On("LoadRunnerIDs", mock.Anything).Return([]string{}, nil)
+		apiMock.On("DeleteJob", mock.Anything).Return(nil)
+		s.NoError(environment.Delete(tests.ErrCleanupDestroyReason))
+	})
+}
diff --git a/internal/environment/nomad_manager_test.go b/internal/environment/nomad_manager_test.go
index 375ebdf8..aefcd5e4 100644
--- a/internal/environment/nomad_manager_test.go
+++ b/internal/environment/nomad_manager_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/stretchr/testify/suite"
 	"os"
 	"testing"
+	"time"
 )
 
 type CreateOrUpdateTestSuite struct {
@@ -315,6 +316,26 @@ func (s *MainTestSuite) TestNomadEnvironmentManager_Load() {
 
 	runnerManager := runner.NewNomadRunnerManager(apiMock, s.TestCtx)
 
+	s.Run("deletes local environments before loading Nomad environments", func() {
+		call.Return([]*nomadApi.Job{}, nil)
+		environment := &runner.ExecutionEnvironmentMock{}
+		environment.On("ID").Return(dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger))
+		environment.On("Image").Return("")
+		environment.On("CPULimit").Return(uint(0))
+		environment.On("MemoryLimit").Return(uint(0))
+		environment.On("NetworkAccess").Return(false, nil)
+		environment.On("Delete", mock.Anything).Return(nil)
+		runnerManager.StoreEnvironment(environment)
+
+		m, err := NewNomadEnvironmentManager(runnerManager, apiMock, "")
+		s.Require().NoError(err)
+
+		err = m.load()
+		s.Require().NoError(err)
+		environment.AssertExpectations(s.T())
+	})
+	runnerManager.DeleteEnvironment(tests.DefaultEnvironmentIDAsInteger)
+
 	s.Run("Stores fetched environments", func() {
 		_, job := helpers.CreateTemplateJob()
 		call.Return([]*nomadApi.Job{job}, nil)
@@ -354,6 +375,66 @@ func (s *MainTestSuite) TestNomadEnvironmentManager_Load() {
 	})
 }
 
+func (s *MainTestSuite) TestNomadEnvironmentManager_KeepEnvironmentsSynced() {
+	apiMock := &nomad.ExecutorAPIMock{}
+	runnerManager := runner.NewNomadRunnerManager(apiMock, s.TestCtx)
+	m, err := NewNomadEnvironmentManager(runnerManager, apiMock, "")
+	s.Require().NoError(err)
+
+	s.Run("stops when context is done", func() {
+		apiMock.On("LoadEnvironmentJobs").Return([]*nomadApi.Job{}, context.DeadlineExceeded)
+		ctx, cancel := context.WithCancel(s.TestCtx)
+		cancel()
+
+		var done bool
+		go func() {
+			<-time.After(tests.ShortTimeout)
+			if !done {
+				s.FailNow("KeepEnvironmentsSynced is ignoring the context")
+			}
+		}()
+
+		m.KeepEnvironmentsSynced(func(_ context.Context) error { return nil }, ctx)
+		done = true
+	})
+	apiMock.ExpectedCalls = []*mock.Call{}
+	apiMock.Calls = []mock.Call{}
+
+	s.Run("retries loading environments", func() {
+		ctx, cancel := context.WithCancel(s.TestCtx)
+
+		apiMock.On("LoadEnvironmentJobs").Return([]*nomadApi.Job{}, context.DeadlineExceeded).Once()
+		apiMock.On("LoadEnvironmentJobs").Return([]*nomadApi.Job{}, nil).Run(func(_ mock.Arguments) {
+			cancel()
+		}).Once()
+
+		m.KeepEnvironmentsSynced(func(_ context.Context) error { return nil }, ctx)
+		apiMock.AssertExpectations(s.T())
+	})
+	apiMock.ExpectedCalls = []*mock.Call{}
+	apiMock.Calls = []mock.Call{}
+
+	s.Run("retries synchronizing runners", func() {
+		apiMock.On("LoadEnvironmentJobs").Return([]*nomadApi.Job{}, nil)
+		ctx, cancel := context.WithCancel(s.TestCtx)
+
+		count := 0
+		synchronizeRunners := func(ctx context.Context) error {
+			count++
+			if count >= 2 {
+				cancel()
+				return nil
+			}
+			return context.DeadlineExceeded
+		}
+		m.KeepEnvironmentsSynced(synchronizeRunners, ctx)
+
+		if count < 2 {
+			s.Fail("KeepEnvironmentsSynced is not retrying to synchronize the runners")
+		}
+	})
+}
+
 func mockWatchAllocations(ctx context.Context, apiMock *nomad.ExecutorAPIMock) {
 	call := apiMock.On("WatchEventStream", mock.Anything, mock.Anything, mock.Anything)
 	call.Run(func(args mock.Arguments) {
diff --git a/internal/runner/nomad_manager_test.go b/internal/runner/nomad_manager_test.go
index 30de4a98..502fc40f 100644
--- a/internal/runner/nomad_manager_test.go
+++ b/internal/runner/nomad_manager_test.go
@@ -267,7 +267,7 @@ func (s *ManagerTestSuite) TestUpdateRunnersLogsErrorFromWatchAllocation() {
 	}()
 	<-time.After(10 * time.Millisecond)
 
-	s.Require().Equal(1, len(hook.Entries))
+	s.Require().Equal(3, len(hook.Entries))
 	s.Equal(logrus.ErrorLevel, hook.LastEntry().Level)
 	err, ok := hook.LastEntry().Data[logrus.ErrorKey].(error)
 	s.Require().True(ok)