Skip to content

Commit 025254b

Browse files
committed
feat: improve test coverage for imagebuilder and scheduling, refactor prereq to cmd/job
1 parent 73be875 commit 025254b

File tree

10 files changed

+637
-35
lines changed

10 files changed

+637
-35
lines changed
Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
package prereq
15+
package job
1616

1717
import (
1818
"bufio"
@@ -245,8 +245,6 @@ func ensureKubectlInstalled() error {
245245
// configureDockerCredentialHelper configures Docker to authenticate to Google Container Registry and Artifact Registry.
246246
func configureDockerCredentialHelper() error {
247247
logging.Info("Configuring Docker credential helper for Google Container Registry and Artifact Registry...")
248-
// Note: It's hard to reliably check if this is already configured without parsing Docker's config.json.
249-
// Running it idempotent-ly is generally safe and ensures configuration.
250248
gcrResult := shell.ExecuteCommand("gcloud", "auth", "configure-docker", "gcr.io", "--quiet")
251249
if gcrResult.ExitCode != 0 {
252250
return fmt.Errorf("failed to configure Docker for gcr.io: %s", gcrResult.Stderr)
@@ -255,7 +253,6 @@ func configureDockerCredentialHelper() error {
255253
if usCentral1Result.ExitCode != 0 {
256254
return fmt.Errorf("failed to configure Docker for us-central1-docker.pkg.dev: %s", usCentral1Result.Stderr)
257255
}
258-
// Add other common Artifact Registry hosts if needed, or make it dynamic
259256
logging.Info("Docker credential helper configured for Google registries.")
260257
return nil
261258
}
@@ -316,20 +313,15 @@ func checkGCloudSDK(newState *PrereqState) error {
316313
}
317314

318315
func checkGCloudProject(newState *PrereqState, projectID *string) error {
319-
// This function also updates the passed projectID pointer, which is essential.
320316
if !newState.GCloudProjectConfigured {
321317
if err := ensureGCloudProjectConfigured(projectID); err != nil {
322318
return err
323319
}
324320
newState.GCloudProjectConfigured = true
325321
} else {
326-
// Even if already configured, ensure the `projectID` passed to the function is updated
327-
// if it was empty initially and inferred from the state.
328322
if *projectID == "" && newState.LastCheckedProjectID != "" {
329323
*projectID = newState.LastCheckedProjectID
330324
}
331-
// If projectID was provided as a flag, ensure that the stored one is the current one
332-
// This is also handled inside ensureGCloudProjectConfigured
333325
if err := ensureGCloudProjectConfigured(projectID); err != nil {
334326
return err
335327
}
@@ -379,7 +371,6 @@ func checkDockerCredsConfigured(newState *PrereqState) error {
379371

380372
func checkArtifactRegistryAPIEnabled(newState *PrereqState, projectID *string) error {
381373
if !newState.ArtifactRegistryAPIEnabled {
382-
// This depends on projectID being set, so it must come after ensureGCloudProjectConfigured
383374
if *projectID == "" {
384375
return fmt.Errorf("project ID is not set after configuration, cannot enable Artifact Registry API")
385376
}
@@ -392,23 +383,19 @@ func checkArtifactRegistryAPIEnabled(newState *PrereqState, projectID *string) e
392383
}
393384

394385
// EnsurePrerequisites checks all necessary gcloud and kubectl prerequisites.
395-
// It modifies the projectID pointer if it infers it or gets it from user input.
396386
func EnsurePrerequisites(projectID *string) error {
397387
if os.Getenv("GCLUSTER_SKIP_PREREQ_CHECKS") == "true" {
398388
logging.Info("Skipping prerequisite checks due to GCLUSTER_SKIP_PREREQ_CHECKS environment variable.")
399389
return nil
400390
}
401391

402392
state := loadPrereqState()
403-
newState := state // Copy the loaded state to modify and save later
393+
newState := state
404394

405-
// Determine the current project ID, which might be inferred later if not provided.
406-
// We need this for the staleness check against the stored LastCheckedProjectID.
407395
var actualProjectID string
408396
if *projectID != "" {
409397
actualProjectID = *projectID
410398
} else {
411-
// Attempt to get project from gcloud config if not provided by flag
412399
projectResult := shell.ExecuteCommand("gcloud", "config", "get-value", "project")
413400
if projectResult.ExitCode == 0 {
414401
actualProjectID = strings.TrimSpace(projectResult.Stdout)
@@ -417,12 +404,11 @@ func EnsurePrerequisites(projectID *string) error {
417404

418405
if isStateStale(state, actualProjectID) {
419406
logging.Info("State is stale or project changed, re-running all prerequisite checks.")
420-
newState = PrereqState{} // Reset new state if stale, effectively re-running all checks
407+
newState = PrereqState{}
421408
} else {
422409
logging.Info("Prerequisite state is fresh, skipping already completed checks.")
423410
}
424411

425-
// Define the sequence of checks
426412
checks := []func(*PrereqState, *string) error{
427413
func(ns *PrereqState, pID *string) error { return checkGCloudSDK(ns) },
428414
func(ns *PrereqState, pID *string) error { return checkGCloudProject(ns, pID) },
@@ -439,7 +425,6 @@ func EnsurePrerequisites(projectID *string) error {
439425
}
440426
}
441427

442-
// Update timestamp and project ID before saving
443428
newState.LastCheckedTimestamp = time.Now()
444429
newState.LastCheckedProjectID = *projectID
445430
savePrereqState(newState)

cmd/job/prereq_test.go

Lines changed: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package job
16+
17+
import (
18+
"encoding/json"
19+
"os"
20+
"path/filepath"
21+
"strings"
22+
"testing"
23+
"time"
24+
)
25+
26+
func TestIsStateStale(t *testing.T) {
27+
now := time.Now()
28+
freshTime := now.Add(-1 * time.Hour)
29+
staleTime := now.Add(-48 * time.Hour)
30+
31+
tests := []struct {
32+
name string
33+
state PrereqState
34+
currentProjectID string
35+
wantStale bool
36+
}{
37+
{
38+
name: "Fresh state, same project",
39+
state: PrereqState{
40+
LastCheckedTimestamp: freshTime,
41+
LastCheckedProjectID: "test-project",
42+
},
43+
currentProjectID: "test-project",
44+
wantStale: false,
45+
},
46+
{
47+
name: "Stale state (time), same project",
48+
state: PrereqState{
49+
LastCheckedTimestamp: staleTime,
50+
LastCheckedProjectID: "test-project",
51+
},
52+
currentProjectID: "test-project",
53+
wantStale: true,
54+
},
55+
{
56+
name: "Fresh state, different project",
57+
state: PrereqState{
58+
LastCheckedTimestamp: freshTime,
59+
LastCheckedProjectID: "old-project",
60+
},
61+
currentProjectID: "new-project",
62+
wantStale: true,
63+
},
64+
}
65+
66+
for _, tt := range tests {
67+
t.Run(tt.name, func(t *testing.T) {
68+
got := isStateStale(tt.state, tt.currentProjectID)
69+
if got != tt.wantStale {
70+
t.Errorf("isStateStale() = %v, want %v", got, tt.wantStale)
71+
}
72+
})
73+
}
74+
}
75+
76+
func TestStateFilePath(t *testing.T) {
77+
tempDir, err := os.MkdirTemp("", "prereq-test-home")
78+
if err != nil {
79+
t.Fatal(err)
80+
}
81+
defer os.RemoveAll(tempDir)
82+
83+
origHome := os.Getenv("HOME")
84+
os.Setenv("HOME", tempDir)
85+
defer os.Setenv("HOME", origHome)
86+
87+
path, err := stateFilePath()
88+
if err != nil {
89+
t.Fatalf("stateFilePath() error = %v", err)
90+
}
91+
92+
expectedPrefix := filepath.Join(tempDir, stateDirName)
93+
if !strings.HasPrefix(path, expectedPrefix) {
94+
t.Errorf("expected path to start with %s, got %s", expectedPrefix, path)
95+
}
96+
}
97+
98+
func TestLoadPrereqState_Success(t *testing.T) {
99+
tempDir, err := os.MkdirTemp("", "prereq-load-test")
100+
if err != nil {
101+
t.Fatal(err)
102+
}
103+
defer os.RemoveAll(tempDir)
104+
105+
origHome := os.Getenv("HOME")
106+
os.Setenv("HOME", tempDir)
107+
defer os.Setenv("HOME", origHome)
108+
109+
stateDir := filepath.Join(tempDir, stateDirName)
110+
if err := os.MkdirAll(stateDir, 0755); err != nil {
111+
t.Fatal(err)
112+
}
113+
statePath := filepath.Join(stateDir, stateFileName)
114+
115+
state := PrereqState{
116+
GCloudSDKInstalled: true,
117+
LastCheckedProjectID: "test-project",
118+
LastCheckedTimestamp: time.Now(),
119+
}
120+
data, _ := json.Marshal(state)
121+
if err := os.WriteFile(statePath, data, 0644); err != nil {
122+
t.Fatal(err)
123+
}
124+
125+
loaded := loadPrereqState()
126+
if !loaded.GCloudSDKInstalled {
127+
t.Errorf("expected loaded state to be true for GCloudSDKInstalled, got false")
128+
}
129+
if loaded.LastCheckedProjectID != "test-project" {
130+
t.Errorf("expected loaded state project ID to be 'test-project', got %s", loaded.LastCheckedProjectID)
131+
}
132+
}
133+
134+
func TestSavePrereqState_Success(t *testing.T) {
135+
tempDir, err := os.MkdirTemp("", "prereq-save-test")
136+
if err != nil {
137+
t.Fatal(err)
138+
}
139+
defer os.RemoveAll(tempDir)
140+
141+
origHome := os.Getenv("HOME")
142+
os.Setenv("HOME", tempDir)
143+
defer os.Setenv("HOME", origHome)
144+
145+
state := PrereqState{
146+
GCloudAuthenticated: true,
147+
LastCheckedProjectID: "test-project",
148+
LastCheckedTimestamp: time.Now(),
149+
}
150+
151+
savePrereqState(state)
152+
153+
stateDir := filepath.Join(tempDir, stateDirName)
154+
statePath := filepath.Join(stateDir, stateFileName)
155+
156+
if _, err := os.Stat(statePath); os.IsNotExist(err) {
157+
t.Fatalf("state file was not created at %s", statePath)
158+
}
159+
160+
data, err := os.ReadFile(statePath)
161+
if err != nil {
162+
t.Fatalf("failed to read created state file: %v", err)
163+
}
164+
165+
var savedState PrereqState
166+
if err := json.Unmarshal(data, &savedState); err != nil {
167+
t.Fatalf("failed to unmarshal saved state: %v", err)
168+
}
169+
170+
if !savedState.GCloudAuthenticated {
171+
t.Errorf("expected saved state to be true for GCloudAuthenticated, got false")
172+
}
173+
}
174+
175+
func TestAskForConfirmation_Yes(t *testing.T) {
176+
pipeRead, pipeWrite, err := os.Pipe()
177+
if err != nil {
178+
t.Fatal(err)
179+
}
180+
defer pipeRead.Close()
181+
defer pipeWrite.Close()
182+
183+
origStdin := os.Stdin
184+
os.Stdin = pipeRead
185+
defer func() { os.Stdin = origStdin }()
186+
187+
if _, err := pipeWrite.Write([]byte("y\n")); err != nil {
188+
t.Fatal(err)
189+
}
190+
191+
got := askForConfirmation("Test prompt")
192+
if !got {
193+
t.Error("expected true for 'y', got false")
194+
}
195+
}
196+
197+
func TestAskForConfirmation_No(t *testing.T) {
198+
pipeRead, pipeWrite, err := os.Pipe()
199+
if err != nil {
200+
t.Fatal(err)
201+
}
202+
defer pipeRead.Close()
203+
defer pipeWrite.Close()
204+
205+
origStdin := os.Stdin
206+
os.Stdin = pipeRead
207+
defer func() { os.Stdin = origStdin }()
208+
209+
if _, err := pipeWrite.Write([]byte("n\n")); err != nil {
210+
t.Fatal(err)
211+
}
212+
213+
got := askForConfirmation("Test prompt")
214+
if got {
215+
t.Error("expected false for 'n', got true")
216+
}
217+
}
218+
219+
func TestLoadPrereqState_CorruptedFile(t *testing.T) {
220+
tempDir, err := os.MkdirTemp("", "prereq-corrupt-test")
221+
if err != nil {
222+
t.Fatal(err)
223+
}
224+
defer os.RemoveAll(tempDir)
225+
226+
origHome := os.Getenv("HOME")
227+
os.Setenv("HOME", tempDir)
228+
defer os.Setenv("HOME", origHome)
229+
230+
stateDir := filepath.Join(tempDir, stateDirName)
231+
if err := os.MkdirAll(stateDir, 0755); err != nil {
232+
t.Fatal(err)
233+
}
234+
statePath := filepath.Join(stateDir, stateFileName)
235+
236+
if err := os.WriteFile(statePath, []byte("invalid-json"), 0644); err != nil {
237+
t.Fatal(err)
238+
}
239+
240+
loaded := loadPrereqState()
241+
if loaded.GCloudSDKInstalled {
242+
t.Errorf("expected empty state for corrupted file, got GCloudSDKInstalled=true")
243+
}
244+
}
245+
246+
func TestSavePrereqState_WriteError(t *testing.T) {
247+
tempDir, err := os.MkdirTemp("", "prereq-write-error-test")
248+
if err != nil {
249+
t.Fatal(err)
250+
}
251+
defer os.RemoveAll(tempDir)
252+
253+
origHome := os.Getenv("HOME")
254+
os.Setenv("HOME", tempDir)
255+
defer os.Setenv("HOME", origHome)
256+
257+
stateDir := filepath.Join(tempDir, stateDirName)
258+
if err := os.MkdirAll(stateDir, 0755); err != nil {
259+
t.Fatal(err)
260+
}
261+
262+
// Make directory read-only to force write error
263+
if err := os.Chmod(stateDir, 0555); err != nil {
264+
t.Fatal(err)
265+
}
266+
defer os.Chmod(stateDir, 0755) // Restore to allow cleanup
267+
268+
state := PrereqState{
269+
GCloudAuthenticated: true,
270+
}
271+
272+
// This should log an error but not panic
273+
savePrereqState(state)
274+
}
275+
276+
277+
278+

cmd/job/submit.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import (
2121
"hpc-toolkit/pkg/logging"
2222
"hpc-toolkit/pkg/orchestrator"
2323
"hpc-toolkit/pkg/orchestrator/gke"
24-
"hpc-toolkit/pkg/prereq"
2524
"strings"
2625

2726
"github.com/spf13/cobra"
@@ -78,7 +77,7 @@ and JobSet/Kueue specific configurations like workload name, queue, nodes, and r
7877

7978
PreRunE: func(cmd *cobra.Command, args []string) error {
8079
logging.Info("Running prerequisite checks for 'gcluster job submit'...")
81-
err := prereq.EnsurePrerequisites(&projectID)
80+
err := EnsurePrerequisites(&projectID)
8281
if err != nil {
8382
return fmt.Errorf("prerequisite checks failed for 'gcluster job submit'. Please ensure your gcloud configuration and cluster context are valid: %w", err)
8483
}

0 commit comments

Comments
 (0)