From ef7063a22e17d01eed9584e0103bf2c8e159532a Mon Sep 17 00:00:00 2001
From: Inanc Gumus <m@inanc.io>
Date: Mon, 11 Oct 2021 17:39:56 +0300
Subject: [PATCH] Add a test for handleRemainingVUs

---
 lib/executor/ramping_vus.go      |  1 -
 lib/executor/ramping_vus_test.go | 74 ++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/lib/executor/ramping_vus.go b/lib/executor/ramping_vus.go
index 52a28d2d7c1..43558385e02 100644
--- a/lib/executor/ramping_vus.go
+++ b/lib/executor/ramping_vus.go
@@ -635,7 +635,6 @@ func (rs rampingVUsRunState) handleVUs(ctx context.Context) {
 	}
 }
 
-// TODO: removing this has no effect on tests?
 func (rs rampingVUsRunState) handleRemainingVUs(ctx context.Context) {
 	var (
 		handleNewMaxAllowedVUs = rs.maxAllowedVUsHandlerStrategy()
diff --git a/lib/executor/ramping_vus_test.go b/lib/executor/ramping_vus_test.go
index 71b2ee36a77..50798b7aac6 100644
--- a/lib/executor/ramping_vus_test.go
+++ b/lib/executor/ramping_vus_test.go
@@ -280,6 +280,80 @@ func TestRampingVUsGracefulRampDown(t *testing.T) {
 	}
 }
 
+func TestRampingVUsHandleRemainingVUs(t *testing.T) {
+	t.Parallel()
+
+	// there will be this number of VUs at max:
+	const maxVus = 2
+
+	// each VU will sleep an amount of time per iteration = sleep time.
+	// the sleep time will be more than the time of all stages plus
+	// graceful stop. so the VUs won't be able to finish their iterations.
+	//
+	// however:
+	// because of the GracefulRampDown, one of the VUs will be allowed to
+	// complete its iteration. and the other one will be interrupted.
+	const (
+		wantVuInterrupted uint32 = 1 // one VU will be interrupted
+		wantVuFinished    uint32 = 1 // one VU will finish an iteration
+	)
+
+	cfg := RampingVUsConfig{
+		BaseConfig: BaseConfig{
+			// give the test run more time to finish
+			GracefulStop: types.NullDurationFrom(50 * time.Millisecond),
+		},
+
+		// give VUs this amount of time to finish their iterations:
+		GracefulRampDown: types.NullDurationFrom(30 * time.Millisecond),
+
+		Stages: []Stage{
+			// ramp up to maxVUs immediately
+			{
+				Duration: types.NullDurationFrom(10 * time.Millisecond),
+				Target:   null.IntFrom(int64(maxVus)),
+			},
+			// ramp down to zero VUs and let the GracefulStop
+			// period to kick in before the first VU wakes up
+			{
+				Duration: types.NullDurationFrom(40 * time.Millisecond),
+				Target:   null.IntFrom(int64(0)),
+			},
+		},
+	}
+
+	var gotVuInterrupted, gotVuFinished uint32
+
+	iteration := func(ctx context.Context) error {
+		select {
+		case <-time.After(65 * time.Millisecond):
+			atomic.AddUint32(&gotVuFinished, 1)
+		case <-ctx.Done():
+			// the context used has the deadline of:
+			// stage durations + GracefulStop.
+			//
+			// here, the context becomes cancelled and one of the
+			// VUs is being interrupted.
+			atomic.AddUint32(&gotVuInterrupted, 1)
+		}
+		return nil
+	}
+
+	// run the executor
+	et, err := lib.NewExecutionTuple(nil, nil)
+	require.NoError(t, err)
+	ctx, cancel, executor, _ := setupExecutor(
+		t, cfg,
+		lib.NewExecutionState(lib.Options{}, et, maxVus, maxVus),
+		simpleRunner(iteration),
+	)
+	defer cancel()
+	require.NoError(t, executor.Run(ctx, nil, nil))
+
+	assert.Equal(t, wantVuInterrupted, atomic.LoadUint32(&gotVuInterrupted))
+	assert.Equal(t, wantVuFinished, atomic.LoadUint32(&gotVuFinished))
+}
+
 // Ensure there's no wobble of VUs during graceful ramp-down, without segments.
 // See https://github.com/k6io/k6/issues/1296
 func TestRampingVUsRampDownNoWobble(t *testing.T) {