[sk8s] Fix setting the maintenance mode in bot_config.
Change-Id: I917086191cc7c210ff942f08931b2b6bb517d286
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/306898
Commit-Queue: Joe Gregorio <jcgregorio@google.com>
Reviewed-by: Ravi Mistry <rmistry@google.com>
diff --git a/sk8s/go/bot_config/machine/machine.go b/sk8s/go/bot_config/machine/machine.go
index 64d1db7..6b32415 100644
--- a/sk8s/go/bot_config/machine/machine.go
+++ b/sk8s/go/bot_config/machine/machine.go
@@ -156,16 +156,20 @@
}
})
- // Also start a second loop that does a firestore onsnapshot watcher that gets the dims we should
- // be reporting to swarming.
+ m.startStoreWatch(ctx)
+ return nil
+}
+
+// startStoreWatch starts a loop that does a firestore onsnapshot watcher
+// that gets the dims and state we should be reporting to swarming.
+func (m *Machine) startStoreWatch(ctx context.Context) {
go func() {
for desc := range m.store.Watch(ctx, m.MachineID) {
m.storeWatchArrivalCounter.Inc(1)
m.SetDimensionsForSwarming(desc.Dimensions)
+ m.SetMaintenanceMode(desc.Mode == machine.ModeRecovery || desc.Mode == machine.ModeMaintenance)
}
}()
-
- return nil
}
// SetDimensionsForSwarming sets the dimensions that should be reported to swarming. Should only
diff --git a/sk8s/go/bot_config/machine/machine_test.go b/sk8s/go/bot_config/machine/machine_test.go
index af81da9..487b811 100644
--- a/sk8s/go/bot_config/machine/machine_test.go
+++ b/sk8s/go/bot_config/machine/machine_test.go
@@ -162,10 +162,77 @@
assert.Equal(t, int64(1), m.storeWatchArrivalCounter.Get())
assert.Equal(t, int64(0), m.interrogateAndSendFailures.Get())
- // Confirm the firestore write make it all the way to Dims().
+ // Confirm the firestore write made it all the way to Dims().
assert.Equal(t, machine.SwarmingDimensions{"foo": {"bar"}}, m.DimensionsForSwarming())
}
+func TestStart_FirestoreWritesGetReflectedToMachine(t *testing.T) {
+ // Manual because we are testing pubsub.
+ unittest.ManualTest(t)
+
+ ctx, _, instanceConfig := setupConfig(t)
+ ctx, cancel := context.WithCancel(ctx)
+
+ // Set the SWARMING_BOT_ID env variable.
+ oldVar := os.Getenv(swarming.SwarmingBotIDEnvVar)
+ err := os.Setenv(swarming.SwarmingBotIDEnvVar, "my-test-bot-001")
+ require.NoError(t, err)
+ defer func() {
+ err = os.Setenv(swarming.SwarmingBotIDEnvVar, oldVar)
+ require.NoError(t, err)
+ }()
+
+ const imageName = "gcr.io/skia-public/rpi-swarming-client:2020-05-09T19_28_20Z-jcgregorio-4fef3ca-clean"
+
+ // Set the IMAGE_NAME env variable.
+ oldImageVar := os.Getenv(swarming.KubernetesImageEnvVar)
+ err = os.Setenv(swarming.KubernetesImageEnvVar, imageName)
+ require.NoError(t, err)
+ defer func() {
+ err = os.Setenv(swarming.KubernetesImageEnvVar, oldImageVar)
+ require.NoError(t, err)
+ }()
+
+ // Create a Machine instance.
+ start := time.Date(2020, time.May, 1, 0, 0, 0, 0, time.UTC)
+ m, err := New(ctx, true, instanceConfig, start)
+ require.NoError(t, err)
+ assert.Equal(t, "my-test-bot-001", m.MachineID)
+
+ assert.False(t, m.GetMaintenanceMode())
+
+ // Start just the Firestore watcher.
+ m.startStoreWatch(ctx)
+
+ // Write a description into firestore. We expect the dimensions and mode to
+ // bubble down to the machine
+ err = m.store.Update(ctx, "my-test-bot-001", func(machine.Description) machine.Description {
+ ret := machine.NewDescription()
+ ret.Mode = machine.ModeMaintenance
+ ret.Dimensions["foo"] = []string{"bar"}
+ return ret
+ })
+ require.NoError(t, err)
+
+ assert.Equal(t, machine.SwarmingDimensions{"foo": {"bar"}}, m.DimensionsForSwarming())
+ assert.True(t, m.GetMaintenanceMode())
+
+ // Now change the mode.
+ err = m.store.Update(ctx, "my-test-bot-001", func(machine.Description) machine.Description {
+ ret := machine.NewDescription()
+ ret.Mode = machine.ModeAvailable
+ return ret
+ })
+ require.NoError(t, err)
+
+ // Confirm we go out of maintenance mode.
+ assert.False(t, m.GetMaintenanceMode())
+
+ // Cancel Go routine inside startStoreWatch.
+ cancel()
+
+}
+
func Test_FakeExe_AdbShellGetProp_Success(t *testing.T) {
unittest.FakeExeTest(t)
if os.Getenv(executil.OverrideEnvironmentVariable) == "" {
diff --git a/sk8s/go/bot_config/server/server_test.go b/sk8s/go/bot_config/server/server_test.go
index 1eec30b..e0a2a86 100644
--- a/sk8s/go/bot_config/server/server_test.go
+++ b/sk8s/go/bot_config/server/server_test.go
@@ -55,6 +55,34 @@
require.NoError(t, err)
assert.Equal(t, someRackName, dict["sk_rack"])
assert.Equal(t, "bar", dict["foo"])
+ _, ok := dict["maintenance"]
+ assert.False(t, ok)
+}
+
+func TestGetState_MaintenanceAppearsInStateResponse(t *testing.T) {
+ unittest.SmallTest(t)
+
+ const someRackName = "some-rack-name"
+
+ err := os.Setenv("MY_RACK_NAME", someRackName)
+ require.NoError(t, err)
+
+ r := httptest.NewRequest("POST", "/get_state", strings.NewReader("{\"foo\":\"bar\"}"))
+
+ m := &botmachine.Machine{}
+ m.SetMaintenanceMode(true)
+ s, err := New(m)
+ require.NoError(t, err)
+ w := httptest.NewRecorder()
+
+ s.getState(w, r)
+
+ res := w.Result()
+ assert.Equal(t, 200, res.StatusCode)
+ var dict map[string]interface{}
+ err = json.NewDecoder(res.Body).Decode(&dict)
+ require.NoError(t, err)
+ assert.True(t, dict["maintenance"].(bool))
}
func TestGetState_ErrOnInvalidJSON(t *testing.T) {
