update unit test to look for gpuType in worker pool

beam-cloud · Jan 19, 2024 · 64e6e67 · 64e6e67
1 parent 32d9e20
commit 64e6e67
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -36,7 +36,7 @@ func NewSchedulerForTest() (*Scheduler, error) {
 		return nil, err
 	}
 
-	poolJson := []byte(`{"worker":{"pools":{"beam-cpu":{},"beam-a10g":{},"beam-t4":{}}}}}`)
+	poolJson := []byte(`{"worker":{"pools":{"beam-cpu":{},"beam-a10g":{"gpuType": "A10G"},"beam-t4":{"gpuType": "T4"}}}}}`)
 	configManager.LoadConfig(common.YAMLConfigFormat, rawbytes.Provider(poolJson))
 	config := configManager.GetConfig()
 
@@ -192,9 +192,9 @@ func TestGetController(t *testing.T) {
 		wb, _ := NewSchedulerForTest()
 
 		cpuRequest := &types.ContainerRequest{Gpu: ""}
-		cpuController, err := wb.getController(cpuRequest)
-		if err != nil || cpuController.Name() != "beam-cpu" {
-			t.Errorf("Expected beam-cpu controller, got %v, error: %v", cpuController, err)
+		defaultController, err := wb.getController(cpuRequest)
+		if err != nil || defaultController.Name() != "default" {
+			t.Errorf("Expected default controller, got %v, error: %v", defaultController, err)
 		}
 
 		a10gRequest := &types.ContainerRequest{Gpu: "A10G"}