package managerserver

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	pbm "github.com/tidbcloud/aws-shared-provider/pbgen/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/notifier"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"google.golang.org/grpc"
)

type watchBackendV2TestStream struct {
	grpc.ServerStream
	ctx       context.Context
	cancel    context.CancelFunc
	responses []*pbm.BackendResponseV2
}

func (s *watchBackendV2TestStream) Context() context.Context {
	return s.ctx
}

func (s *watchBackendV2TestStream) Send(response *pbm.BackendResponseV2) error {
	s.responses = append(s.responses, response)
	s.cancel()
	return nil
}

func TestWatchBackendV2SendsInitialResponseWithoutEndpoints(t *testing.T) {
	const (
		token = "keyspace-token"
		pool  = "tidb-pool"
	)
	streamCtx, cancel := context.WithCancel(context.Background())
	stream := &watchBackendV2TestStream{
		ctx:    streamCtx,
		cancel: cancel,
	}
	server := &grpcServer{
		ctx:      context.Background(),
		notifier: notifier.New(),
		tManagers: map[string]tidbmanager.TidbManager{
			pool: &clientfake.FakeTidbManager{
				GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
					token: {
						TidbPool:    pool,
						Replicas:    2,
						ModRevision: 42,
					},
				},
			},
		},
	}

	err := server.WatchBackendV2(&pbm.WatchBackendRequestV2{Token: token}, stream)
	require.ErrorIs(t, err, context.Canceled)
	require.Len(t, stream.responses, 1)
	require.Empty(t, stream.responses[0].Backends)
	require.Equal(t, pool, stream.responses[0].TidbPool)
	require.Equal(t, int32(2), stream.responses[0].Replicas)
	require.Equal(t, int64(42), stream.responses[0].ModRevision)
}