1 files changed, 55 insertions, 0 deletions
diff --git a/workhorse/internal/redis/keywatcher_test.go b/workhorse/internal/redis/keywatcher_test.go
index f1ee77e2194..99892bc64b8 100644
--- a/workhorse/internal/redis/keywatcher_test.go
+++ b/workhorse/internal/redis/keywatcher_test.go
@@ -160,3 +160,58 @@ func TestWatchKeyMassivelyParallel(t *testing.T) {
 	processMessages(runTimes, "somethingelse")
 	wg.Wait()
 }
+
+func TestShutdown(t *testing.T) {
+	conn, td := setupMockPool()
+	defer td()
+	defer func() { shutdown = make(chan struct{}) }()
+
+	conn.Command("GET", runnerKey).Expect("something")
+
+	wg := &sync.WaitGroup{}
+	wg.Add(2)
+
+	go func() {
+		val, err := WatchKey(runnerKey, "something", 10*time.Second)
+
+		require.NoError(t, err, "Expected no error")
+		require.Equal(t, WatchKeyStatusNoChange, val, "Expected value not to change")
+		wg.Done()
+	}()
+
+	go func() {
+		for countWatchers(runnerKey) == 0 {
+			time.Sleep(time.Millisecond)
+		}
+
+		require.Equal(t, 1, countWatchers(runnerKey))
+
+		Shutdown()
+		wg.Done()
+	}()
+
+	wg.Wait()
+
+	for countWatchers(runnerKey) == 1 {
+		time.Sleep(time.Millisecond)
+	}
+
+	require.Equal(t, 0, countWatchers(runnerKey))
+
+	// Adding a key after the shutdown should result in an immediate response
+	var val WatchKeyStatus
+	var err error
+	done := make(chan struct{})
+	go func() {
+		val, err = WatchKey(runnerKey, "something", 10*time.Second)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		require.NoError(t, err, "Expected no error")
+		require.Equal(t, WatchKeyStatusNoChange, val, "Expected value not to change")
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("timeout waiting for WatchKey")
+	}
+}