package kas

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"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/testhelpers"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"
)

func TestAgentConnPool_GetNonExistentConnection(t *testing.T) {
	ctrl := gomock.NewController(t)
	routingConn := mock_rpc.NewMockClientConnInterface(ctrl)
	pool := newAgentConnPool(routingConn)

	assertWrapper(t, pool, routingConn)
}

func TestAgentConnPool_AddGetConnection(t *testing.T) {
	ctrl := gomock.NewController(t)
	routingConn := mock_rpc.NewMockClientConnInterface(ctrl)
	pool := newAgentConnPool(routingConn)

	newConn := mock_rpc.NewMockClientConnInterface(ctrl)
	err := pool.Add(testhelpers.AgentkKey1, newConn)
	require.NoError(t, err)

	conn := pool.Get(testhelpers.AgentkKey1)
	assert.Same(t, newConn, conn)
}

func TestAgentConnPool_AddRemoveGetConnection(t *testing.T) {
	ctrl := gomock.NewController(t)
	routingConn := mock_rpc.NewMockClientConnInterface(ctrl)
	pool := newAgentConnPool(routingConn)

	newConn := mock_rpc.NewMockClientConnInterface(ctrl)
	err := pool.Add(testhelpers.AgentkKey1, newConn)
	require.NoError(t, err)

	pool.Remove(testhelpers.AgentkKey1)

	assertWrapper(t, pool, routingConn)
}

func assertWrapper(t *testing.T, pool *agentConnPool, delegate grpc.ClientConnInterface) {
	conn := pool.Get(testhelpers.AgentkKey1)
	wrapper := conn.(*routingMetadataWrapper)
	assert.Same(t, delegate, wrapper.delegate)
	assert.Equal(t, testhelpers.AgentkKey1, wrapper.agentKey)
}
