package router

import (
	"context"
	"io"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/matcher"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_tunnel_rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_tunnel_tunserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/info"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/rpc"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tunserver"
	otelmetric "go.opentelemetry.io/otel/metric"
	metricnoop "go.opentelemetry.io/otel/metric/noop"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"k8s.io/apimachinery/pkg/util/wait"
)

// "slow" tests in this file are marked for concurrent execution with t.Parallel()

const (
	serviceName    = "gitlab.service1"
	methodName     = "DoSomething"
	fullMethodName = "/" + serviceName + "/" + methodName
)

var (
	_ Handler          = &Registry{}
	_ Querier          = &Registry{}
	_ tunserver.Tunnel = (*tunnelHolder)(nil)
)

func TestStopUnregistersAllConnections(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	gomock.InOrder(
		connectServer.EXPECT().
			Recv().
			Return(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Descriptor_{
					Descriptor_: descriptor(),
				},
			}, nil),
		rb.EXPECT().
			Register(gomock.Any(), testhelpers.AgentkKey1),
		rb.EXPECT().
			Do(gomock.Any()).
			Do(func(ctx context.Context) error {
				cancel()
				return nil
			}),
		rb.EXPECT().
			Unregister(testhelpers.AgentkKey1),
		rb.EXPECT().
			Do(gomock.Any()),
	)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		handleErr := r.HandleTunnel(connectServer.Context(), context.Background(), testhelpers.AgentkInfoObj().Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		})
		assert.NoError(t, handleErr)
	})
	err = r.Run(ctx)
	assert.NoError(t, err)
}

func TestTunnelDoneRegistersUnusedTunnel(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	reg := make(chan struct{})
	gomock.InOrder(
		connectServer.EXPECT().
			Recv().
			Return(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Descriptor_{
					Descriptor_: descriptor(),
				},
			}, nil),
		rb.EXPECT(). // HandleTunnel()
				Register(gomock.Any(), gomock.Any()).
				Do(func(ttl time.Duration, agentKeys ...api.AgentKey) {
				close(reg)
			}),
		rb.EXPECT(). // FindTunnel()
				Unregister(gomock.Any()),
		rb.EXPECT(). // Done()
				Register(gomock.Any(), gomock.Any()),
		rb.EXPECT(). // FindTunnel()
				Unregister(gomock.Any()),
		rb.EXPECT(). // Done()
				Register(gomock.Any(), gomock.Any()),
		rb.EXPECT(). // stopInternal()
				Unregister(gomock.Any()),
	)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())

	agentInfo := testhelpers.AgentkInfoObj()
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		handleErr := r.HandleTunnel(connectServer.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		})
		assert.NoError(t, handleErr)
	})
	wg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	<-reg
	found, th := r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
	assert.True(t, found)
	tun, err := th.Get(context.Background())
	require.NoError(t, err)
	tun.Done(context.Background())
	th.Done(context.Background())
	found, th = r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
	assert.True(t, found)
	tun, err = th.Get(context.Background())
	require.NoError(t, err)
	tun.Done(context.Background())
	th.Done(context.Background())
	cancel()
}

func TestTunnelDoneDonePanics(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	reg := make(chan struct{})
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()).
			Do(func(ttl time.Duration, agentKeys ...api.AgentKey) {
				close(reg)
			}),
		rb.EXPECT().
			Unregister(gomock.Any()),
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	agentInfo := testhelpers.AgentkInfoObj()
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		handleErr := r.HandleTunnel(connectServer.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		})
		assert.NoError(t, handleErr)
	})
	wg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	<-reg
	found, th := r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
	assert.True(t, found)
	tun, err := th.Get(context.Background())
	require.NoError(t, err)
	tun.Done(context.Background())
	assert.PanicsWithError(t, "unreachable: ready -> done should never happen", func() {
		tun.Done(context.Background())
	})
	th.Done(context.Background())
	cancel()
}

func TestHandleTunnelIsUnblockedByContext(t *testing.T) {
	t.Parallel()
	ctxConn, cancelConn := context.WithTimeout(context.Background(), 50*time.Millisecond) // will unblock HandleTunnel()
	defer cancelConn()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	gomock.InOrder(
		connectServer.EXPECT().
			Recv().
			Return(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Descriptor_{
					Descriptor_: descriptor(),
				},
			}, nil),
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	err = r.HandleTunnel(connectServer.Context(), ctxConn, testhelpers.AgentkInfoObj().Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
		return NewTunnelForwarder(connectServer, tunnelRetErr)
	})
	assert.NoError(t, err)
	cancel()
}

// Two tunnels with the same agent key. Both register. Then one of them is retrieved via FindTunnel()
// and then its context is canceled. If this test gets stuck, we have a problem.
// Reproducer for https://gitlab.com/gitlab-org/cluster-integration/gitlab-agent/-/issues/183.
func TestHandleTunnelIsUnblockedByContext_WithTwoTunnels(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer1 := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer2 := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer1.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer2.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	d1 := descriptor()
	connectServer1.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: d1,
			},
		}, nil)
	connectServer2.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wgReg wait.Group
	defer wgReg.Wait()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	wgReg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	var wg wait.Group
	defer wg.Wait()
	agentInfo := testhelpers.AgentkInfoObj()
	ctx1, cancel1 := context.WithCancel(context.Background())
	defer cancel1()
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(connectServer1.Context(), ctx1, agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer1, tunnelRetErr)
		}))
	})
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(connectServer2.Context(), ctx2, agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer2, tunnelRetErr)
		}))
	})
	// wait for both to register
	assert.Eventually(t, func() bool {
		r.mu.Lock()
		defer r.mu.Unlock()
		return len(r.tunsByAgentKey[agentInfo.Key].tunHolders) == 2
	}, time.Second, 10*time.Millisecond)
	found, th := r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
	assert.True(t, found)
	tun, err := th.Get(context.Background())
	require.NoError(t, err)
	// cancel context for the found tunnel
	switch tun.(*tunnelHolder).forwarder.(*TunnelForwarder).Tunnel {
	case connectServer1:
		cancel1()
	case connectServer2:
		cancel2()
	default:
		t.FailNow()
	}
	assert.Eventually(t, func() bool {
		r.mu.Lock()
		defer r.mu.Unlock()
		return tun.(*tunnelHolder).state == tunserver.StateContextDone
	}, time.Second, 10*time.Millisecond)
	tun.Done(context.Background())
	th.Done(context.Background())
}

func TestHandleTunnelReturnErrOnRecvErr(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(nil, status.Error(codes.DataLoss, "expected err"))
	rb := NewMockRegistrationBuilder(ctrl)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	err = r.HandleTunnel(connectServer.Context(), context.Background(), testhelpers.AgentkInfoObj().Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
		return NewTunnelForwarder(connectServer, tunnelRetErr)
	})
	assert.EqualError(t, err, "rpc error: code = DataLoss desc = expected err")
}

func TestHandleTunnelReturnErrOnInvalidMsg(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Header{
				Header: &rpc.Header{},
			},
		}, nil)
	rb := NewMockRegistrationBuilder(ctrl)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	err = r.HandleTunnel(connectServer.Context(), context.Background(), testhelpers.AgentkInfoObj().Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
		return NewTunnelForwarder(connectServer, tunnelRetErr)
	})
	assert.EqualError(t, err, "rpc error: code = InvalidArgument desc = invalid oneof value type: *rpc.ConnectRequest_Header")
}

func TestHandleTunnelIsMatchedToIncomingConnection(t *testing.T) {
	t.Parallel()
	incomingStream, cb, tunnel, r := setupStreams(t, true)
	agentInfo := testhelpers.AgentkInfoObj()
	var wg wait.Group
	defer wg.Wait()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	wg.Start(func() {
		err := r.Run(ctx)
		assert.NoError(t, err)
	})
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(tunnel.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) { return NewTunnelForwarder(tunnel, tunnelRetErr) }))
	})
	time.Sleep(50 * time.Millisecond)
	found, th := r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
	defer th.Done(context.Background())
	assert.True(t, found)
	tun, err := th.Get(context.Background())
	require.NoError(t, err)
	defer tun.Done(context.Background())
	err = tun.ForwardStream(testlogger.New(t), incomingStream, cb)
	require.NoError(t, err)
}

func TestHandleTunnelIsNotMatchedToIncomingConnectionForMissingMethod(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	agentInfo := testhelpers.AgentkInfoObj()
	var wg wait.Group
	defer wg.Wait()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	wg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(connectServer.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		}))
	})
	time.Sleep(50 * time.Millisecond)
	ctx2, cancel2 := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel2()
	found, th := r.FindTunnel(context.Background(), agentInfo.Key, "missing_service", "missing_method")
	defer th.Done(context.Background())
	assert.False(t, found)
	_, err = th.Get(ctx2)
	assert.Equal(t, context.DeadlineExceeded, err)
}

func TestForwardStreamIsMatchedToHandleTunnel(t *testing.T) {
	t.Parallel()
	incomingStream, cb, tunnel, r := setupStreams(t, false)
	agentInfo := testhelpers.AgentkInfoObj()
	var wg wait.Group
	defer wg.Wait()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	wg.Start(func() {
		err := r.Run(ctx)
		assert.NoError(t, err)
	})
	wg.Start(func() {
		_, th := r.FindTunnel(context.Background(), agentInfo.Key, serviceName, methodName)
		defer th.Done(context.Background())
		tun, err := th.Get(context.Background())
		if !assert.NoError(t, err) {
			return
		}
		defer tun.Done(context.Background())
		err = tun.ForwardStream(testlogger.New(t), incomingStream, cb)
		assert.NoError(t, err)
	})
	time.Sleep(50 * time.Millisecond)
	err := r.HandleTunnel(tunnel.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) { return NewTunnelForwarder(tunnel, tunnelRetErr) })
	require.NoError(t, err)
}

func TestForwardStreamIsNotMatchedToHandleTunnelForMissingMethod(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	agentInfo := testhelpers.AgentkInfoObj()
	var wg wait.Group
	defer wg.Wait()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	wg.Start(func() {
		runErr := r.Run(ctx)
		assert.NoError(t, runErr)
	})
	wg.Start(func() {
		found, th := r.FindTunnel(context.Background(), agentInfo.Key, "missing_service", "missing_method")
		defer th.Done(context.Background())
		assert.False(t, found)
		_, findErr := th.Get(context.Background())
		assert.EqualError(t, findErr, "rpc error: code = Canceled desc = kas is shutting down")
	})
	time.Sleep(50 * time.Millisecond)
	ctx2, cancel2 := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel2()
	err = r.HandleTunnel(connectServer.Context(), ctx2, agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
		return NewTunnelForwarder(connectServer, tunnelRetErr)
	})
	assert.NoError(t, err)
}

func TestFindTunnelIsUnblockedByContext(t *testing.T) {
	t.Parallel()
	ctxConn, cancelConn := context.WithTimeout(context.Background(), 50*time.Millisecond) // will unblock FindTunnel()
	defer cancelConn()

	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	rb := NewMockRegistrationBuilder(ctrl)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	found, th := r.FindTunnel(context.Background(), testhelpers.AgentkKey1, serviceName, methodName)
	defer th.Done(context.Background())
	assert.False(t, found)
	_, err = th.Get(ctxConn)
	assert.Equal(t, context.DeadlineExceeded, err)
}

func TestGC(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	tunnelTracker.EXPECT().
		GC(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, ids []api.AgentKey) (int, error) {
			cancel()
			return 3, nil
		})
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	agentInfo := testhelpers.AgentkInfoObj()
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, 50*time.Millisecond, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(connectServer.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		}))
	})

	err = r.Run(ctx)
	require.NoError(t, err)
}

func TestRefreshRegistrations(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
	rb := NewMockRegistrationBuilder(ctrl)
	rb.EXPECT().
		Do(gomock.Any()).
		MinTimes(1)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	gomock.InOrder(
		rb.EXPECT().
			Register(gomock.Any(), gomock.Any()),
		rb.EXPECT().
			Refresh(gomock.Any(), gomock.Any()).
			Do(func(tttl time.Duration, agentKeys ...api.AgentKey) {
				cancel()
			}),
		rb.EXPECT().
			Unregister(gomock.Any()),
	)
	agentInfo := testhelpers.AgentkInfoObj()
	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), 50*time.Millisecond, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	var wg wait.Group
	defer wg.Wait()
	wg.Start(func() {
		assert.NoError(t, r.HandleTunnel(connectServer.Context(), context.Background(), agentInfo.Key, func(tunnelRetErr chan<- error) (Forwarder, error) {
			return NewTunnelForwarder(connectServer, tunnelRetErr)
		}))
	})

	err = r.Run(ctx)
	require.NoError(t, err)
}

func setupStreams(t *testing.T, expectRegisterTunnel bool) (*mock_rpc.MockServerStream, *mock_tunnel_tunserver.MockDataCallback, *mock_tunnel_rpc.MockReverseTunnel_ConnectServer, *Registry) {
	const metaKey = "Cba"
	meta := metadata.MD{}
	meta.Set(metaKey, "3", "4")
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	sts := mock_rpc.NewMockServerTransportStream(ctrl)
	cb := mock_tunnel_tunserver.NewMockDataCallback(ctrl)

	incomingCtx := grpc.NewContextWithServerTransportStream(context.Background(), sts)
	incomingCtx = metadata.NewIncomingContext(incomingCtx, meta)
	incomingStream := mock_rpc.NewMockServerStream(ctrl)
	incomingStream.EXPECT().
		Context().
		Return(incomingCtx).
		MinTimes(1)

	rb := NewMockRegistrationBuilder(ctrl)
	tunnelTracker := NewMockTracker(ctrl)
	tunnelTracker.EXPECT().
		RegistrationBuilder().
		Return(rb)
	connectServer := mock_tunnel_rpc.NewMockReverseTunnel_ConnectServer(ctrl)
	connectServer.EXPECT().
		Context().
		Return(context.Background()).
		MinTimes(1)
	connectServer.EXPECT().
		Recv().
		Return(&rpc.ConnectRequest{
			Msg: &rpc.ConnectRequest_Descriptor_{
				Descriptor_: descriptor(),
			},
		}, nil)
	if expectRegisterTunnel {
		rb.EXPECT().
			Do(gomock.Any()).
			MinTimes(1)
		gomock.InOrder(
			rb.EXPECT().
				Register(gomock.Any(), testhelpers.AgentkKey1),
			rb.EXPECT().
				Unregister(testhelpers.AgentkKey1),
		)
	}
	frame := rpc.RawFrame{}
	gomock.InOrder(
		sts.EXPECT().
			Method().
			Return(fullMethodName).
			MinTimes(1),
		connectServer.EXPECT().
			Send(matcher.ProtoEq(t, &rpc.ConnectResponse{
				Msg: &rpc.ConnectResponse_RequestInfo{
					RequestInfo: &rpc.RequestInfo{
						MethodName: fullMethodName,
						Meta: []*rpc.MetadataKV{
							mock_tunnel_rpc.NewMetadataKV("cba", "3", "4"),
						},
					},
				},
			}, mock_tunnel_rpc.EquateMetadataKV())),
		incomingStream.EXPECT().
			RecvMsg(gomock.Any()).
			Do(testhelpers.RecvMsg(&frame)),
		connectServer.EXPECT().
			Send(matcher.ProtoEq(t, &rpc.ConnectResponse{
				Msg: &rpc.ConnectResponse_Message{
					Message: &rpc.Message{
						Data: frame.Data.Materialize(),
					},
				},
			})),
		incomingStream.EXPECT().
			RecvMsg(gomock.Any()).
			Return(io.EOF),
		connectServer.EXPECT().
			Send(matcher.ProtoEq(t, &rpc.ConnectResponse{
				Msg: &rpc.ConnectResponse_CloseSend{
					CloseSend: &rpc.CloseSend{},
				},
			})),
	)
	gomock.InOrder(
		connectServer.EXPECT().
			RecvMsg(gomock.Any()).
			Do(testhelpers.RecvMsg(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Header{
					Header: &rpc.Header{
						Meta: []*rpc.MetadataKV{
							mock_tunnel_rpc.NewMetadataKV("resp", "1", "2"),
						},
					},
				},
			})),
		cb.EXPECT().
			Header(matcher.ProtoEq(t, []*rpc.MetadataKV{
				mock_tunnel_rpc.NewMetadataKV("resp", "1", "2"),
			}, mock_tunnel_rpc.EquateMetadataKV())),
		connectServer.EXPECT().
			RecvMsg(gomock.Any()).
			Do(testhelpers.RecvMsg(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Message{
					Message: &rpc.Message{
						Data: []byte{5, 6, 7},
					},
				},
			})),
		cb.EXPECT().
			Message([]byte{5, 6, 7}),
		connectServer.EXPECT().
			RecvMsg(gomock.Any()).
			Do(testhelpers.RecvMsg(&rpc.ConnectRequest{
				Msg: &rpc.ConnectRequest_Trailer{
					Trailer: &rpc.Trailer{
						Meta: []*rpc.MetadataKV{
							mock_tunnel_rpc.NewMetadataKV("trailer", "8", "9"),
						},
					},
				},
			})),
		cb.EXPECT().
			Trailer(matcher.ProtoEq(t, []*rpc.MetadataKV{
				mock_tunnel_rpc.NewMetadataKV("trailer", "8", "9"),
			}, mock_tunnel_rpc.EquateMetadataKV())),
		connectServer.EXPECT().
			RecvMsg(gomock.Any()).
			Return(io.EOF),
	)

	r, err := NewRegistry(testlogger.New(t), mockAPI, nt(), nm(), time.Minute, time.Minute, time.Minute, tunnelTracker)
	require.NoError(t, err)
	return incomingStream, cb, connectServer, r
}

func descriptor() *rpc.Descriptor {
	return &rpc.Descriptor{
		ApiDescriptor: &info.APIDescriptor{
			Services: []*info.Service{
				{
					Name: serviceName,
					Methods: []*info.Method{
						{
							Name: methodName,
						},
					},
				},
			},
		},
	}
}

func nt() trace.Tracer {
	return noop.NewTracerProvider().Tracer("test")
}

func nm() otelmetric.Meter {
	return metricnoop.NewMeterProvider().Meter("test")
}
