From c7be349a2bf0f25c1365fd516b5a986f14b26cfb Mon Sep 17 00:00:00 2001
From: Kevin Su <pingsutw@apache.org>
Date: Sat, 12 Aug 2023 12:24:02 -0700
Subject: [PATCH] fix tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 .../k8s/kfoperators/common/common_operator.go | 18 -----------
 .../common/common_operator_test.go            | 32 ++++++-------------
 go/tasks/plugins/k8s/kfoperators/mpi/mpi.go   |  3 +-
 3 files changed, 10 insertions(+), 43 deletions(-)

diff --git a/go/tasks/plugins/k8s/kfoperators/common/common_operator.go b/go/tasks/plugins/k8s/kfoperators/common/common_operator.go
index f4b1224bd..acaf09680 100644
--- a/go/tasks/plugins/k8s/kfoperators/common/common_operator.go
+++ b/go/tasks/plugins/k8s/kfoperators/common/common_operator.go
@@ -31,24 +31,6 @@ type ReplicaEntry struct {
 	RestartPolicy commonOp.RestartPolicy
 }
 
-// ExtractMPICurrentCondition will return the first job condition for MPI
-func ExtractMPICurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.JobCondition, error) {
-	if jobConditions != nil {
-		sort.Slice(jobConditions, func(i, j int) bool {
-			return jobConditions[i].LastTransitionTime.Time.After(jobConditions[j].LastTransitionTime.Time)
-		})
-
-		for _, jc := range jobConditions {
-			if jc.Status == v1.ConditionTrue {
-				return jc, nil
-			}
-		}
-		return commonOp.JobCondition{}, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions)
-	}
-
-	return commonOp.JobCondition{}, nil
-}
-
 // ExtractCurrentCondition will return the first job condition for tensorflow/pytorch
 func ExtractCurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.JobCondition, error) {
 	if jobConditions != nil {
diff --git a/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go b/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go
index 933fbf5f8..dc844b866 100644
--- a/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go
+++ b/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go
@@ -18,7 +18,7 @@ import (
 	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func TestExtractMPICurrentCondition(t *testing.T) {
+func TestExtractCurrentCondition(t *testing.T) {
 	jobCreated := commonOp.JobCondition{
 		Type:   commonOp.JobCreated,
 		Status: corev1.ConditionTrue,
@@ -31,35 +31,21 @@ func TestExtractMPICurrentCondition(t *testing.T) {
 		jobCreated,
 		jobRunningActive,
 	}
-	currentCondition, err := ExtractMPICurrentCondition(jobConditions)
+	currentCondition, err := ExtractCurrentCondition(jobConditions)
 	assert.NoError(t, err)
 	assert.Equal(t, currentCondition, jobCreated)
 
 	jobConditions = nil
-	currentCondition, err = ExtractMPICurrentCondition(jobConditions)
-	assert.Error(t, err)
+	currentCondition, err = ExtractCurrentCondition(jobConditions)
 	assert.Equal(t, currentCondition, commonOp.JobCondition{})
-	assert.Equal(t, err, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions))
-}
 
-func TestExtractCurrentCondition(t *testing.T) {
-	jobCreated := commonOp.JobCondition{
-		Type:   commonOp.JobCreated,
-		Status: corev1.ConditionTrue,
-	}
-	jobRunningActive := commonOp.JobCondition{
-		Type:   commonOp.JobRunning,
-		Status: corev1.ConditionFalse,
-	}
-	jobConditions := []commonOp.JobCondition{
-		jobCreated,
-		jobRunningActive,
-	}
-	currentCondition, err := ExtractCurrentCondition(jobConditions)
-	assert.NoError(t, err)
-	assert.Equal(t, currentCondition, jobCreated)
+	jobCreating := commonOp.JobCondition{}
+	jobConditions = []commonOp.JobCondition{jobCreating}
+	currentCondition, err = ExtractCurrentCondition(jobConditions)
+	assert.Equal(t, currentCondition, commonOp.JobCondition{})
 
-	jobConditions = nil
+	jobUnknown := commonOp.JobCondition{Type: "unknown"}
+	jobConditions = []commonOp.JobCondition{jobUnknown}
 	currentCondition, err = ExtractCurrentCondition(jobConditions)
 	assert.Error(t, err)
 	assert.Equal(t, currentCondition, commonOp.JobCondition{})
diff --git a/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go b/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go
index 92c864340..7123dff52 100644
--- a/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go
+++ b/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go
@@ -209,7 +209,7 @@ func (mpiOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err
 	}
-	currentCondition, err := common.ExtractMPICurrentCondition(app.Status.Conditions)
+	currentCondition, err := common.ExtractCurrentCondition(app.Status.Conditions)
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err
 	}
@@ -223,7 +223,6 @@ func (mpiOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext
 	}
 
 	return common.GetMPIPhaseInfo(currentCondition, occurredAt, taskPhaseInfo)
-
 }
 
 func init() {