Fix scheduler tests

Signed-off-by: Eduardo Apolinario <[email protected]>
flyteorg · Feb 27, 2025 · 8651211 · 8651211
1 parent 098c2b1
commit 8651211
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/flyteadmin/scheduler/core/gocron_scheduler_test.go b/flyteadmin/scheduler/core/gocron_scheduler_test.go
@@ -114,7 +114,7 @@ func setupWithSchedules(t *testing.T, subscope string, schedules []models.Schedu
 	rateLimiter := rate.NewLimiter(1, 10)
 	executor := new(mocks.Executor)
 	snapshot := &snapshoter.SnapshotV1{}
-	executor.OnExecuteMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	executor.EXPECT().Execute(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	g := NewGoCronScheduler(context.Background(), schedules, schedulerScope, snapshot, rateLimiter, executor, useUtcTz)
 	goCronScheduler, ok := g.(*GoCronScheduler)
 	goCronScheduler.UpdateSchedules(context.Background(), schedules)
@@ -186,7 +186,7 @@ func TestGetTimedFuncWithSchedule(t *testing.T) {
 		ctx := context.Background()
 		g := setup(t, "failure_case", false)
 		executor := new(mocks.Executor)
-		executor.OnExecuteMatch(mock.Anything, mock.Anything, mock.Anything).Return(fmt.Errorf("failure case"))
+		executor.EXPECT().Execute(mock.Anything, mock.Anything, mock.Anything).Return(fmt.Errorf("failure case"))
 		g.executor = executor
 		timeFunc := g.GetTimedFuncWithSchedule()
 		assert.NotNil(t, timeFunc)

diff --git a/flyteadmin/scheduler/schedule_executor_test.go b/flyteadmin/scheduler/schedule_executor_test.go
@@ -88,7 +88,7 @@ func TestSuccessfulSchedulerExec(t *testing.T) {
 			Active:              &activeV2,
 		})
 
-		scheduleEntitiesRepo.OnGetAllMatch(mock.Anything).Return(schedules, nil)
+		scheduleEntitiesRepo.EXPECT().GetAll(mock.Anything).Return(schedules, nil)
 		go func() {
 			err := scheduleExecutor.Run(context.Background())
 			assert.Nil(t, err)
@@ -113,7 +113,7 @@ func TestSuccessfulSchedulerExec(t *testing.T) {
 			KickoffTimeInputArg: "kickoff_time",
 			Active:              &activeV2,
 		})
-		scheduleEntitiesRepo.OnGetAllMatch(mock.Anything).Return(schedules, nil)
+		scheduleEntitiesRepo.EXPECT().GetAll(mock.Anything).Return(schedules, nil)
 		time.Sleep(30 * time.Second)
 	})
 
@@ -139,7 +139,7 @@ func TestSuccessfulSchedulerExec(t *testing.T) {
 			KickoffTimeInputArg: "kickoff_time",
 			Active:              &activeV2,
 		})
-		scheduleEntitiesRepo.OnGetAllMatch(mock.Anything).Return(schedules, nil)
+		scheduleEntitiesRepo.EXPECT().GetAll(mock.Anything).Return(schedules, nil)
 
 		go func() {
 			err := scheduleExecutor.Run(context.Background())
@@ -166,14 +166,14 @@ func TestSuccessfulSchedulerExec(t *testing.T) {
 			KickoffTimeInputArg: "kickoff_time",
 			Active:              &activeV2,
 		})
-		scheduleEntitiesRepo.OnGetAllMatch(mock.Anything).Return(schedules, nil)
+		scheduleEntitiesRepo.EXPECT().GetAll(mock.Anything).Return(schedules, nil)
 		time.Sleep(30 * time.Second)
 	})
 
 	t.Run("unable to read schedules", func(t *testing.T) {
 		scheduleExecutor := setupScheduleExecutor(t, "unable_read_schedules")
 		scheduleEntitiesRepo := db.SchedulableEntityRepo().(*schedMocks.SchedulableEntityRepoInterface)
-		scheduleEntitiesRepo.OnGetAllMatch(mock.Anything).Return(nil, fmt.Errorf("unable to read schedules"))
+		scheduleEntitiesRepo.EXPECT().GetAll(mock.Anything).Return(nil, fmt.Errorf("unable to read schedules"))
 
 		go func() {
 			err := scheduleExecutor.Run(context.Background())

diff --git a/flyteadmin/tests/scheduler_test.go b/flyteadmin/tests/scheduler_test.go
@@ -50,7 +50,7 @@ func TestScheduleJob(t *testing.T) {
 	rateLimiter := rate.NewLimiter(1, 10)
 	executor := new(mocks.Executor)
 	snapshot := &snapshoter.SnapshotV1{}
-	executor.OnExecuteMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	executor.EXPECT().Execute(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	g := scheduler.NewGoCronScheduler(context.Background(), []models.SchedulableEntity{}, schedulerScope, snapshot, rateLimiter, executor, false)
 	c.Start()