package mysqlbackend

import (
	"context"
	"testing"
)

type recordingLeaseProvider struct {
	kinds []ConnLeaseKind
}

func (p *recordingLeaseProvider) AcquireConnLease(kind ConnLeaseKind) (func(), error) {
	p.kinds = append(p.kinds, kind)
	return func() {}, nil
}

func TestAcquireConnLeaseFromContextUsesPreferredKind(t *testing.T) {
	provider := &recordingLeaseProvider{}
	ctx := WithConnLeaseProvider(context.Background(), provider)
	ctx = WithPreferredConnLeaseKind(ctx, ConnLeaseKindBulk)

	release, err := AcquireConnLeaseFromContext(ctx, ConnLeaseKindStandard)
	if err != nil {
		t.Fatalf("AcquireConnLeaseFromContext failed: %v", err)
	}
	if release == nil {
		t.Fatal("expected release func")
	}
	if len(provider.kinds) != 1 || provider.kinds[0] != ConnLeaseKindBulk {
		t.Fatalf("expected preferred bulk lease kind, got %v", provider.kinds)
	}
}

func TestAcquireConnLeaseFromContextSkipsHeldPreferredKind(t *testing.T) {
	provider := &recordingLeaseProvider{}
	ctx := WithConnLeaseProvider(context.Background(), provider)
	ctx = WithPreferredConnLeaseKind(ctx, ConnLeaseKindBulk)
	ctx = WithHeldConnLease(ctx, ConnLeaseKindBulk)

	release, err := AcquireConnLeaseFromContext(ctx, ConnLeaseKindStandard)
	if err != nil {
		t.Fatalf("AcquireConnLeaseFromContext failed: %v", err)
	}
	if release != nil {
		t.Fatal("expected nil release when preferred lease is already held")
	}
	if len(provider.kinds) != 0 {
		t.Fatalf("expected held preferred lease to skip provider acquisition, got %v", provider.kinds)
	}
}