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

Code before:

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

		sleepBeforeServingLeaseRenew(t, sleepDuration)

		prevUnavailableCount := getLeaseKeepAliveMetric(t, follower, "Unavailable")
		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))

		_, err = keepAliveClient.Recv()
		require.Equal(t, rpctypes.ErrGRPCTimeout, err)

		require.Eventually(t, func() bool {
			return getLeaseKeepAliveMetric(t, follower, "Unavailable") == prevUnavailableCount+1
		}, 3*time.Second, 100*time.Millisecond)
	})

	// Client set up with WithRequireLeader() will receive NoLeader error right after
	// monitorLeader() detects leader missing and cancels the server stream with ErrGRPCNoLeader.
	t.Run("catches NoLeader error with WithRequireLeader", func(t *testing.T) {
		leader, follower := setupLeaseForwardingCluster(t)
		followerClient := integration.ToGRPC(follower.Client).Lease

		prevUnavailableCount := getLeaseKeepAliveMetric(t, follower, "Unavailable")
		leader.Stop(t)

		keepAliveClient, err := followerClient.LeaseKeepAlive(clientv3.WithRequireLeader(t.Context()))
		require.NoError(t, err)

		_, err = keepAliveClient.Recv()
		require.Equal(t, rpctypes.ErrGRPCNoLeader, err)
		require.Equal(t, prevUnavailableCount+1, getLeaseKeepAliveMetric(t, follower, "Unavailable"))
	})
}

func setupLeaseForwardingCluster(t *testing.T) (*integration.Member, *integration.Member) {
	t.Helper()
	cluster := integration.NewCluster(t, &integration.ClusterConfig{Size: 2})
	t.Cleanup(func() { cluster.Terminate(t) })

	leaderIdx := cluster.WaitLeader(t)
	followerIdx := 1 - leaderIdx
	return cluster.Members[leaderIdx], cluster.Members[followerIdx]
}

func getLeaseKeepAliveMetric(t *testing.T, member *integration.Member, grpcCode string) int64 {
	t.Helper()
	metricVal, err := member.Metric(
		"grpc_server_handled_total",
		`grpc_method="LeaseKeepAlive"`,

Review comment (suggestion):
This is different than previous test. The original test killed 2 members to ensure there cannot be a leader, while you just kill leader which will cause a temporary unavailability but a new leader can be elected from other two members.

Code after:

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

		sleepBeforeServingLeaseRenew(t, sleepDuration)

		prevUnavailableCount := getLeaseKeepAliveMetric(t, follower, "Unavailable")
		require.NoError(t, keepAliveClient.Send(&pb.LeaseKeepAliveRequest{ID: leaseID}))

		_, err = keepAliveClient.Recv()
		require.Equal(t, rpctypes.ErrGRPCTimeout, err)

		require.Eventually(t, func() bool {
			return getLeaseKeepAliveMetric(t, follower, "Unavailable") == prevUnavailableCount+1
		}, 3*time.Second, 100*time.Millisecond)
	})

	// Client set up with WithRequireLeader() will receive NoLeader error right after
	// monitorLeader() detects leader missing and cancels the server stream with ErrGRPCNoLeader.
	t.Run("catches NoLeader error with WithRequireLeader", func(t *testing.T) {
		leader, follower, anotherFollower := setupLeaseForwardingCluster(t)
		followerClient := integration.ToGRPC(follower.Client).Lease

		prevUnavailableCount := getLeaseKeepAliveMetric(t, follower, "Unavailable")
		leader.Stop(t)
		anotherFollower.Stop(t)

		keepAliveClient, err := followerClient.LeaseKeepAlive(clientv3.WithRequireLeader(t.Context()))
		require.NoError(t, err)

		_, err = keepAliveClient.Recv()
		require.Equal(t, rpctypes.ErrGRPCNoLeader, err)
		require.Equal(t, prevUnavailableCount+1, getLeaseKeepAliveMetric(t, follower, "Unavailable"))
	})
}

func setupLeaseForwardingCluster(t *testing.T) (*integration.Member, *integration.Member, *integration.Member) {
	t.Helper()
	cluster := integration.NewCluster(t, &integration.ClusterConfig{Size: 3})
	t.Cleanup(func() { cluster.Terminate(t) })

	leaderIdx := cluster.WaitLeader(t)
	return cluster.Members[leaderIdx], cluster.Members[(leaderIdx+1)%3], cluster.Members[(leaderIdx+2)%3]
}

func getLeaseKeepAliveMetric(t *testing.T, member *integration.Member, grpcCode string) int64 {
	t.Helper()
	metricVal, err := member.Metric(
		"grpc_server_handled_total",
		`grpc_method="LeaseKeepAlive"`,