From 2bb5bea49c787eeef4d3cb4839ba730d56eff102 Mon Sep 17 00:00:00 2001
From: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Date: Wed, 12 Feb 2025 22:09:39 +0000
Subject: [PATCH] Fix unit test

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
---
 pkg/runtime/framework/core/framework_test.go    | 7 +++++--
 pkg/runtime/framework/plugins/jobset/builder.go | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/pkg/runtime/framework/core/framework_test.go b/pkg/runtime/framework/core/framework_test.go
index a20b030407..d616025c07 100644
--- a/pkg/runtime/framework/core/framework_test.go
+++ b/pkg/runtime/framework/core/framework_test.go
@@ -86,10 +86,12 @@ func TestNew(t *testing.T) {
 				watchExtensionPlugins: []framework.WatchExtensionPlugin{
 					&coscheduling.CoScheduling{},
 					&jobset.JobSet{},
+					&mpi.MPI{},
 				},
 				componentBuilderPlugins: []framework.ComponentBuilderPlugin{
 					&coscheduling.CoScheduling{},
 					&jobset.JobSet{},
+					&mpi.MPI{},
 				},
 				terminalConditionPlugins: []framework.TerminalConditionPlugin{
 					&jobset.JobSet{},
@@ -528,11 +530,12 @@ func TestWatchExtensionPlugins(t *testing.T) {
 		registry    fwkplugins.Registry
 		wantPlugins []framework.WatchExtensionPlugin
 	}{
-		"coscheding and jobset are performed": {
+		"coscheding, jobset, and mpi are performed": {
 			registry: fwkplugins.NewRegistry(),
 			wantPlugins: []framework.WatchExtensionPlugin{
 				&coscheduling.CoScheduling{},
 				&jobset.JobSet{},
+				&mpi.MPI{},
 			},
 		},
 		"an empty registry": {
@@ -541,7 +544,7 @@ func TestWatchExtensionPlugins(t *testing.T) {
 	}
 	cmpOpts := []cmp.Option{
 		cmpopts.SortSlices(func(a, b framework.Plugin) bool { return a.Name() < b.Name() }),
-		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, jobset.JobSet{}),
+		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, jobset.JobSet{}, mpi.MPI{}),
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
diff --git a/pkg/runtime/framework/plugins/jobset/builder.go b/pkg/runtime/framework/plugins/jobset/builder.go
index 8cc8bc10da..0d4d29d525 100644
--- a/pkg/runtime/framework/plugins/jobset/builder.go
+++ b/pkg/runtime/framework/plugins/jobset/builder.go
@@ -137,7 +137,7 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
 }
 
 // Launcher updates JobSet values for the launcher Job.
-func (b *Builder) Launcher(info *runtime.Info, trainJob *kubeflowv2.TrainJob) *Builder {
+func (b *Builder) Launcher(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
 		if rJob.Name == constants.JobLauncher {