fix unit test + linting

Signed-off-by: punker <[email protected]>
flyteorg · Jan 30, 2025 · 644ce84 · 644ce84
1 parent 633aa57
commit 644ce84
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
@@ -746,7 +746,7 @@ func TestReplicaCounts(t *testing.T) {
 		contains           []commonOp.ReplicaType
 		notContains        []commonOp.ReplicaType
 	}{
-		{"NoWorkers", 0, true, nil, nil},
+		{"NoWorkers", 0, false, []commonOp.ReplicaType{kubeflowv1.PyTorchJobReplicaTypeMaster}, []commonOp.ReplicaType{}},
 		{"Works", 1, false, []commonOp.ReplicaType{kubeflowv1.PyTorchJobReplicaTypeMaster, kubeflowv1.PyTorchJobReplicaTypeWorker}, []commonOp.ReplicaType{}},
 	} {
 		t.Run(test.name, func(t *testing.T) {
@@ -1266,7 +1266,7 @@ func TestBuildResourcePytorchV1WithDifferentWorkersNumber(t *testing.T) {
 			pytorchJob, ok := res.(*kubeflowv1.PyTorchJob)
 			assert.True(t, ok)
 
-			if taskConfig.WorkerReplicas.Replicas == 0 {
+			if taskConfig.GetWorkerReplicas().GetReplicas() == 0 {
 				// Should only contain master spec
 				assert.Equal(t, 1, len(pytorchJob.Spec.PyTorchReplicaSpecs))
 				assert.Contains(t, pytorchJob.Spec.PyTorchReplicaSpecs, kubeflowv1.PyTorchJobReplicaTypeMaster)