fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Sep 15, 2023 · e040a02 · e040a02
1 parent a417872
commit e040a02
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/go/tasks/plugins/hive/execution_state.go b/go/tasks/plugins/hive/execution_state.go
@@ -489,7 +489,7 @@ func Finalize(ctx context.Context, tCtx core.TaskExecutionContext, _ ExecutionSt
 }
 
 func InTerminalState(e ExecutionState) bool {
-	return e.Phase == PhaseQuerySucceeded || e.Phase == PhaseQueryFailed
+	return e.Phase.IsTerminal()
 }
 
 func IsNotYetSubmitted(e ExecutionState) bool {

diff --git a/go/tasks/plugins/hive/execution_state_test.go b/go/tasks/plugins/hive/execution_state_test.go
@@ -52,6 +52,7 @@ func TestInTerminalState(t *testing.T) {
 			e := ExecutionState{Phase: tt.phase}
 			res := InTerminalState(e)
 			assert.Equal(t, tt.isTerminal, res)
+			assert.Equal(t, tt.phase.IsTerminal(), res)
 		})
 	}
 }
@@ -338,7 +339,10 @@ func TestKickOffQuery(t *testing.T) {
 
 	var getOrCreateCalled = false
 	mockCache := &mocks2.AutoRefresh{}
-	mockCache.OnGetOrCreate(mock.Anything, mock.Anything).Run(func(_ mock.Arguments) {
+	mockCache.OnGetOrCreate(mock.Anything,
+		ExecutionStateCacheItem{
+			ExecutionState: ExecutionState{Phase: PhaseSubmitted, CommandID: "453298043"},
+			Identifier:     "my_wf_exec_project:my_wf_exec_domain:my_wf_exec_name"}).Run(func(_ mock.Arguments) {
 		getOrCreateCalled = true
 	}).Return(ExecutionStateCacheItem{}, nil)