Repository: etcd-io/etcd
File: tests/integration/v3_lease_test.go (Go)

Code before:
func TestV3LeaseKeepAliveForwardingCatchError(t *testing.T) {
	if len(gofail.List()) == 0 {
		t.Skip("please run 'make gofail-enable' before running the test")
	}
	integration.BeforeTest(t)

	getMetricVal := func(follower *integration.Member, grpcCode string) int64 {
		metricVal, err := follower.Metric(
			"grpc_server_handled_total",
			`grpc_method="LeaseKeepAlive"`,
			fmt.Sprintf(`grpc_code="%v"`, grpcCode),
		)
		require.NoError(t, err)
		count, err := strconv.ParseInt(metricVal, 10, 32)
		require.NoError(t, err)
		return count
	}

	// This test case shows the current incorrect behavior.
	// When the client cancels the request during the follower server is forwarding
	// LeaseKeepAlive request to the leader, the server would generate error code
	// ErrGRPCNoLeader (gRPC Unavailable) but it should have generated [gRPC Canceled]
	// instead.
	// After we merge the fix, we need to update this test case to reflect the correct behavior.
	t.Run("client cancels while forwarding", func(t *testing.T) {
		follower, cancel, keepAliveClient, leaseID := setupLeaseKeepAliveForwarding(t)
		defer keepAliveClient.CloseSend()

		// It's insufficient to use the error return by keepAliveClient.Recv() and assert
		// the error code is Canceled. Such assertion always holds, because after the client
		// initialized cancel(), it simply returns Canceled and ignores the actual response
		// returned by server.
		// So in this test case, we'd need to use the metrics collected at the server side
		// to ensure the worker function `EtcdServer.LeaseRenew` properly handles the context
		// cancellation.
		prevCanceledCount := getMetricVal(follower, "Canceled")
		prevUnavailableCount := getMetricVal(follower, "Unavailable")

		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))
		// Make it more deterministic that the follower has entered the forwarding path,
		// then cancel the client context.
		time.Sleep(50 * time.Millisecond)
		cancel()

		// Wait for the server-side accounting to observe the canceled stream.
		// The server will incorrectly generate an Unavailable event (should be Canceled instead)
		require.Eventually(t, func() bool {
			return getMetricVal(follower, "Unavailable") == prevUnavailableCount+1
		}, 3*time.Second, 100*time.Millisecond)

		// We didn't see the Canceled count increase

Review comment (suggestion):
Setup function should not be 80% of test code. Tests should have 3 separate phases, Setup,Execute,Assert. While there can be some code reuse, the test should not only what you are asserting, but also actions that lead to it. Without them here you don't know what lead to the results. 

Please consider moving some code from setupLeaseKeepAliveForwarding to the test. Also remove comments.

Code after:
func TestV3LeaseKeepAliveForwardingCatchError(t *testing.T) {
	if len(gofail.List()) == 0 {
		t.Skip("please run 'make gofail-enable' before running the test")
	}
	integration.BeforeTest(t)
	// Longer than leaseHandler.ServeHTTP()'s default timeout duration
	sleepDuration := 8 * time.Second

	t.Run("forwarding succeeds", func(t *testing.T) {
		leaderClient, follower := setupLeaseForwardingCluster(t)

		grantResp, err := leaderClient.LeaseGrant(t.Context(), &pb.LeaseGrantRequest{TTL: 30})
		require.NoError(t, err)
		leaseID := grantResp.ID

		keepAliveClient, err := integration.ToGRPC(follower.Client).Lease.LeaseKeepAlive(t.Context())
		require.NoError(t, err)
		defer keepAliveClient.CloseSend()

		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))
		resp, err := keepAliveClient.Recv()
		require.NoError(t, err)
		require.Equal(t, leaseID, resp.ID)
		require.Positive(t, resp.TTL)
	})

	// Shows current behavior: client cancel during forwarding incorrectly returns Unavailable.
	t.Run("client cancels while forwarding", func(t *testing.T) {
		leaderClient, follower := setupLeaseForwardingCluster(t)

		grantResp, err := leaderClient.LeaseGrant(t.Context(), &pb.LeaseGrantRequest{TTL: 30})
		require.NoError(t, err)
		leaseID := grantResp.ID

		ctx, cancel := context.WithCancel(t.Context())
		keepAliveClient, err := integration.ToGRPC(follower.Client).Lease.LeaseKeepAlive(ctx)
		require.NoError(t, err)
		defer keepAliveClient.CloseSend()

		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))
		_, err = keepAliveClient.Recv()
		require.NoError(t, err)

		sleepBeforeServingLeaseRenew(t, sleepDuration)

		// Use server metrics to verify behavior since client.Recv() always returns Canceled
		// after cancel() regardless of the actual server response.
		prevCanceledCount := getLeaseKeepAliveMetric(t, follower, "Canceled")
		prevUnavailableCount := getLeaseKeepAliveMetric(t, follower, "Unavailable")

		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))