package source

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
	"go.mongodb.org/mongo-driver/v2/mongo/writeconcern"
)

type trackedSource struct {
	database *mongo.Database
	released bool
}

func (s *trackedSource) Acquire(ctx context.Context) (*mongo.Database, context.Context, func(), error) {
	return s.database, ctx, func() { s.released = true }, nil
}

func TestMajorityPreservesSourceLifetime(t *testing.T) {
	client, err := mongo.Connect(options.Client().ApplyURI("mongodb://localhost:27017").SetWriteConcern(&writeconcern.WriteConcern{W: 1}).SetReadPreference(readpref.Secondary()))
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, client.Disconnect(context.Background())) })
	underlying := &trackedSource{database: client.Database("handoff")}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	database, scoped, release, err := Majority(underlying).Acquire(ctx)
	require.NoError(t, err)
	require.Same(t, client, database.Client())
	require.Equal(t, "handoff", database.Name())
	require.Equal(t, ctx, scoped)
	require.False(t, underlying.released)
	release()
	require.True(t, underlying.released)
	_, _, _, err = Majority(Fixed(nil)).Acquire(ctx)
	require.Error(t, err)
}