Skip to content

Commit 22d74d7

Browse files
Neelabh94Neelabh94
authored
feat: Support DNS endpoints for GKE clusters (#5990)
1 parent a86a912 commit 22d74d7

10 files changed

Lines changed: 130 additions & 72 deletions

File tree

‎cmd/job/cancel_test.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import (
2424
)
2525

2626
func TestCancelCmd_Success(t *testing.T) {
27-
resetSubmitCmdFlags() // Reset shared flags
27+
setupSubmitTestEnv(t) // Reset shared flags
2828

2929
// Mock the orchestrator factory
3030
oldFactory := gkeOrchestratorFactory
@@ -52,6 +52,9 @@ func TestCancelCmd_Success(t *testing.T) {
5252
type mockCancelExecutor struct{}
5353

5454
func (m *mockCancelExecutor) ExecuteCommand(name string, args ...string) shell.CommandResult {
55+
if name == "gcloud" && len(args) >= 3 && args[0] == "container" && args[1] == "clusters" && args[2] == "describe" {
56+
return shell.CommandResult{ExitCode: 0, Stdout: "{}"}
57+
}
5558
return shell.CommandResult{ExitCode: 0}
5659
}
5760

@@ -88,7 +91,7 @@ func (m *mockKubeClient) GetCurrentNamespace() (string, error) {
8891
}
8992

9093
func TestCancelCmd_MissingArgs(t *testing.T) {
91-
resetSubmitCmdFlags()
94+
setupSubmitTestEnv(t)
9295

9396
_, err := executeCommand(JobCmd, "cancel")
9497
if err == nil {
@@ -101,7 +104,7 @@ func TestCancelCmd_MissingArgs(t *testing.T) {
101104
}
102105

103106
func TestCancelCmd_JobNotFound(t *testing.T) {
104-
resetSubmitCmdFlags()
107+
setupSubmitTestEnv(t)
105108

106109
oldFactory := gkeOrchestratorFactory
107110
defer func() { gkeOrchestratorFactory = oldFactory }()

‎cmd/job/inspect_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ type mockJobOrchestrator struct {
2525
inspectCalled bool
2626
}
2727

28+
func (m *mockJobOrchestrator) Initialize(clusterName, location, projectID string) (string, error) {
29+
return location, nil
30+
}
31+
2832
func (m *mockJobOrchestrator) SubmitJob(job orchestrator.JobDefinition) error { return nil }
2933
func (m *mockJobOrchestrator) ListJobs(opts orchestrator.ListOptions) ([]orchestrator.JobStatus, error) {
3034
return nil, nil

‎cmd/job/job.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,10 +59,17 @@ var JobCmd = &cobra.Command{
5959
if location == "" {
6060
return fmt.Errorf("location is required; please specify it using the --location flag or set a default value using 'gcluster job config set location <value>'")
6161
}
62+
6263
if projectID == "" {
6364
return fmt.Errorf("project ID is required; please specify it using the --project flag or set a default value using 'gcluster job config set project <value>'")
6465
}
6566

67+
resolvedLoc, err := orc.Initialize(clusterName, location, projectID)
68+
if err != nil {
69+
return err
70+
}
71+
location = resolvedLoc
72+
6673
return nil
6774
},
6875
}

‎cmd/job/list_test.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import (
2222
)
2323

2424
func TestListWorkloadsCmd_Success(t *testing.T) {
25-
resetSubmitCmdFlags() // Reset shared flags
25+
setupSubmitTestEnv(t) // Reset shared flags
2626

2727
// Mock the orchestrator factory
2828
oldFactory := gkeOrchestratorFactory
@@ -48,7 +48,17 @@ func TestListWorkloadsCmd_Success(t *testing.T) {
4848
}
4949

5050
func TestListWorkloadsCmd_InvalidStatus(t *testing.T) {
51-
resetSubmitCmdFlags()
51+
setupSubmitTestEnv(t)
52+
53+
oldFactory := gkeOrchestratorFactory
54+
defer func() { gkeOrchestratorFactory = oldFactory }()
55+
56+
gkeOrchestratorFactory = func() orchestrator.JobOrchestrator {
57+
g := gke.NewGKEOrchestrator()
58+
g.SetExecutor(&mockCancelExecutor{}) // Use the mock from cancel_test.go if available
59+
g.SetKubeClient(&mockKubeClient{namespace: "default"})
60+
return g
61+
}
5262

5363
_, err := executeCommand(JobCmd, "list", "--status", "InvalidStatus", "--cluster", "test-cluster", "--location", "us-central1-a", "--project", "test-project")
5464
if err == nil {

‎cmd/job/logs_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ import (
2525
type mockLogsExecutor struct{}
2626

2727
func (m *mockLogsExecutor) ExecuteCommand(name string, args ...string) shell.CommandResult {
28+
if name == "gcloud" && len(args) >= 3 && args[0] == "container" && args[1] == "clusters" && args[2] == "describe" {
29+
return shell.CommandResult{ExitCode: 0, Stdout: "{}"}
30+
}
2831
if name == "kubectl" && len(args) > 0 && args[0] == "logs" {
2932
return shell.CommandResult{ExitCode: 0, Stdout: "mock logs output"}
3033
}
@@ -36,7 +39,7 @@ func (m *mockLogsExecutor) ExecuteCommandStream(name string, args ...string) err
3639
}
3740

3841
func TestLogsCmd_Success(t *testing.T) {
39-
resetSubmitCmdFlags() // Reset shared flags
42+
setupSubmitTestEnv(t) // Reset shared flags
4043

4144
// Mock the orchestrator factory
4245
oldFactory := gkeOrchestratorFactory

‎cmd/job/submit_test.go‎

Lines changed: 37 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ func TestSubmitCmd_PathwaysDryRun(t *testing.T) {
6868
}
6969

7070
// Reset flags before each test
71-
resetSubmitCmdFlags()
71+
setupSubmitTestEnv(t)
7272

7373
output, err := executeCommand(JobCmd,
7474
"submit",
@@ -143,7 +143,7 @@ func TestSubmitCmd_RegularDryRun(t *testing.T) {
143143
}
144144

145145
// Reset flags before each test
146-
resetSubmitCmdFlags()
146+
setupSubmitTestEnv(t)
147147

148148
output, err := executeCommand(JobCmd,
149149
"submit",
@@ -181,7 +181,7 @@ func TestSubmitCmd_RegularDryRun(t *testing.T) {
181181
}
182182

183183
func TestSubmitCmd_TPUWithNumNodes_Fails(t *testing.T) {
184-
resetSubmitCmdFlags()
184+
setupSubmitTestEnv(t)
185185

186186
oldStore := store
187187
defer func() { store = oldStore }()
@@ -221,7 +221,7 @@ func TestSubmitCmd_TPUWithNumNodes_Fails(t *testing.T) {
221221
}
222222

223223
func TestSubmitCmd_LongWorkloadName_Fails(t *testing.T) {
224-
resetSubmitCmdFlags()
224+
setupSubmitTestEnv(t)
225225

226226
oldStore := store
227227
defer func() { store = oldStore }()
@@ -252,7 +252,7 @@ func TestSubmitCmd_LongWorkloadName_Fails(t *testing.T) {
252252
}
253253
}
254254

255-
func resetSubmitCmdFlags() {
255+
func setupSubmitTestEnv(t *testing.T) {
256256
imageName = ""
257257
baseImage = ""
258258
buildContext = ""
@@ -288,10 +288,26 @@ func resetSubmitCmdFlags() {
288288
pathwaysProxyEnv = nil
289289
pathwaysServerEnv = nil
290290
pathwaysWorkerEnv = nil
291+
292+
oldFactory := gkeOrchestratorFactory
293+
defer t.Cleanup(func() {
294+
gkeOrchestratorFactory = oldFactory
295+
})
296+
gkeOrchestratorFactory = func() orchestrator.JobOrchestrator {
297+
return &mockOrchestrator{}
298+
}
291299
}
292300

293301
type mockOrchestrator struct {
294302
orchestrator.JobOrchestrator
303+
initializeFunc func(string, string, string) (string, error)
304+
}
305+
306+
func (m *mockOrchestrator) Initialize(clusterName, location, projectID string) (string, error) {
307+
if m.initializeFunc != nil {
308+
return m.initializeFunc(clusterName, location, projectID)
309+
}
310+
return location, nil
295311
}
296312

297313
func (m *mockOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
@@ -348,7 +364,7 @@ func TestParseDurationToSeconds(t *testing.T) {
348364
}
349365

350366
func TestSubmitCmd_MissingRepoEnvVar(t *testing.T) {
351-
resetSubmitCmdFlags()
367+
setupSubmitTestEnv(t)
352368

353369
origRepo := os.Getenv("GCLUSTER_IMAGE_REPO")
354370
os.Setenv("GCLUSTER_IMAGE_REPO", "")
@@ -386,7 +402,7 @@ func TestSubmitCmd_MissingRepoEnvVar(t *testing.T) {
386402
}
387403

388404
func TestSubmitCmd_MissingUserEnvVar(t *testing.T) {
389-
resetSubmitCmdFlags()
405+
setupSubmitTestEnv(t)
390406

391407
origUser := os.Getenv("USER")
392408
origUsername := os.Getenv("USERNAME")
@@ -431,7 +447,7 @@ func TestSubmitCmd_MissingUserEnvVar(t *testing.T) {
431447
}
432448

433449
func TestSubmitCmd_InvalidGKENAPProvisioning(t *testing.T) {
434-
resetSubmitCmdFlags()
450+
setupSubmitTestEnv(t)
435451

436452
oldStore := store
437453
defer func() { store = oldStore }()
@@ -471,7 +487,7 @@ func TestSubmitCmd_InvalidGKENAPProvisioning(t *testing.T) {
471487
}
472488

473489
func TestSubmitCmd_ReservationModelWithoutName_Fails(t *testing.T) {
474-
resetSubmitCmdFlags()
490+
setupSubmitTestEnv(t)
475491

476492
oldStore := store
477493
defer func() { store = oldStore }()
@@ -511,7 +527,7 @@ func TestSubmitCmd_ReservationModelWithoutName_Fails(t *testing.T) {
511527
}
512528

513529
func TestSubmitCmd_NonReservationModelWithName_Fails(t *testing.T) {
514-
resetSubmitCmdFlags()
530+
setupSubmitTestEnv(t)
515531

516532
oldStore := store
517533
defer func() { store = oldStore }()
@@ -574,7 +590,7 @@ func TestSubmitCmd_DryRunMissingDir_Approved(t *testing.T) {
574590
defer func() { shell.PromptYesNo = oldPrompt }()
575591
shell.PromptYesNo = func(prompt string) bool { return true }
576592

577-
resetSubmitCmdFlags()
593+
setupSubmitTestEnv(t)
578594

579595
_, err = executeCommand(JobCmd,
580596
"submit",
@@ -624,7 +640,7 @@ func TestSubmitCmd_DryRunMissingDir_Rejected(t *testing.T) {
624640
defer func() { shell.PromptYesNo = oldPrompt }()
625641
shell.PromptYesNo = func(prompt string) bool { return false }
626642

627-
resetSubmitCmdFlags()
643+
setupSubmitTestEnv(t)
628644

629645
output, err := executeCommand(JobCmd,
630646
"submit",
@@ -663,7 +679,7 @@ func TestSubmitCmd_DryRunIsDir_Existing(t *testing.T) {
663679
defer func() { store = oldStore }()
664680
store = &MockPrereqStore{State: PrereqState{LastCheckedTimestamp: time.Now()}}
665681

666-
resetSubmitCmdFlags()
682+
setupSubmitTestEnv(t)
667683

668684
_, err = executeCommand(JobCmd,
669685
"submit",
@@ -692,7 +708,7 @@ func TestSubmitCmd_DryRunIsDir_TrailingSlash(t *testing.T) {
692708
defer func() { store = oldStore }()
693709
store = &MockPrereqStore{State: PrereqState{LastCheckedTimestamp: time.Now()}}
694710

695-
resetSubmitCmdFlags()
711+
setupSubmitTestEnv(t)
696712

697713
_, err := executeCommand(JobCmd,
698714
"submit",
@@ -736,7 +752,7 @@ func TestSubmitCmd_ValidEnvVars(t *testing.T) {
736752
return &mockOrchestrator{}
737753
}
738754

739-
resetSubmitCmdFlags()
755+
setupSubmitTestEnv(t)
740756

741757
_, err := executeCommand(JobCmd,
742758
"submit",
@@ -770,7 +786,7 @@ func TestSubmitCmd_InvalidEnvFormat_Fails(t *testing.T) {
770786
},
771787
}
772788

773-
resetSubmitCmdFlags()
789+
setupSubmitTestEnv(t)
774790

775791
_, err := executeCommand(JobCmd,
776792
"submit",
@@ -815,7 +831,7 @@ func TestSubmitCmd_PathwaysEnv_Success(t *testing.T) {
815831
return &mockOrchestrator{}
816832
}
817833

818-
resetSubmitCmdFlags()
834+
setupSubmitTestEnv(t)
819835

820836
_, err := executeCommand(JobCmd,
821837
"submit",
@@ -859,7 +875,7 @@ func TestSubmitCmd_PathwaysEnv_InvalidFormat_Fails(t *testing.T) {
859875
return &mockOrchestrator{}
860876
}
861877

862-
resetSubmitCmdFlags()
878+
setupSubmitTestEnv(t)
863879

864880
_, err := executeCommand(JobCmd,
865881
"submit",
@@ -923,7 +939,7 @@ func TestSubmitCmd_InvalidEnvKey_Fails(t *testing.T) {
923939
},
924940
}
925941

926-
resetSubmitCmdFlags()
942+
setupSubmitTestEnv(t)
927943

928944
_, err := executeCommand(JobCmd,
929945
"submit",
@@ -971,7 +987,7 @@ func TestSubmitCmd_PathwaysMTCFlags(t *testing.T) {
971987
return &mockOrchestrator{}
972988
}
973989

974-
resetSubmitCmdFlags()
990+
setupSubmitTestEnv(t)
975991

976992
_, err := executeCommand(JobCmd,
977993
"submit",
@@ -1025,7 +1041,7 @@ func TestSubmitCmd_PathwaysHeadless(t *testing.T) {
10251041
return &mockOrchestrator{}
10261042
}
10271043

1028-
resetSubmitCmdFlags()
1044+
setupSubmitTestEnv(t)
10291045

10301046
_, err := executeCommand(JobCmd,
10311047
"submit",

0 commit comments

Comments
 (0)