package server

import (
	"bytes"
	"context"
	"crypto/rand"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/encrypt"
	"github.com/mem9-ai/drive9/pkg/leader"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/tenant"
	"github.com/mem9-ai/drive9/pkg/testtidb"
)

// stallManagedSharedDBPoolWave blocks CreateManagedSharedDBPoolTenantWaveWithResize
// until release is called. TiDB unistore does not honor LOCK TABLES.
func stallManagedSharedDBPoolWave(t *testing.T) (started <-chan struct{}, release func()) {
	t.Helper()
	startedCh := make(chan struct{})
	block := make(chan struct{})
	var once sync.Once
	meta.TestHookBeforeManagedSharedDBPoolWave = func() {
		once.Do(func() { close(startedCh) })
		<-block
	}
	released := false
	release = func() {
		if released {
			return
		}
		released = true
		close(block)
	}
	t.Cleanup(func() {
		meta.TestHookBeforeManagedSharedDBPoolWave = nil
		release()
	})
	return startedCh, release
}

// stallManagedSharedDBPoolWaveCommit blocks after uncommitted wave staging so
// tests can assert in-flight capacity is still invisible to allocators.
func stallManagedSharedDBPoolWaveCommit(t *testing.T) (started <-chan struct{}, release func()) {
	t.Helper()
	startedCh := make(chan struct{})
	block := make(chan struct{})
	var once sync.Once
	meta.TestHookBeforeManagedSharedDBPoolWaveCommit = func() {
		once.Do(func() { close(startedCh) })
		<-block
	}
	released := false
	release = func() {
		if released {
			return
		}
		released = true
		close(block)
	}
	t.Cleanup(func() {
		meta.TestHookBeforeManagedSharedDBPoolWaveCommit = nil
		release()
	})
	return startedCh, release
}

func waitManagedSharedDBPoolWaveStalled(t *testing.T, started <-chan struct{}) {
	t.Helper()
	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("managed shared DB wave did not reach the staging stall")
	}
}

func newFollowerLeaderManager(t *testing.T, metaStore *meta.Store) *leader.Manager {
	t.Helper()
	ctx := context.Background()
	conn, err := metaStore.DB().Conn(ctx)
	if err != nil {
		t.Fatalf("open leader lock connection: %v", err)
	}
	lockName := fmt.Sprintf("d9:test-follower:%d", time.Now().UnixNano())
	var acquired sql.NullInt64
	if err := conn.QueryRowContext(ctx, `SELECT GET_LOCK(?, 0)`, lockName).Scan(&acquired); err != nil {
		_ = conn.Close()
		t.Fatalf("hold leader lock: %v", err)
	}
	if !acquired.Valid || acquired.Int64 != 1 {
		_ = conn.Close()
		t.Fatalf("hold leader lock result = %+v, want 1", acquired)
	}
	t.Cleanup(func() {
		var released sql.NullInt64
		_ = conn.QueryRowContext(context.Background(), `SELECT RELEASE_LOCK(?)`, lockName).Scan(&released)
		_ = conn.Close()
	})
	return leader.NewManager(metaStore.DB(), leader.WithLockName(lockName))
}

func TestRetryTenantPoolClaimCASSucceedsOnEighthAttempt(t *testing.T) {
	attempts := 0
	got, err := retryTenantPoolClaimCAS(func() (int, error) {
		attempts++
		if attempts < tenantPoolClaimCASRetryLimit {
			return 0, meta.ErrNotFound
		}
		return 42, nil
	})
	if err != nil {
		t.Fatalf("retryTenantPoolClaimCAS: %v", err)
	}
	if got != 42 || attempts != tenantPoolClaimCASRetryLimit {
		t.Fatalf("result=%d attempts=%d, want 42/%d", got, attempts, tenantPoolClaimCASRetryLimit)
	}
}

func TestRetryTenantPoolClaimCASReturnsMissAfterLimit(t *testing.T) {
	attempts := 0
	got, err := retryTenantPoolClaimCAS(func() (int, error) {
		attempts++
		return 0, meta.ErrNotFound
	})
	if err != nil {
		t.Fatalf("err = %v, want nil for exhausted claim miss", err)
	}
	if got != 0 {
		t.Fatalf("result = %d, want zero value", got)
	}
	if attempts != tenantPoolClaimCASRetryLimit {
		t.Fatalf("attempts=%d, want %d", attempts, tenantPoolClaimCASRetryLimit)
	}
}

func TestRetryTenantPoolClaimCASDoesNotRetryBusinessError(t *testing.T) {
	wantErr := errors.New("quota headroom exceeded")
	attempts := 0
	_, err := retryTenantPoolClaimCAS(func() (int, error) {
		attempts++
		return 0, wantErr
	})
	if !errors.Is(err, wantErr) {
		t.Fatalf("err = %v, want %v", err, wantErr)
	}
	if attempts != 1 {
		t.Fatalf("attempts=%d, want 1", attempts)
	}
}

func TestFreeTenantPoolClaimRejectsUncountedCandidateAtLimit(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-free-claim-limit")
	ctx := context.Background()
	if err := rt.meta.UpsertTenantTiDBCloudOrgBinding(ctx, &meta.TenantTiDBCloudOrgBinding{
		TenantID: rt.tenantID, OrganizationID: "org-free-claim-limit", ClusterID: "cluster-quota-1",
	}); err != nil {
		t.Fatal(err)
	}
	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if !errors.Is(err, tenant.ErrTiDBCloudFreeTenantLimitReached) {
		t.Fatalf("claim error = %v, want tenant limit", err)
	}
	if res != nil || claimed {
		t.Fatalf("claim result = %+v claimed=%v, want rejection", res, claimed)
	}
	binding, err := rt.meta.GetTenantTiDBCloudOrgBinding(ctx, candidateID)
	if err != nil || binding.PoolStatus != meta.TenantPoolBindingFree {
		t.Fatalf("candidate pool binding = %+v, err=%v, want free", binding, err)
	}
}

func TestFreeTenantPoolClaimWithHeadroomBecomesCounted(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-free-claim-headroom")
	ctx := context.Background()
	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil || !claimed || res == nil || res.TenantID != candidateID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	count, err := rt.meta.CountTiDBCloudOrganizationTenants(ctx, "org-free-claim-headroom")
	if err != nil || count != 1 {
		t.Fatalf("free count = %d, err=%v, want 1", count, err)
	}
	var legacySpendingLimit sql.NullInt64
	if err := rt.meta.DB().QueryRowContext(ctx,
		"SELECT tidbcloud_spending_limit FROM tenant_quota_config WHERE tenant_id = ?", candidateID).
		Scan(&legacySpendingLimit); err != nil {
		t.Fatal(err)
	}
	if !legacySpendingLimit.Valid || legacySpendingLimit.Int64 != 0 {
		t.Fatalf("legacy spending limit = %+v, want 0", legacySpendingLimit)
	}
	cloudOpts := rt.prov.lastOptionsSnapshot()
	if cloudOpts.TiDBCloudSpendingLimitMonthly == nil || *cloudOpts.TiDBCloudSpendingLimitMonthly != 0 {
		t.Fatalf("physical spending limit = %#v, want 0", cloudOpts.TiDBCloudSpendingLimitMonthly)
	}
}

func TestPaidTenantPoolClaimAppliesAndPersistsDedicatedSpendingLimit(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-paid-claim-spending")
	rt.prov.billingFree = false
	ctx := context.Background()
	spendingLimit := int64(12_000)

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"},
		&quotaRequest{quotaFields: quotaFields{TiDBCloudSpendingLimit: &spendingLimit}})
	if err != nil || !claimed || res == nil || res.TenantID != candidateID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	cloudOpts := rt.prov.lastOptionsSnapshot()
	if cloudOpts.TiDBCloudSpendingLimitMonthly == nil || *cloudOpts.TiDBCloudSpendingLimitMonthly != spendingLimit {
		t.Fatalf("physical spending limit = %#v, want %d", cloudOpts.TiDBCloudSpendingLimitMonthly, spendingLimit)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, candidateID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != spendingLimit {
		t.Fatalf("persisted spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, spendingLimit)
	}
}

func TestPaidTenantPoolClaimPersistsProvisionerDefaultSpendingLimit(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-paid-claim-default-spending")
	rt.prov.billingFree = false
	defaultSpendingLimit := int64(5_000)
	rt.prov.cloudCfg = &tenant.QuotaCloudConfig{
		TiDBCloudSpendingLimitMonthly: &defaultSpendingLimit,
	}
	ctx := context.Background()

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil || !claimed || res == nil || res.TenantID != candidateID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, candidateID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != defaultSpendingLimit {
		t.Fatalf("persisted spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, defaultSpendingLimit)
	}
}

func TestFreeTenantPoolClaimDoesNotReapplyQuotaAfterMarkUsed(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-free-claim-single-quota-write")
	ctx := context.Background()
	mediaLimit := int64(12)
	if err := rt.meta.UpsertTenantExtractConfig(ctx, &meta.TenantExtractConfig{
		TenantID: candidateID, MediaType: meta.ExtractMediaTypeImage, Enabled: true,
		APIBase: "https://extract.example.com", APIKeyCipher: []byte("cipher"), Model: "model",
	}); err != nil {
		t.Fatal(err)
	}
	rt.prov.markPoolUsedHook = func() error {
		return rt.meta.UpsertTenantExtractConfig(ctx, &meta.TenantExtractConfig{
			TenantID: candidateID, MediaType: meta.ExtractMediaTypeImage, Enabled: false,
			APIBase: "https://extract.example.com", APIKeyCipher: []byte("cipher"), Model: "model",
		})
	}

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"},
		&quotaRequest{quotaFields: quotaFields{MaxMediaLLMFiles: &mediaLimit}})
	if err != nil || !claimed || res == nil || res.TenantID != candidateID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, candidateID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxMediaLLMFiles != mediaLimit {
		t.Fatalf("max media LLM files = %d, want %d", cfg.MaxMediaLLMFiles, mediaLimit)
	}
}

func TestFreeTenantPoolClaimFailureRollsBackQuotaAndPoolStatus(t *testing.T) {
	rt, candidateID := newFreeNativePoolClaimRuntime(t, "org-free-claim-rollback")
	ctx := context.Background()
	wantErr := errors.New("mark pool used failed")
	rt.prov.markPoolUsedErr = wantErr

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if !errors.Is(err, wantErr) {
		t.Fatalf("claim error = %v, want %v", err, wantErr)
	}
	if res != nil || claimed {
		t.Fatalf("claim result = %+v claimed=%v, want failure", res, claimed)
	}
	var quotaRows int
	if err := rt.meta.DB().QueryRowContext(ctx,
		"SELECT COUNT(*) FROM tenant_quota_config WHERE tenant_id = ?", candidateID).Scan(&quotaRows); err != nil {
		t.Fatal(err)
	}
	if quotaRows != 0 {
		t.Fatalf("quota rows = %d, want 0", quotaRows)
	}
	binding, err := rt.meta.GetTenantTiDBCloudOrgBinding(ctx, candidateID)
	if err != nil || binding.PoolStatus != meta.TenantPoolBindingFree {
		t.Fatalf("pool binding = %+v, err=%v, want free", binding, err)
	}
	count, err := rt.meta.CountTiDBCloudOrganizationTenants(ctx, "org-free-claim-rollback")
	if err != nil || count != 0 {
		t.Fatalf("free count = %d, err=%v, want 0", count, err)
	}
}

func TestFreeSharedTenantPoolClaimBecomesCountedAndHasQuota(t *testing.T) {
	rt, tenantID := newFreeSharedPoolClaimRuntime(t, "org-free-shared-claim")
	rt.server.defaultTenantProvider = tenant.ProviderTiDBCloudNative
	ctx := context.Background()
	requestedSpendingLimit := int64(12_000)

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"},
		&quotaRequest{quotaFields: quotaFields{TiDBCloudSpendingLimit: &requestedSpendingLimit}})
	if err != nil || !claimed || res == nil || res.TenantID != tenantID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	count, err := rt.meta.CountTiDBCloudOrganizationTenants(ctx, "org-free-shared-claim")
	if err != nil || count != 1 {
		t.Fatalf("free count = %d, err=%v, want 1", count, err)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxStorageBytes != DefaultTiDBCloudFreeMaxStorageBytes ||
		cfg.MaxFileSizeBytes != DefaultTiDBCloudFreeMaxFileSizeBytes ||
		cfg.MaxFileCount != DefaultTiDBCloudFreeMaxFileCount {
		t.Fatalf("free shared claim quota = %+v", cfg)
	}
}

func TestSharedTenantPoolClaimLLMQuotaRequiresEnabledTenantExtractConfig(t *testing.T) {
	rt, tenantID := newFreeSharedPoolClaimRuntime(t, "org-shared-claim-llm-config")
	ctx := context.Background()
	mediaLimit := int64(12)

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"},
		&quotaRequest{quotaFields: quotaFields{MaxMediaLLMFiles: &mediaLimit}})
	if !errors.Is(err, meta.ErrTenantExtractConfigNotEnabled) {
		t.Fatalf("claim error = %v, want enabled tenant extract config error", err)
	}
	if res != nil || claimed {
		t.Fatalf("claim result = %+v claimed=%v, want rejection", res, claimed)
	}
	membership, err := rt.meta.GetTenantPoolMembership(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if membership.PoolStatus != meta.TenantPoolBindingFree {
		t.Fatalf("pool status = %s, want free after rejected claim", membership.PoolStatus)
	}
	if err := rt.meta.UpsertTenantExtractConfig(ctx, &meta.TenantExtractConfig{
		TenantID: tenantID, MediaType: meta.ExtractMediaTypeImage, Enabled: true,
		APIBase: "https://extract.example.com", APIKeyCipher: []byte("cipher"), Model: "model",
	}); err != nil {
		t.Fatal(err)
	}
	res, _, claimed, _, err = rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"},
		&quotaRequest{quotaFields: quotaFields{MaxMediaLLMFiles: &mediaLimit}})
	if err != nil || !claimed || res == nil || res.TenantID != tenantID {
		t.Fatalf("claim with enabled extract config = %+v claimed=%v err=%v", res, claimed, err)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxMediaLLMFiles != mediaLimit {
		t.Fatalf("max media LLM files = %d, want %d", cfg.MaxMediaLLMFiles, mediaLimit)
	}
}

func TestMaterializeSharedTenantQuotaLLMLimitRequiresEnabledTenantExtractConfig(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNativeShared)
	rt.server.sharedDBSpendingLimit = 9_000
	ctx := context.Background()
	mediaLimit := int64(12)
	requestedSpendingLimit := int64(12_000)

	err := rt.server.materializeSharedTenantQuota(ctx, rt.tenantID, provisionTenantOptions{
		Quota: &quotaRequest{quotaFields: quotaFields{
			MaxMediaLLMFiles:       &mediaLimit,
			TiDBCloudSpendingLimit: &requestedSpendingLimit,
		}},
	})
	if !errors.Is(err, meta.ErrTenantExtractConfigNotEnabled) {
		t.Fatalf("materialize quota error = %v, want enabled tenant extract config error", err)
	}
	if err := rt.meta.UpsertTenantExtractConfig(ctx, &meta.TenantExtractConfig{
		TenantID: rt.tenantID, MediaType: meta.ExtractMediaTypeImage, Enabled: true,
		APIBase: "https://extract.example.com", APIKeyCipher: []byte("cipher"), Model: "model",
	}); err != nil {
		t.Fatal(err)
	}
	if err := rt.server.materializeSharedTenantQuota(ctx, rt.tenantID, provisionTenantOptions{
		Quota: &quotaRequest{quotaFields: quotaFields{
			MaxMediaLLMFiles:       &mediaLimit,
			TiDBCloudSpendingLimit: &requestedSpendingLimit,
		}},
	}); err != nil {
		t.Fatalf("materialize with enabled extract config: %v", err)
	}
	cfg, err := rt.meta.GetQuotaConfig(ctx, rt.tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxMediaLLMFiles != mediaLimit {
		t.Fatalf("max media LLM files = %d, want %d", cfg.MaxMediaLLMFiles, mediaLimit)
	}
	if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != 9_000 {
		t.Fatalf("shared compatibility spending limit = %#v, want 9000", cfg.TiDBCloudSpendingLimit)
	}
}

func TestSharedTenantQuotaConfigAppliesRequestedFileLimits(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNativeShared)
	maxStorageSize := int64(2048)
	maxFileSize := int64(64)
	maxFileCount := int64(123)

	cfg, err := rt.server.sharedTenantQuotaConfig(rt.tenantID, provisionTenantOptions{
		Quota: &quotaRequest{quotaFields: quotaFields{
			MaxStorageSize: &maxStorageSize,
			MaxFileSize:    &maxFileSize,
			MaxFileCount:   &maxFileCount,
		}},
	})
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxStorageBytes != maxStorageSize*quotaStorageSizeBytes ||
		cfg.MaxFileSizeBytes != maxFileSize*quotaStorageSizeBytes ||
		cfg.MaxFileCount != maxFileCount {
		t.Fatalf("shared quota config = %+v, want requested storage/file limits", cfg)
	}
}

func TestSharedTenantQuotaConfigKeepsFreeCompatibilityShadow(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNativeShared)

	cfg, err := rt.server.sharedTenantQuotaConfig(rt.tenantID, provisionTenantOptions{
		TiDBCloudAccess: &tiDBCloudAccessProfile{IsFree: true},
	})
	if err != nil {
		t.Fatal(err)
	}
	if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != 0 {
		t.Fatalf("free shared compatibility spending limit = %#v, want 0", cfg.TiDBCloudSpendingLimit)
	}
}

func TestPaidSharedTenantPoolClaimKeepsLegacyQuotaMaterialized(t *testing.T) {
	rt, tenantID := newFreeSharedPoolClaimRuntime(t, "org-paid-shared-compat")
	rt.prov.billingFree = false
	ctx := context.Background()

	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil || !claimed || res == nil || res.TenantID != tenantID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}
	var spendingLimit sql.NullInt64
	if err := rt.meta.DB().QueryRowContext(ctx,
		`SELECT tidbcloud_spending_limit FROM tenant_quota_config WHERE tenant_id = ?`,
		tenantID).Scan(&spendingLimit); err != nil {
		t.Fatal(err)
	}
	_, wantSpendingLimit := rt.server.managedSharedDBPolicy()
	if !spendingLimit.Valid || spendingLimit.Int64 != wantSpendingLimit {
		t.Fatalf("legacy shared spending limit = %+v, want %d for old-pod compatibility", spendingLimit, wantSpendingLimit)
	}
}

func TestFailedWaveCleanupDoesNotFailClaimedSharedTenant(t *testing.T) {
	rt, tenantID := newFreeSharedPoolClaimRuntime(t, "org-claimed-wave-cleanup")
	ctx := context.Background()
	cred := tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}
	res, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx, cred, nil)
	if err != nil || !claimed || res == nil || res.TenantID != tenantID {
		t.Fatalf("claim result = %+v claimed=%v err=%v", res, claimed, err)
	}

	rt.server.cleanupCreatedPoolTenants(ctx, []*provisionTenantResult{{
		TenantID: tenantID, Provider: tenant.ProviderTiDBCloudNativeShared,
	}}, cred, "concurrent_claim_test")

	tenantRow, err := rt.meta.GetTenant(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if tenantRow.Status != meta.TenantActive {
		t.Fatalf("claimed tenant status after wave cleanup = %s, want %s", tenantRow.Status, meta.TenantActive)
	}
	membership, err := rt.meta.GetTenantPoolMembership(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if membership.PoolStatus != meta.TenantPoolBindingUsed {
		t.Fatalf("claimed membership status after wave cleanup = %s, want %s",
			membership.PoolStatus, meta.TenantPoolBindingUsed)
	}
}

func newFreeSharedPoolClaimRuntime(t *testing.T, organizationID string) (*quotaRuntime, string) {
	t.Helper()
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	rt.server.defaultTenantProvider = tenant.ProviderTiDBCloudNativeShared
	rt.prov.billingFree = true
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: organizationID, Role: tenant.TiDBCloudRoleOrgOwner,
	}}
	ctx := context.Background()
	now := time.Now().UTC()
	poolID := "pool-" + organizationID
	tenantID := "tenant-" + organizationID
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: poolID, OrganizationID: organizationID, Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
	passwordCipher, err := rt.server.pool.Encrypt(ctx, []byte("shared-pass"))
	if err != nil {
		t.Fatal(err)
	}
	dbID, err := rt.meta.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: organizationID, Host: "shared.example.com", Port: 4000,
		User: "root", PasswordCipher: passwordCipher, Name: "shared_db", MaxTenants: 10,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := rt.server.insertPendingPoolTenant(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared, now); err != nil {
		t.Fatal(err)
	}
	if err := rt.server.materializeSharedTenantQuota(ctx, tenantID, provisionTenantOptions{}); err != nil {
		t.Fatal(err)
	}
	fsID, err := rt.meta.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if err := rt.meta.CompleteSharedTenantPoolMember(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared,
			SchemaShape: meta.SchemaShapeShared, Status: meta.SharedDBStatusActive},
		&meta.TenantPoolMembership{TenantID: tenantID, TiDBCloudOrganizationID: organizationID,
			PoolID: poolID, PoolStatus: meta.TenantPoolBindingFree, CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatal(err)
	}
	return rt, tenantID
}

func TestSharedTenantPoolWaveMaterializesLegacySpendingLimit(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	rt.server.defaultTenantProvider = tenant.ProviderTiDBCloudNativeShared
	rt.server.sharedDBSpendingLimit = 23_000
	rt.prov.cloudProvider = "aws"
	rt.prov.region = "us-east-1"
	ctx := context.Background()
	now := time.Now().UTC()
	const (
		organizationID = "org-shared-wave-legacy-spending"
		poolID         = "pool-shared-wave-legacy-spending"
	)
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: poolID, OrganizationID: organizationID, Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}

	partitions, err := rt.server.stageSharedPoolTenantWave(ctx, poolID, organizationID,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, 1, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(partitions) != 1 || len(partitions[0].tenantIDs) != 1 {
		t.Fatalf("partitions = %+v, want one tenant", partitions)
	}
	var spendingLimit sql.NullInt64
	if err := rt.meta.DB().QueryRowContext(ctx,
		`SELECT tidbcloud_spending_limit FROM tenant_quota_config WHERE tenant_id = ?`,
		partitions[0].tenantIDs[0]).Scan(&spendingLimit); err != nil {
		t.Fatal(err)
	}
	if !spendingLimit.Valid || spendingLimit.Int64 != 23_000 {
		t.Fatalf("legacy spending limit = %+v, want 23000", spendingLimit)
	}
}

func newFreeNativePoolClaimRuntime(t *testing.T, organizationID string) (*quotaRuntime, string) {
	t.Helper()
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	rt.prov.billingFree = true
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: organizationID, Role: tenant.TiDBCloudRoleOrgOwner,
	}}
	rt.server.tidbCloudFreePlanLimits = TiDBCloudFreePlanLimits{
		TenantCount:      1,
		MaxStorageBytes:  DefaultTiDBCloudFreeMaxStorageBytes,
		MaxFileSizeBytes: DefaultTiDBCloudFreeMaxFileSizeBytes,
		MaxFileCount:     DefaultTiDBCloudFreeMaxFileCount,
	}
	ctx := context.Background()
	now := time.Now().UTC()
	poolID := "pool-" + organizationID
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: poolID, OrganizationID: organizationID, Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
	return rt, insertAdminPoolFreeTenant(t, rt, poolID, organizationID, 1)
}

func TestAdminTenantCreatePoolClaimMapsQuotaValidationErrors(t *testing.T) {
	for i, tt := range []struct {
		name       string
		free       bool
		body       map[string]any
		wantStatus int
		wantError  error
	}{
		{
			name: "free quota exceeds plan",
			free: true,
			body: map[string]any{
				"max_file_count": DefaultTiDBCloudFreeMaxFileCount + 1,
			},
			wantStatus: http.StatusPaymentRequired,
			wantError:  tenant.ErrTiDBCloudFreeQuotaExceeded,
		},
		{
			name: "free spending limit is forbidden",
			free: true,
			body: map[string]any{
				"tidbcloud_spending_limit": 100,
			},
			wantStatus: http.StatusPaymentRequired,
			wantError:  tenant.ErrTiDBCloudFreeSpendingLimitForbidden,
		},
		{
			name: "paid spending limit is out of range",
			body: map[string]any{
				"tidbcloud_spending_limit": 1,
			},
			wantStatus: http.StatusBadRequest,
			wantError:  errTiDBCloudSpendingLimitInvalid,
		},
	} {
		t.Run(tt.name, func(t *testing.T) {
			rt, _ := newFreeNativePoolClaimRuntime(t, fmt.Sprintf("org-claim-err-%d", i))
			rt.prov.billingFree = tt.free
			ts := httptest.NewServer(rt.server)
			t.Cleanup(ts.Close)
			for key, value := range map[string]any{
				"public_key":  "public-1",
				"private_key": "private-1",
			} {
				tt.body[key] = value
			}
			resp := postJSON(t, ts.URL+"/v1/admin/tenants", tt.body, "")
			defer func() { _ = resp.Body.Close() }()
			raw, err := io.ReadAll(resp.Body)
			if err != nil {
				t.Fatal(err)
			}
			if resp.StatusCode != tt.wantStatus {
				t.Fatalf("status = %d body=%s, want %d", resp.StatusCode, raw, tt.wantStatus)
			}
			if !strings.Contains(string(raw), tt.wantError.Error()) {
				t.Fatalf("body = %s, want error %q", raw, tt.wantError)
			}
		})
	}
}

func TestTenantPoolClaimUsesNativeInventoryBeforeExternalSharedPool(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	ctx := context.Background()
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: "org-mixed-inventory", Role: tenant.TiDBCloudRoleOrgOwner,
	}}
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: "pool-mixed-inventory", OrganizationID: "org-mixed-inventory", Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	nativeTenantID := insertAdminPoolFreeTenant(t, rt, "pool-mixed-inventory", "org-mixed-inventory", 1)
	if _, err := rt.meta.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-mixed-inventory", Host: "shared.example.com", Port: 4000,
		User: "root", PasswordCipher: []byte("cipher"), Name: "shared_db", MaxTenants: 100,
	}); err != nil {
		t.Fatalf("RegisterSharedDB: %v", err)
	}

	res, pool, claimed, sharedPoolMatched, err := rt.server.claimAdminTenantFromPool(ctx,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil {
		t.Fatalf("claimAdminTenantFromPool: %v", err)
	}
	if !claimed || sharedPoolMatched || pool == nil || res == nil {
		t.Fatalf("claim = result=%+v pool=%+v claimed=%v sharedMatched=%v", res, pool, claimed, sharedPoolMatched)
	}
	if res.TenantID != nativeTenantID || res.Provider != tenant.ProviderTiDBCloudNative {
		t.Fatalf("claimed result = %+v, want native tenant %s", res, nativeTenantID)
	}
	deadline := time.Now().Add(time.Second)
	for {
		if _, ok := rt.server.tenantPoolReplenishJobs.Load("pool-mixed-inventory"); ok {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("native request did not trigger request-scoped tenant-pool replenishment")
		}
		time.Sleep(time.Millisecond)
	}
}

func TestTenantPoolClaimConsumesMixedInventoryInGlobalAgeOrder(t *testing.T) {
	rt := newQuotaRuntimeWithOptions(t, tenant.ProviderTiDBCloudNative, quotaRuntimeOptions{
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared,
	})
	ctx := context.Background()
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: "org-mixed-age", Role: tenant.TiDBCloudRoleProjectOwner,
	}}
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: "pool-mixed-age", OrganizationID: "org-mixed-age", Size: 2,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	nativeTenantID := insertAdminPoolFreeTenant(t, rt, "pool-mixed-age", "org-mixed-age", 1)
	passwordCipher, err := rt.server.pool.Encrypt(ctx, []byte("shared-pass"))
	if err != nil {
		t.Fatal(err)
	}
	dbID, err := rt.meta.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-mixed-age", Host: "shared.example.com", Port: 4000,
		User: "root", PasswordCipher: passwordCipher, Name: "shared_db", MaxTenants: 10,
	})
	if err != nil {
		t.Fatalf("RegisterSharedDB: %v", err)
	}
	sharedTenantID := "pool-mixed-age-shared"
	sharedCreatedAt := now.Add(time.Hour)
	if err := rt.server.insertPendingPoolTenant(ctx, sharedTenantID, tenant.ProviderTiDBCloudNativeShared, sharedCreatedAt); err != nil {
		t.Fatalf("insertPendingPoolTenant: %v", err)
	}
	if err := rt.server.materializeSharedTenantQuota(ctx, sharedTenantID, provisionTenantOptions{}); err != nil {
		t.Fatalf("materializeSharedTenantQuota: %v", err)
	}
	fsID, err := rt.meta.EnsureFsID(ctx, sharedTenantID)
	if err != nil {
		t.Fatalf("EnsureFsID: %v", err)
	}
	if err := rt.meta.CompleteSharedTenantPoolMember(ctx, sharedTenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared,
			SchemaShape: meta.SchemaShapeShared, Status: meta.SharedDBStatusActive},
		&meta.TenantPoolMembership{TenantID: sharedTenantID, TiDBCloudOrganizationID: "org-mixed-age",
			PoolID: "pool-mixed-age", PoolStatus: meta.TenantPoolBindingFree,
			CreatedAt: sharedCreatedAt, UpdatedAt: sharedCreatedAt}); err != nil {
		t.Fatalf("CompleteSharedTenantPoolMember: %v", err)
	}
	free, err := rt.meta.CountFreeTenantPoolBindings(ctx, "org-mixed-age")
	if err != nil {
		t.Fatalf("CountFreeTenantPoolBindings: %v", err)
	}
	if free != 2 {
		t.Fatalf("mixed native/shared free inventory = %d, want 2", free)
	}

	cred := tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}
	first, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx, cred, nil)
	if err != nil || !claimed {
		t.Fatalf("first claim = %+v, claimed=%v, err=%v", first, claimed, err)
	}
	if first.TenantID != nativeTenantID || first.Provider != tenant.ProviderTiDBCloudNative {
		t.Fatalf("first claim = %+v, want older native tenant %s", first, nativeTenantID)
	}
	if _, ok := rt.server.tenantPoolReplenishJobs.Load("pool-mixed-age"); ok {
		t.Fatal("shared-default request unexpectedly triggered bulk tenant-pool replenishment")
	}
	second, _, claimed, _, err := rt.server.claimAdminTenantFromPool(ctx, cred, nil)
	if err != nil || !claimed {
		t.Fatalf("second claim = %+v, claimed=%v, err=%v", second, claimed, err)
	}
	if second.TenantID != sharedTenantID || second.Provider != tenant.ProviderTiDBCloudNativeShared {
		t.Fatalf("second claim = %+v, want shared tenant %s", second, sharedTenantID)
	}
	if second.TenantDSN != "" {
		t.Fatalf("shared claim TenantDSN = %q, want empty", second.TenantDSN)
	}
}

func TestSharedTenantPoolRefillPlansTenPoolsInOneBatch(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	workerCtx, cancel := context.WithCancel(context.Background())
	srv := &Server{
		meta: metaStore, pool: poolManager, provisioner: prov,
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared,
		sharedDBMaxTenants:    1, managedSharedDBCloudBatchSize: 10,
		forkWorkerCtx: workerCtx, forkWorkerCancel: cancel,
	}
	defer srv.Close()
	now := time.Now().UTC()
	if err := metaStore.CreateTenantPool(context.Background(), &meta.TenantPool{PoolID: "pool-shared-10",
		OrganizationID: "org-shared-ten-pools", Size: 10, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatal(err)
	}
	results, err := srv.createFreePoolTenants(context.Background(), "pool-shared-10", 10,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil {
		t.Fatalf("createFreePoolTenants: %v", err)
	}
	if len(results) != 10 {
		t.Fatalf("results = %d, want 10", len(results))
	}
	for _, result := range results {
		if result.Status != meta.TenantPending {
			t.Fatalf("tenant %s status = %q, want pending while connection metadata is incomplete", result.TenantID, result.Status)
		}
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 1 {
		t.Fatalf("shared batch calls = %d, want 1", got)
	}
	if got := prov.sharedPoolBatchMembers.Load(); got != 10 {
		t.Fatalf("shared batch members = %d, want 10", got)
	}
	var managedPools int
	if err := metaStore.DB().QueryRowContext(context.Background(), `SELECT COUNT(*) FROM db_pool
		WHERE org_id = ? AND status = ?`, "org-shared-ten-pools", meta.SharedDBStatusPending).Scan(&managedPools); err != nil {
		t.Fatal(err)
	}
	if managedPools != 10 {
		t.Fatalf("managed pools = %d, want 10", managedPools)
	}
	slots, err := metaStore.CountTenantPoolFreeSlots(context.Background(), "org-shared-ten-pools")
	if err != nil {
		t.Fatal(err)
	}
	if slots != 10 {
		t.Fatalf("free slots = %d, want 10", slots)
	}
}

func TestSharedTenantPoolRefillReusesExistingPhysicalCapacity(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(poolManager.Close)
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 100,
		TokenSecret: make([]byte, 32), Leader: newFollowerLeaderManager(t, metaStore)})
	t.Cleanup(srv.Close)
	ctx := context.Background()
	now := time.Now().UTC()
	logicalPool := &meta.TenantPool{PoolID: "pool-reuse-existing", OrganizationID: "org-reuse-existing",
		Size: 15, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	passwordCipher, err := poolManager.Encrypt(ctx, []byte("shared-pass"))
	if err != nil {
		t.Fatalf("encrypt shared password: %v", err)
	}
	dbID, err := metaStore.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: logicalPool.OrganizationID,
		Host:                    "shared.example.com",
		Port:                    4000,
		User:                    "root",
		PasswordCipher:          passwordCipher,
		Name:                    "tidbcloud_fs",
		MaxTenants:              100,
	})
	if err != nil {
		t.Fatalf("RegisterSharedDB: %v", err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET tenant_count = 9 WHERE db_id = ?`, dbID); err != nil {
		t.Fatalf("seed tenant_count: %v", err)
	}

	results, err := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, 1, tenant.CredentialProvisionRequest{}, nil)
	if err != nil {
		t.Fatalf("createFreeSharedPoolTenants: %v", err)
	}
	if len(results) != 1 || results[0].Status != meta.TenantActive {
		t.Fatalf("results = %+v, want one active tenant", results)
	}
	var physicalPools, tenantCount int
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(tenant_count), 0)
		FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalPools, &tenantCount); err != nil {
		t.Fatalf("load physical pool inventory: %v", err)
	}
	if physicalPools != 1 || tenantCount != 10 {
		t.Fatalf("physical pools = %d, tenant count = %d, want 1 and 10", physicalPools, tenantCount)
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 0 {
		t.Fatalf("shared Cloud batch calls = %d, want 0 while existing capacity remains", got)
	}

	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET tenant_count = 93 WHERE db_id = ?`, dbID); err != nil {
		t.Fatalf("move existing pool below capacity: %v", err)
	}
	_, err = srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, 5, tenant.CredentialProvisionRequest{}, nil)
	if err != nil {
		t.Fatalf("createFreeSharedPoolTenants below capacity: %v", err)
	}
	var softCapReached bool
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT tenant_count, soft_cap_reached FROM db_pool WHERE db_id = ?`, dbID).
		Scan(&tenantCount, &softCapReached); err != nil {
		t.Fatalf("load remaining physical capacity: %v", err)
	}
	if tenantCount != 98 || softCapReached {
		t.Fatalf("tenant count = %d, soft cap reached = %t, want 98 and false", tenantCount, softCapReached)
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 0 {
		t.Fatalf("shared Cloud batch calls = %d, want 0 before physical capacity is full", got)
	}

	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET tenant_count = 99 WHERE db_id = ?`, dbID); err != nil {
		t.Fatalf("move existing pool near capacity: %v", err)
	}
	results, err = srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, 3, tenant.CredentialProvisionRequest{}, nil)
	if err != nil {
		t.Fatalf("createFreeSharedPoolTenants across capacity boundary: %v", err)
	}
	active, pending := 0, 0
	for _, result := range results {
		switch result.Status {
		case meta.TenantActive:
			active++
		case meta.TenantPending:
			pending++
		}
	}
	if active != 1 || pending != 2 {
		t.Fatalf("result statuses: active=%d pending=%d, want 1 and 2", active, pending)
	}
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(tenant_count), 0)
		FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalPools, &tenantCount); err != nil {
		t.Fatalf("load physical pool inventory after spill: %v", err)
	}
	if physicalPools != 2 || tenantCount != 102 {
		t.Fatalf("physical pools = %d, tenant count = %d, want 2 and 102", physicalPools, tenantCount)
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 1 {
		t.Fatalf("shared Cloud batch calls = %d, want 1 for the capacity remainder", got)
	}
	if got := prov.sharedPoolBatchMembers.Load(); got != 1 {
		t.Fatalf("shared Cloud batch physical pools = %d, want 1 for the capacity remainder", got)
	}

	start := make(chan struct{})
	errCh := make(chan error, 2)
	var waveWG sync.WaitGroup
	for i := 0; i < 2; i++ {
		waveWG.Add(1)
		go func() {
			defer waveWG.Done()
			<-start
			waveResults, waveErr := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, 3, tenant.CredentialProvisionRequest{}, nil)
			if waveErr == nil {
				for _, result := range waveResults {
					if result.Status != meta.TenantPending {
						waveErr = fmt.Errorf("concurrent wave tenant %s status = %q, want pending", result.TenantID, result.Status)
						break
					}
				}
			}
			errCh <- waveErr
		}()
	}
	close(start)
	waveWG.Wait()
	close(errCh)
	for waveErr := range errCh {
		if waveErr != nil {
			t.Fatalf("concurrent refill wave: %v", waveErr)
		}
	}
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(tenant_count), 0)
		FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalPools, &tenantCount); err != nil {
		t.Fatalf("load physical pool inventory after concurrent waves: %v", err)
	}
	if physicalPools != 2 || tenantCount != 108 {
		t.Fatalf("physical pools = %d, tenant count = %d, want 2 and 108 after concurrent waves", physicalPools, tenantCount)
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 1 {
		t.Fatalf("shared Cloud batch calls = %d, want no extra create from concurrent reuse", got)
	}
}

func TestSharedTenantPoolRefillMakesWaveVisibleOnlyAfterMembershipReservation(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	now := time.Now().UTC()
	if err := metaStore.CreateTenantPool(context.Background(), &meta.TenantPool{
		PoolID: "pool-stage-wave", OrganizationID: "org-stage-wave", Size: 6,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
	batchStarted := make(chan struct{}, 1)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		sharedPoolBatchStarted: batchStarted}
	workerCtx, workerCancel := context.WithCancel(context.Background())
	srv := &Server{meta: metaStore, pool: poolManager, provisioner: prov,
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, sharedDBMaxTenants: 2,
		managedSharedDBCloudBatchSize: 10, forkWorkerCtx: workerCtx, forkWorkerCancel: workerCancel}

	started, releaseWave := stallManagedSharedDBPoolWaveCommit(t)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	finished := false
	go func() {
		_, err := srv.createFreeSharedPoolTenants(ctx, "pool-stage-wave", 6,
			tenant.CredentialProvisionRequest{}, nil)
		done <- err
	}()
	t.Cleanup(func() {
		cancel()
		releaseWave()
		if !finished {
			select {
			case <-done:
			case <-time.After(5 * time.Second):
				t.Error("shared refill did not stop after releasing wave stall")
			}
		}
		srv.Close()
	})

	waitManagedSharedDBPoolWaveStalled(t, started)
	select {
	case <-batchStarted:
		t.Fatal("Cloud batch started before tenant membership reservations committed")
	default:
	}
	if _, err := metaStore.FindSharedDBForAllocation(context.Background(), "org-stage-wave"); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("direct allocation saw refill-staged capacity before membership commit: %v", err)
	}
	releaseWave()
	select {
	case <-batchStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("Cloud batch did not start after durable wave staging committed")
	}
	select {
	case err := <-done:
		finished = true
		if err != nil {
			t.Fatalf("createFreeSharedPoolTenants: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("shared refill did not finish after durable wave staging")
	}
	free, err := metaStore.CountTenantPoolFreeSlots(context.Background(), "org-stage-wave")
	if err != nil {
		t.Fatal(err)
	}
	if free != 6 {
		t.Fatalf("durably reserved free slots = %d, want 6", free)
	}
	if _, err := metaStore.FindSharedDBForAllocation(context.Background(), "org-stage-wave"); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("direct allocation saw capacity already reserved by the completed refill wave: %v", err)
	}
	active, err := metaStore.CountFreeTenantPoolBindings(context.Background(), "org-stage-wave")
	if err != nil {
		t.Fatal(err)
	}
	planned, err := metaStore.CountTenantPoolPlannedSlots(context.Background(), "org-stage-wave")
	if err != nil {
		t.Fatal(err)
	}
	if active+planned != 6 {
		t.Fatalf("active + planned slots after atomic wave commit = %d + %d, want 6", active, planned)
	}
}

func TestSharedTenantPoolRefillDoesNotCommitAfterLogicalPoolDelete(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(poolManager.Close)
	poolManager.SetMetaStore(metaStore)
	ctx := context.Background()
	logicalPool := &meta.TenantPool{PoolID: "pool-delete-during-wave", OrganizationID: "org-delete-during-wave",
		Size: 2, Status: meta.TenantPoolActive, CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC()}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatal(err)
	}
	batchStarted := make(chan struct{}, 1)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		sharedPoolBatchStarted: batchStarted}
	workerCtx, workerCancel := context.WithCancel(context.Background())
	srv := &Server{meta: metaStore, pool: poolManager, provisioner: prov,
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, sharedDBMaxTenants: 1,
		managedSharedDBCloudBatchSize: 10, forkWorkerCtx: workerCtx, forkWorkerCancel: workerCancel}
	t.Cleanup(srv.Close)

	started, releaseWave := stallManagedSharedDBPoolWave(t)
	waveDone := make(chan error, 1)
	go func() {
		_, createErr := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, logicalPool.Size,
			tenant.CredentialProvisionRequest{}, nil)
		waveDone <- createErr
	}()
	waitManagedSharedDBPoolWaveStalled(t, started)
	select {
	case err := <-waveDone:
		t.Fatalf("refill returned before wave stall was released: %v", err)
	default:
	}
	if err := metaStore.UpdateTenantPoolStatus(ctx, logicalPool.PoolID, meta.TenantPoolDeleting); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.DeleteTenantPoolAndDetachUsedMembers(ctx, logicalPool.PoolID); err != nil {
		t.Fatalf("delete logical pool while refill is blocked: %v", err)
	}
	releaseWave()
	select {
	case err := <-waveDone:
		if err == nil {
			t.Fatal("refill committed after logical pool deletion")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("refill did not stop after logical pool deletion")
	}
	select {
	case <-batchStarted:
		t.Fatal("Cloud batch started after logical pool deletion won")
	default:
	}
	if _, err := metaStore.GetTenantPoolByID(ctx, logicalPool.PoolID); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("logical pool after delete = %v, want not found", err)
	}
	var physicalRows, memberships int
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalRows); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM tenant_pool_memberships WHERE pool_id = ?`, logicalPool.PoolID).Scan(&memberships); err != nil {
		t.Fatal(err)
	}
	if physicalRows != 0 || memberships != 0 {
		t.Fatalf("durable refill residue after delete: physical=%d memberships=%d", physicalRows, memberships)
	}
}

func TestSharedTenantPoolRefillDoesNotOverfillAfterLogicalPoolShrink(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(poolManager.Close)
	poolManager.SetMetaStore(metaStore)
	ctx := context.Background()
	logicalPool := &meta.TenantPool{PoolID: "pool-shrink-during-wave", OrganizationID: "org-shrink-during-wave",
		Size: 2, Status: meta.TenantPoolActive, CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC()}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatal(err)
	}
	batchStarted := make(chan struct{}, 2)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		sharedPoolBatchStarted: batchStarted}
	workerCtx, workerCancel := context.WithCancel(context.Background())
	srv := &Server{meta: metaStore, pool: poolManager, provisioner: prov,
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, sharedDBMaxTenants: 1,
		managedSharedDBCloudBatchSize: 10, forkWorkerCtx: workerCtx, forkWorkerCancel: workerCancel}
	t.Cleanup(srv.Close)

	started, releaseWave := stallManagedSharedDBPoolWave(t)
	waveDone := make(chan error, 1)
	go func() {
		_, createErr := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, logicalPool.Size,
			tenant.CredentialProvisionRequest{}, nil)
		waveDone <- createErr
	}()
	waitManagedSharedDBPoolWaveStalled(t, started)
	select {
	case err := <-waveDone:
		t.Fatalf("refill returned before wave stall was released: %v", err)
	default:
	}
	if err := metaStore.UpdateTenantPoolSize(ctx, logicalPool.PoolID, 1); err != nil {
		t.Fatalf("shrink logical pool while refill is blocked: %v", err)
	}
	releaseWave()
	select {
	case err := <-waveDone:
		if err == nil {
			t.Fatal("stale refill committed after logical pool shrink")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("stale refill did not stop after logical pool shrink")
	}
	select {
	case <-batchStarted:
		t.Fatal("Cloud batch started for stale oversized refill")
	default:
	}
	gotPool, err := metaStore.GetTenantPoolByID(ctx, logicalPool.PoolID)
	if err != nil {
		t.Fatal(err)
	}
	if gotPool.Size != 1 {
		t.Fatalf("logical pool size = %d, want 1", gotPool.Size)
	}
	if free, err := metaStore.CountTenantPoolFreeSlots(ctx, logicalPool.OrganizationID); err != nil || free != 0 {
		t.Fatalf("free slots after rejected stale refill = %d, err=%v", free, err)
	}
	if _, err := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, 1,
		tenant.CredentialProvisionRequest{}, nil); err != nil {
		t.Fatalf("correctly sized refill after shrink: %v", err)
	}
	select {
	case <-batchStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("Cloud batch did not start for correctly sized refill")
	}
	if free, err := metaStore.CountTenantPoolFreeSlots(ctx, logicalPool.OrganizationID); err != nil || free != 1 {
		t.Fatalf("free slots after correctly sized refill = %d, err=%v", free, err)
	}
}

func TestSharedTenantPoolWaveIgnoresOrphanPendingMembershipCapacity(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(poolManager.Close)
	poolManager.SetMetaStore(metaStore)
	ctx := context.Background()
	now := time.Now().UTC()
	logicalPool := &meta.TenantPool{
		PoolID: "pool-orphan-pending", OrganizationID: "org-orphan-pending", Size: 2,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	prov := &fakeProvisioner{
		provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		identityOrg: logicalPool.OrganizationID,
	}
	srv := NewWithConfig(Config{
		Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 1,
		TokenSecret: make([]byte, 32), Leader: newFollowerLeaderManager(t, metaStore),
	})
	t.Cleanup(srv.Close)

	const orphanTenantID = "tenant-orphan-pending-membership"
	if err := srv.insertPendingPoolTenant(ctx, orphanTenantID, tenant.ProviderTiDBCloudNativeShared, now); err != nil {
		t.Fatalf("insert orphan tenant: %v", err)
	}
	if err := metaStore.UpsertTenantPoolMembership(ctx, &meta.TenantPoolMembership{
		TenantID: orphanTenantID, TiDBCloudOrganizationID: logicalPool.OrganizationID,
		PoolID: logicalPool.PoolID, PoolStatus: meta.TenantPoolBindingFree,
		CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatalf("insert orphan membership: %v", err)
	}

	results, err := srv.createFreeSharedPoolTenants(ctx, logicalPool.PoolID, logicalPool.Size,
		tenant.CredentialProvisionRequest{}, nil)
	if err != nil {
		t.Fatalf("create valid shared wave: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("valid wave members = %d, want 2", len(results))
	}
	if got := prov.sharedPoolBatchMembers.Load(); got != 2 {
		t.Fatalf("physical Cloud creates = %d, want 2", got)
	}
	var memberships int
	if err := metaStore.DB().QueryRowContext(ctx,
		`SELECT COUNT(*) FROM tenant_pool_memberships WHERE pool_id = ?`, logicalPool.PoolID).Scan(&memberships); err != nil {
		t.Fatal(err)
	}
	if memberships != 3 {
		t.Fatalf("memberships = %d, want orphan plus two valid members", memberships)
	}
	var profiles int
	if err := metaStore.DB().QueryRowContext(ctx,
		`SELECT COUNT(*) FROM tenant_auto_embedding_profiles p
		 JOIN tenant_pool_memberships m ON m.tenant_id = p.tenant_id
		 WHERE m.pool_id = ?`, logicalPool.PoolID).Scan(&profiles); err != nil {
		t.Fatal(err)
	}
	if profiles != 0 {
		t.Fatalf("shared wave wrote %d tenant_auto_embedding_profiles rows, want 0", profiles)
	}
}

func TestSharedTenantPoolRefillPlansWithServerOwnedSharedCredential(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		sharedPoolBatchErr: errors.New("stop after managed pool planning")}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, TokenSecret: make([]byte, 32),
		Leader: newFollowerLeaderManager(t, metaStore)})
	defer srv.Close()
	now := time.Now().UTC()
	if err := metaStore.CreateTenantPool(context.Background(), &meta.TenantPool{PoolID: "pool-shared-credential",
		OrganizationID: "org-shared", Size: 1, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatal(err)
	}
	_, err = srv.createFreeSharedPoolTenants(context.Background(), "pool-shared-credential", 1,
		tenant.CredentialProvisionRequest{PublicKey: "customer-public", PrivateKey: "customer-private"}, nil)
	if err == nil {
		t.Fatal("createFreeSharedPoolTenants succeeded, want forced batch failure")
	}
	rows, err := metaStore.ListSharedDBsByStatus(context.Background(), meta.SharedDBStatusPending, 10)
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("managed pools = %d, want 1", len(rows))
	}
	wantProvisioningKey := sharedDBProvisioningKey(tenant.CredentialProvisionRequest{PublicKey: "shared-public"})
	if !bytes.Equal(rows[0].ProvisioningKey, wantProvisioningKey) {
		t.Fatalf("managed pool provisioning key = %x, want server-owned shared credential key %x",
			rows[0].ProvisioningKey, wantProvisioningKey)
	}
	if got := prov.iamCalls.Load(); got != 0 {
		t.Fatalf("identity lookups = %d, want 0", got)
	}
}

func TestSharedTenantPoolRefillReusesCapacityWithoutIdentityLookup(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	passwordCipher, err := poolManager.Encrypt(context.Background(), []byte("shared-pass"))
	if err != nil {
		t.Fatal(err)
	}
	existingDBID, err := metaStore.RegisterSharedDB(context.Background(), &meta.SharedDB{
		TiDBCloudOrganizationID: "org-shared-existing", Host: "shared.example.com", Port: 4000,
		User: "root", PasswordCipher: passwordCipher, Name: "tidbcloud_fs", MaxTenants: 100,
	})
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC()
	if err := metaStore.CreateTenantPool(context.Background(), &meta.TenantPool{
		PoolID: "pool-shared-existing", OrganizationID: "org-shared-existing", Size: 50,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		identityOrg: "org-shared-existing"}
	srv := &Server{meta: metaStore, pool: poolManager, provisioner: prov,
		tidbCloudRBACCache: newTiDBCloudRBACCache(time.Hour)}
	results, err := srv.createFreeSharedPoolTenants(context.Background(), "pool-shared-existing", 50,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil {
		t.Fatalf("createFreeSharedPoolTenants: %v", err)
	}
	if len(results) != 50 {
		t.Fatalf("results = %d, want 50", len(results))
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 0 {
		t.Fatalf("physical batch calls = %d, want no Cloud create while existing capacity remains", got)
	}
	if got := prov.iamCalls.Load(); got != 0 {
		t.Fatalf("identity lookups = %d, want 0 when tenant pool organization is known", got)
	}
	var existingTenantCount, physicalPoolCount int
	if err := metaStore.DB().QueryRowContext(context.Background(),
		`SELECT tenant_count FROM db_pool WHERE db_id = ?`, existingDBID).Scan(&existingTenantCount); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.DB().QueryRowContext(context.Background(),
		`SELECT COUNT(*) FROM db_pool WHERE org_id = ?`, "org-shared-existing").Scan(&physicalPoolCount); err != nil {
		t.Fatal(err)
	}
	if existingTenantCount != 50 || physicalPoolCount != 1 {
		t.Fatalf("existing tenant_count=%d physical pools=%d, want reused existing capacity without a new pool",
			existingTenantCount, physicalPoolCount)
	}
}

func TestSharedTenantPoolDefensivelyReportsProvisioningWhenBatchMetadataComplete(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{
		provider:      tenant.ProviderTiDBCloudNative,
		cloudProvider: "aws",
		region:        "us-east-1",
		managedClusters: []tenant.CloudClusterInfo{{
			OrganizationID: "org-shared-defensive",
		}},
		sharedPoolResults: []*tenant.SharedDBPoolInfo{{
			ClusterID: "cluster-ready", Host: "127.0.0.1", Port: 4000,
			Username: "root", DBName: "tidbcloud_fs",
		}},
	}
	secret := make([]byte, 32)
	if _, err := rand.Read(secret); err != nil {
		t.Fatal(err)
	}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, TokenSecret: secret,
		Leader: newFollowerLeaderManager(t, metaStore)})
	defer srv.Close()
	now := time.Now().UTC()
	if err := metaStore.CreateTenantPool(context.Background(), &meta.TenantPool{
		PoolID: "pool-shared-defensive", OrganizationID: "org-shared-defensive", Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}

	results, err := srv.createFreePoolTenants(context.Background(), "pool-shared-defensive", 1,
		tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}, nil)
	if err != nil {
		t.Fatalf("createFreePoolTenants: %v", err)
	}
	if len(results) != 1 {
		t.Fatalf("results = %d, want 1", len(results))
	}
	if results[0].Status != meta.TenantProvisioning {
		t.Fatalf("result status = %q, want provisioning", results[0].Status)
	}
	got, err := metaStore.GetTenant(context.Background(), results[0].TenantID)
	if err != nil {
		t.Fatal(err)
	}
	if got.Status != meta.TenantProvisioning {
		t.Fatalf("persisted tenant status = %q, want provisioning", got.Status)
	}
}

func TestAdminTenantPoolCreateCleansPartialSharedMembersOnBatchFailure(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		managedClusters:   []tenant.CloudClusterInfo{{OrganizationID: "org-shared-cleanup"}},
		sharedPoolResults: []*tenant.SharedDBPoolInfo{{DBPoolID: 999999, ClusterID: "cluster-unknown", OrganizationID: "org-shared-cleanup"}}}
	secret := make([]byte, 32)
	_, _ = rand.Read(secret)
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, TokenSecret: secret})
	defer srv.Close()
	ts := httptest.NewServer(srv)
	defer ts.Close()

	resp := postJSON(t, ts.URL+"/v1/admin/tenant-pool", map[string]any{
		"public_key": "public", "private_key": "private", "pool_size": 1,
	}, "")
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusBadGateway {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("status = %d, want %d: %s", resp.StatusCode, http.StatusBadGateway, body)
	}
	var memberships, placements, tenantCount int
	if err := metaStore.DB().QueryRowContext(context.Background(), `SELECT COUNT(*) FROM tenant_pool_memberships
		WHERE tidbcloud_organization_id = ?`, "org-shared-cleanup").Scan(&memberships); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.DB().QueryRowContext(context.Background(), `SELECT COUNT(*) FROM tenant_placements p
		JOIN db_pool d ON d.db_id = p.db_id WHERE d.org_id = ?`, "org-shared-cleanup").Scan(&placements); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.DB().QueryRowContext(context.Background(), `SELECT COALESCE(SUM(tenant_count), 0)
		FROM db_pool WHERE org_id = ?`, "org-shared-cleanup").Scan(&tenantCount); err != nil {
		t.Fatal(err)
	}
	if memberships != 0 || placements != 0 || tenantCount != 0 {
		t.Fatalf("partial shared cleanup left memberships=%d placements=%d tenant_count=%d", memberships, placements, tenantCount)
	}
}

func TestDeleteFreeSharedPoolTenantFailsClosedWhenDBPoolIsUnready(t *testing.T) {
	rt := newQuotaRuntimeWithOptions(t, tenant.ProviderTiDBCloudNative, quotaRuntimeOptions{
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared,
	})
	ctx := context.Background()
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID: "pool-unready-delete", OrganizationID: "org-unready-delete", Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
	tenantID := "shared-unready-delete"
	if err := rt.server.insertPendingPoolTenant(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared, now); err != nil {
		t.Fatal(err)
	}
	if err := rt.server.materializeSharedTenantQuota(ctx, tenantID, provisionTenantOptions{}); err != nil {
		t.Fatal(err)
	}
	passwordCipher, err := rt.server.pool.Encrypt(ctx, []byte("root-pass"))
	if err != nil {
		t.Fatal(err)
	}
	spendingTarget := meta.MaxTiDBCloudSpendingLimit
	dbID, err := rt.meta.CreateManagedSharedDBPool(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-unready-delete", ProvisioningKey: make([]byte, 32),
		CloudProvider: "aws", Region: "us-east-1",
		MaxTenants: 100, SpendingLimit: &spendingTarget, PasswordCipher: passwordCipher, Name: "tidbcloud_fs",
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := rt.meta.UpdateManagedSharedDBPoolCloudResult(ctx, &meta.SharedDB{
		ID: dbID, TiDBCloudOrganizationID: "org-unready-delete", ClusterID: "cluster-unready-delete",
		PasswordCipher: passwordCipher, Name: "tidbcloud_fs", TLSMode: "true",
	}); err != nil {
		t.Fatal(err)
	}
	fsID, err := rt.meta.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if err := rt.meta.CompleteSharedTenantPoolMember(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared,
			SchemaShape: meta.SchemaShapeShared, Status: meta.SharedDBStatusActive},
		&meta.TenantPoolMembership{TenantID: tenantID, TiDBCloudOrganizationID: "org-unready-delete",
			PoolID: "pool-unready-delete", PoolStatus: meta.TenantPoolBindingFree, CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatal(err)
	}
	tenantRow, err := rt.meta.GetTenant(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	wantStatus := tenantRow.Status
	var notifyBefore int
	if err := rt.meta.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM tenant_notify_outbox WHERE tenant_id = ?`, tenantID).Scan(&notifyBefore); err != nil {
		t.Fatal(err)
	}
	owned, err := rt.server.deleteFreeSharedPoolTenant(ctx, tenantRow)
	if err == nil || owned {
		t.Fatalf("deleteFreeSharedPoolTenant: owned=%v err=%v, want false/error before durable mark", owned, err)
	}
	tenantRow, err = rt.meta.GetTenant(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if tenantRow.Status != wantStatus {
		t.Fatalf("tenant status after unready delete = %s, want %s", tenantRow.Status, wantStatus)
	}
	if _, err := rt.meta.GetTenantPlacement(ctx, fsID); err != nil {
		t.Fatalf("placement after unready delete = %v, want present", err)
	}
	if _, err := rt.meta.GetTenantPoolMembership(ctx, tenantID); err != nil {
		t.Fatalf("membership after unready delete = %v, want present", err)
	}
	dbPool, err := rt.meta.GetSharedDB(ctx, dbID)
	if err != nil {
		t.Fatal(err)
	}
	if dbPool.TenantCount != 1 {
		t.Fatalf("db pool tenant_count = %d, want 1", dbPool.TenantCount)
	}
	var notifyAfter int
	if err := rt.meta.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM tenant_notify_outbox WHERE tenant_id = ?`, tenantID).Scan(&notifyAfter); err != nil {
		t.Fatal(err)
	}
	if notifyAfter != notifyBefore {
		t.Fatalf("tenant notify rows = %d, want unchanged %d for an unready DB", notifyAfter, notifyBefore)
	}
}

func TestAdminTenantPoolMetadataResumeResultRank(t *testing.T) {
	ordered := []string{"ok", "canceled", "deadline_exceeded", "bad_conn", "error", "unknown"}
	for i := 1; i < len(ordered); i++ {
		prev := ordered[i-1]
		next := ordered[i]
		if adminTenantPoolMetadataResumeResultRank(next) < adminTenantPoolMetadataResumeResultRank(prev) {
			t.Fatalf("rank(%q) < rank(%q)", next, prev)
		}
	}
}

func TestAdminTenantPoolCreateUsesPrivateEndpointDBTLS(t *testing.T) {
	t.Setenv("DRIVE9_TIDBCLOUD_NATIVE_USE_PRIVATE_ENDPOINT", "1")
	rt, schemaInitRecorder := newAdminTenantPoolRuntime(t)
	rt.prov.listPages = []*tenant.ManagedClusterListResult{
		{},
		{Clusters: []tenant.CloudClusterInfo{{OrganizationID: "org-1"}}},
	}
	ts := httptest.NewServer(rt.server)
	t.Cleanup(ts.Close)

	resp := postJSON(t, ts.URL+"/v1/admin/tenant-pool", map[string]any{
		"public_key":  "public-1",
		"private_key": "private-1",
		"pool_size":   1,
	}, "")
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusAccepted {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("status = %d body=%s", resp.StatusCode, body)
	}

	deadline := time.Now().Add(5 * time.Second)
	var rows []meta.TenantWithTiDBCloudOrgBinding
	for {
		var err error
		rows, err = rt.meta.ListTenantPoolFreeSlotsForDelete(context.Background(), "org-1", false, 10)
		if err != nil {
			t.Fatalf("list free slots: %v", err)
		}
		if len(rows) == 1 && rows[0].Tenant.DBHost == "db.example.com" && schemaInitRecorder.schemaInitCalls.Load() >= 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("free slots = %d, schema init calls=%d", len(rows), schemaInitRecorder.schemaInitCalls.Load())
		}
		time.Sleep(10 * time.Millisecond)
	}
	if rows[0].Tenant.DBTLS {
		t.Fatalf("pool tenant DBTLS = true, want false for private endpoint")
	}
	assertTenantUsesPrivateEndpointTLS(t, rows[0].Tenant)
	assertSchemaInitUsesPrivateEndpointTLS(t, schemaInitRecorder.lastSchemaInitDSNSnapshot())
}

func TestTenantPoolMetadataResumeUsesPrivateEndpointDBTLS(t *testing.T) {
	t.Setenv("DRIVE9_TIDBCLOUD_NATIVE_USE_PRIVATE_ENDPOINT", "1")
	rt, schemaInitRecorder := newAdminTenantPoolRuntime(t)
	ctx := context.Background()
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID:         "pool-1",
		OrganizationID: "org-1",
		Size:           1,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	passCipher, err := rt.server.pool.Encrypt(ctx, []byte("pool-pass"))
	if err != nil {
		t.Fatalf("encrypt password: %v", err)
	}
	tenantID := "pool-private-resume-1"
	clusterID := "cluster-private-resume-1"
	if err := rt.meta.InsertTenant(ctx, &meta.Tenant{
		ID:               tenantID,
		Status:           meta.TenantPending,
		DBPasswordCipher: passCipher,
		DBName:           "tidbcloud_fs",
		DBTLS:            true,
		Provider:         tenant.ProviderTiDBCloudNative,
		ClusterID:        clusterID,
		SchemaVersion:    1,
		CreatedAt:        now,
		UpdatedAt:        now,
	}); err != nil {
		t.Fatalf("insert tenant: %v", err)
	}
	if err := rt.meta.UpsertTenantTiDBCloudOrgBinding(ctx, &meta.TenantTiDBCloudOrgBinding{
		TenantID:       tenantID,
		OrganizationID: "org-1",
		ClusterID:      clusterID,
		PoolID:         "pool-1",
		PoolStatus:     meta.TenantPoolBindingFree,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("upsert binding: %v", err)
	}

	rt.server.startPoolClustersMetadataResume(ctx, "pool-1", []*tenant.ClusterInfo{{
		TenantID:       tenantID,
		ClusterID:      clusterID,
		OrganizationID: "org-1",
		Password:       "pool-pass",
		DBName:         "tidbcloud_fs",
		Provider:       tenant.ProviderTiDBCloudNative,
	}}, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})

	deadline := time.Now().Add(5 * time.Second)
	var got *meta.Tenant
	for {
		got, err = rt.meta.GetTenant(ctx, tenantID)
		if err != nil {
			t.Fatalf("get tenant: %v", err)
		}
		if rt.prov.metadataBatchWaitCalls.Load() >= 1 && schemaInitRecorder.schemaInitCalls.Load() >= 1 && got.DBHost == "db.example.com" {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("tenant after resume = status %s host %q, metadata waits=%d, schema init calls=%d", got.Status, got.DBHost, rt.prov.metadataBatchWaitCalls.Load(), schemaInitRecorder.schemaInitCalls.Load())
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got.DBTLS {
		t.Fatalf("resumed pool tenant DBTLS = true, want false for private endpoint")
	}
	assertTenantUsesPrivateEndpointTLS(t, *got)
	assertSchemaInitUsesPrivateEndpointTLS(t, schemaInitRecorder.lastSchemaInitDSNSnapshot())
}

func TestTenantPoolMetadataResumePersistsAfterWaitDeadline(t *testing.T) {
	oldWaitTimeout := tenantPoolMetadataResumeWaitTimeout
	tenantPoolMetadataResumeWaitTimeout = 20 * time.Millisecond
	t.Cleanup(func() {
		tenantPoolMetadataResumeWaitTimeout = oldWaitTimeout
	})

	var waiter *deadlineMetadataResumeProvisioner
	rt, schemaInitRecorder := newAdminTenantPoolRuntimeWithProvisioner(t, func(recorder *adminTenantPoolSchemaInitRecorder) tenant.Provisioner {
		waiter = &deadlineMetadataResumeProvisioner{
			adminTenantPoolSchemaInitRecorder: recorder,
			waitStarted:                       make(chan struct{}),
		}
		return waiter
	})

	ctx := context.Background()
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID:         "pool-1",
		OrganizationID: "org-1",
		Size:           1,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	passCipher, err := rt.server.pool.Encrypt(ctx, []byte("pool-pass"))
	if err != nil {
		t.Fatalf("encrypt password: %v", err)
	}
	tenantID := "pool-deadline-resume-1"
	clusterID := "cluster-deadline-resume-1"
	if err := rt.meta.InsertTenant(ctx, &meta.Tenant{
		ID:               tenantID,
		Status:           meta.TenantPending,
		DBPasswordCipher: passCipher,
		DBName:           "tidbcloud_fs",
		DBTLS:            true,
		Provider:         tenant.ProviderTiDBCloudNative,
		ClusterID:        clusterID,
		SchemaVersion:    1,
		CreatedAt:        now,
		UpdatedAt:        now,
	}); err != nil {
		t.Fatalf("insert tenant: %v", err)
	}
	if err := rt.meta.UpsertTenantTiDBCloudOrgBinding(ctx, &meta.TenantTiDBCloudOrgBinding{
		TenantID:       tenantID,
		OrganizationID: "org-1",
		ClusterID:      clusterID,
		PoolID:         "pool-1",
		PoolStatus:     meta.TenantPoolBindingFree,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("upsert binding: %v", err)
	}

	rt.server.startPoolClustersMetadataResume(ctx, "pool-1", []*tenant.ClusterInfo{{
		TenantID:       tenantID,
		ClusterID:      clusterID,
		OrganizationID: "org-1",
		Password:       "pool-pass",
		DBName:         "tidbcloud_fs",
		Provider:       tenant.ProviderTiDBCloudNative,
	}}, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})

	select {
	case <-waiter.waitStarted:
	case <-time.After(5 * time.Second):
		t.Fatal("metadata resume did not start")
	}

	deadline := time.Now().Add(5 * time.Second)
	for {
		got, err := rt.meta.GetTenant(ctx, tenantID)
		if err != nil {
			t.Fatalf("get tenant: %v", err)
		}
		if rt.prov.metadataBatchWaitCalls.Load() >= 1 && schemaInitRecorder.schemaInitCalls.Load() >= 1 && got.DBHost == "db.example.com" {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("tenant after deadline resume = status %s host %q, metadata waits=%d, schema init calls=%d", got.Status, got.DBHost, rt.prov.metadataBatchWaitCalls.Load(), schemaInitRecorder.schemaInitCalls.Load())
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestTenantPoolMetadataResumePersistsReadyGroupBeforeSlowGroup(t *testing.T) {
	oldGroupSize := tenantPoolMetadataResumeGroupSize
	tenantPoolMetadataResumeGroupSize = 1
	t.Cleanup(func() {
		tenantPoolMetadataResumeGroupSize = oldGroupSize
	})

	var waiter *groupStreamingMetadataResumeProvisioner
	rt, _ := newAdminTenantPoolRuntimeWithProvisioner(t, func(recorder *adminTenantPoolSchemaInitRecorder) tenant.Provisioner {
		waiter = &groupStreamingMetadataResumeProvisioner{
			adminTenantPoolSchemaInitRecorder: recorder,
			slowTenantID:                      "pool-stream-resume-slow",
			slowStarted:                       make(chan struct{}),
			releaseSlow:                       make(chan struct{}),
		}
		return waiter
	})
	var releaseSlowOnce sync.Once
	t.Cleanup(func() {
		releaseSlowOnce.Do(func() { close(waiter.releaseSlow) })
	})
	ctx := context.Background()
	now := time.Now().UTC()
	if err := rt.meta.CreateTenantPool(ctx, &meta.TenantPool{
		PoolID:         "pool-1",
		OrganizationID: "org-1",
		Size:           2,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	makePending := func(tenantID, clusterID string) *tenant.ClusterInfo {
		t.Helper()
		passCipher, err := rt.server.pool.Encrypt(ctx, []byte("pool-pass"))
		if err != nil {
			t.Fatalf("encrypt password: %v", err)
		}
		if err := rt.meta.InsertTenant(ctx, &meta.Tenant{
			ID:               tenantID,
			Status:           meta.TenantPending,
			DBPasswordCipher: passCipher,
			DBName:           "tidbcloud_fs",
			DBTLS:            true,
			Provider:         tenant.ProviderTiDBCloudNative,
			ClusterID:        clusterID,
			SchemaVersion:    1,
			CreatedAt:        now,
			UpdatedAt:        now,
		}); err != nil {
			t.Fatalf("insert tenant %s: %v", tenantID, err)
		}
		if err := rt.meta.UpsertTenantTiDBCloudOrgBinding(ctx, &meta.TenantTiDBCloudOrgBinding{
			TenantID:       tenantID,
			OrganizationID: "org-1",
			ClusterID:      clusterID,
			PoolID:         "pool-1",
			PoolStatus:     meta.TenantPoolBindingFree,
			CreatedAt:      now,
			UpdatedAt:      now,
		}); err != nil {
			t.Fatalf("upsert binding %s: %v", tenantID, err)
		}
		return &tenant.ClusterInfo{
			TenantID:       tenantID,
			ClusterID:      clusterID,
			OrganizationID: "org-1",
			Password:       "pool-pass",
			DBName:         "tidbcloud_fs",
			Provider:       tenant.ProviderTiDBCloudNative,
		}
	}
	slow := makePending(waiter.slowTenantID, "cluster-stream-resume-slow")
	fast := makePending("pool-stream-resume-fast", "cluster-stream-resume-fast")

	rt.server.startPoolClustersMetadataResume(ctx, "pool-1", []*tenant.ClusterInfo{slow, fast}, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})

	select {
	case <-waiter.slowStarted:
	case <-time.After(5 * time.Second):
		t.Fatal("slow metadata resume group did not start")
	}

	deadline := time.Now().Add(5 * time.Second)
	for {
		got, err := rt.meta.GetTenant(ctx, fast.TenantID)
		if err != nil {
			t.Fatalf("get fast tenant: %v", err)
		}
		if got.DBHost == "db.example.com" {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("fast tenant was not persisted while slow group was blocked: status %s host %q", got.Status, got.DBHost)
		}
		time.Sleep(10 * time.Millisecond)
	}
	slowTenant, err := rt.meta.GetTenant(ctx, slow.TenantID)
	if err != nil {
		t.Fatalf("get slow tenant: %v", err)
	}
	if slowTenant.DBHost != "" {
		t.Fatalf("slow tenant host = %q before slow group release, want empty", slowTenant.DBHost)
	}

	releaseSlowOnce.Do(func() { close(waiter.releaseSlow) })
	deadline = time.Now().Add(5 * time.Second)
	for {
		got, err := rt.meta.GetTenant(ctx, slow.TenantID)
		if err != nil {
			t.Fatalf("get slow tenant: %v", err)
		}
		if got.DBHost == "db.example.com" {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("slow tenant was not persisted after release: status %s host %q", got.Status, got.DBHost)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestTenantPoolMetadataResumePersistContextPreservesServerCancellation(t *testing.T) {
	srv := NewWithConfig(Config{})
	ctx, cancel := srv.tenantPoolMetadataResumePersistContext(context.Background())
	defer cancel()

	srv.Close()

	select {
	case <-ctx.Done():
	case <-time.After(time.Second):
		t.Fatal("persist context was not canceled by server close")
	}
}

func TestAdminTenantPoolReplenishSkipsAtFreeWatermark(t *testing.T) {
	rt, _ := newAdminTenantPoolRuntime(t)
	ctx := context.Background()
	now := time.Now().UTC()
	pool := &meta.TenantPool{
		PoolID:         "pool-watermark-skip",
		OrganizationID: "org-1",
		Size:           10,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}
	if err := rt.meta.CreateTenantPool(ctx, pool); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	for i := 1; i <= 8; i++ {
		insertAdminPoolFreeTenant(t, rt, pool.PoolID, pool.OrganizationID, i)
	}
	lockHeld := make(chan struct{})
	releaseLock := make(chan struct{})
	lockDone := make(chan error, 1)
	go func() {
		lockDone <- rt.meta.WithTenantPoolLock(ctx, pool.PoolID, func(context.Context) error {
			close(lockHeld)
			<-releaseLock
			return nil
		})
	}()
	select {
	case <-lockHeld:
	case err := <-lockDone:
		t.Fatalf("hold tenant pool lock: %v", err)
	case <-time.After(time.Second):
		t.Fatal("timed out acquiring tenant pool lock for test")
	}
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseLock) }) }
	t.Cleanup(func() {
		release()
		<-lockDone
	})

	rt.server.replenishTenantPoolAsync(ctx, pool, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})
	workerDone := make(chan struct{})
	go func() {
		rt.server.forkWorkerWG.Wait()
		close(workerDone)
	}()
	select {
	case <-workerDone:
	case <-time.After(time.Second):
		release()
		<-workerDone
		t.Fatal("above-watermark replenish waited for the MySQL tenant pool lock")
	}
	release()

	if got := rt.prov.batchPoolCalls.Load(); got != 0 {
		t.Fatalf("batch pool calls = %d, want 0", got)
	}
	free, err := rt.meta.CountFreeTenantPoolBindings(ctx, pool.OrganizationID)
	if err != nil {
		t.Fatalf("count free: %v", err)
	}
	if free != 8 {
		t.Fatalf("free size = %d, want 8", free)
	}
}

func TestTenantPoolReplenishmentCoalescesConcurrentTriggersByPool(t *testing.T) {
	s := &Server{}
	now := time.Date(2026, 7, 25, 0, 0, 0, 0, time.UTC)

	if !s.beginTenantPoolReplenishmentAt("pool-1", now) {
		t.Fatal("first trigger was not accepted")
	}
	if s.beginTenantPoolReplenishmentAt("pool-1", now) {
		t.Fatal("concurrent trigger for the same pool was accepted")
	}
	if !s.beginTenantPoolReplenishmentAt("pool-2", now) {
		t.Fatal("trigger for a different pool was not accepted")
	}

	s.finishTenantPoolReplenishmentAt("pool-1", now)
	if s.beginTenantPoolReplenishmentAt("pool-1", now.Add(tenantPoolReplenishMinInterval-time.Nanosecond)) {
		t.Fatal("trigger was accepted during the per-pool minimum interval")
	}
	if !s.beginTenantPoolReplenishmentAt("pool-1", now.Add(tenantPoolReplenishMinInterval)) {
		t.Fatal("trigger was not accepted after the per-pool minimum interval")
	}
}

func TestTenantPoolReplenishmentRerunsCoalescedTriggerAfterWorkerFinishes(t *testing.T) {
	s := &Server{}
	now := time.Date(2026, 7, 25, 0, 0, 0, 0, time.UTC)

	first := s.requestTenantPoolReplenishmentAt("pool-1", now)
	if !first.start || first.scheduleAfter != 0 {
		t.Fatalf("first decision = %+v, want immediate start", first)
	}
	coalesced := s.requestTenantPoolReplenishmentAt("pool-1", now)
	if coalesced.start || coalesced.scheduleAfter != 0 {
		t.Fatalf("running decision = %+v, want coalesced trigger", coalesced)
	}
	finished := s.finishTenantPoolReplenishmentAt("pool-1", now)
	if finished.start || finished.scheduleAfter != tenantPoolReplenishMinInterval {
		t.Fatalf("finish decision = %+v, want one delayed rerun", finished)
	}
	cooldown := s.requestTenantPoolReplenishmentAt("pool-1", now.Add(tenantPoolReplenishMinInterval/2))
	if cooldown.start || cooldown.scheduleAfter != 0 {
		t.Fatalf("cooldown decision = %+v, want existing delayed rerun to absorb trigger", cooldown)
	}
	rerun := s.requestTenantPoolReplenishmentAt("pool-1", now.Add(tenantPoolReplenishMinInterval))
	if !rerun.start || rerun.scheduleAfter != 0 {
		t.Fatalf("rerun decision = %+v, want immediate start after cooldown", rerun)
	}
}

func TestTenantPoolReplenishmentSchedulesCoalescedRerunWhenWorkStartFails(t *testing.T) {
	s := &Server{}
	pool := &meta.TenantPool{PoolID: "pool-work-start-failed", OrganizationID: "org-1", Size: 1}
	timerStarts := 0
	workStarter := func(context.Context, func(context.Context)) bool {
		decision := s.requestTenantPoolReplenishmentAt(pool.PoolID, time.Now())
		if decision.start || decision.scheduleAfter != 0 {
			t.Fatalf("coalesced decision = %+v, want no immediate or delayed start", decision)
		}
		return false
	}
	timerStarter := func(context.Context, func(context.Context)) bool {
		timerStarts++
		return true
	}

	s.replenishTenantPoolAsyncWithStarters(context.Background(), pool, tenant.CredentialProvisionRequest{}, workStarter, timerStarter)

	if timerStarts != 1 {
		t.Fatalf("delayed rerun starts = %d, want 1 after work start failure", timerStarts)
	}
}

func TestTenantPoolReplenishmentRetriesSchedulingWhenTimerStartFails(t *testing.T) {
	s := &Server{}
	pool := &meta.TenantPool{PoolID: "pool-timer-start-failed", OrganizationID: "org-1", Size: 1}
	value, _ := s.tenantPoolReplenishJobs.LoadOrStore(pool.PoolID, &tenantPoolWorkGate{})
	gate := value.(*tenantPoolWorkGate)
	gate.mu.Lock()
	gate.rerun = true
	gate.nextAllowed = time.Now().Add(time.Minute)
	gate.mu.Unlock()
	timerStarts := 0
	timerStarter := func(context.Context, func(context.Context)) bool {
		timerStarts++
		return false
	}
	workStarter := func(context.Context, func(context.Context)) bool {
		t.Fatal("cooldown trigger unexpectedly started replenish work")
		return false
	}

	s.replenishTenantPoolAsyncWithStarters(context.Background(), pool, tenant.CredentialProvisionRequest{}, workStarter, timerStarter)
	s.replenishTenantPoolAsyncWithStarters(context.Background(), pool, tenant.CredentialProvisionRequest{}, workStarter, timerStarter)

	if timerStarts != 2 {
		t.Fatalf("timer start attempts = %d, want 2 after the first start failed", timerStarts)
	}
}

func TestLeaderTenantPoolReplenishmentTimerDoesNotConsumeWorkSlot(t *testing.T) {
	s := &Server{}
	pool := &meta.TenantPool{PoolID: "pool-leader-rerun", OrganizationID: "org-leader-rerun", Size: 1}
	value, _ := s.tenantPoolReplenishJobs.LoadOrStore(pool.PoolID, &tenantPoolWorkGate{})
	gate := value.(*tenantPoolWorkGate)
	gate.mu.Lock()
	gate.rerun = true
	gate.scheduled = true
	gate.nextAllowed = time.Now()
	gate.mu.Unlock()

	workStarted := make(chan struct{}, 1)
	workStarter := func(context.Context, func(context.Context)) bool {
		workStarted <- struct{}{}
		return true
	}
	timerStarted := make(chan struct{}, 1)
	timerStarter := func(ctx context.Context, fn func(context.Context)) bool {
		timerStarted <- struct{}{}
		go fn(ctx)
		return true
	}

	s.scheduleTenantPoolReplenishment(context.Background(), pool, tenant.CredentialProvisionRequest{}, 0, workStarter, timerStarter)
	select {
	case <-timerStarted:
	case <-time.After(time.Second):
		t.Fatal("leader rerun timer did not start")
	}
	select {
	case <-workStarted:
	case <-time.After(time.Second):
		t.Fatal("leader rerun did not reacquire a refill work slot")
	}
}

func TestTenantPoolPendingResumeScanCooldownAfterEmptyResult(t *testing.T) {
	s := &Server{}
	now := time.Date(2026, 7, 25, 0, 0, 0, 0, time.UTC)

	if !s.beginTenantPoolResumeScanAt("pool-1", now) {
		t.Fatal("first pending resume scan was not accepted")
	}
	if s.beginTenantPoolResumeScanAt("pool-1", now) {
		t.Fatal("concurrent pending resume scan was accepted")
	}
	s.finishTenantPoolResumeScanAt("pool-1", now, true)
	if s.beginTenantPoolResumeScanAt("pool-1", now.Add(tenantPoolPendingResumeEmptyInterval-time.Nanosecond)) {
		t.Fatal("empty pending resume scan was retried during cooldown")
	}
	if !s.beginTenantPoolResumeScanAt("pool-1", now.Add(tenantPoolPendingResumeEmptyInterval)) {
		t.Fatal("pending resume scan was not accepted after cooldown")
	}
}

func TestTenantPoolPendingResumeRequestsRerunBeforeListing(t *testing.T) {
	s := &Server{}
	job := &tenantPoolResumeJob{}
	s.tenantPoolResumeJobs.Store("pool-1", job)

	s.resumePendingTenantPoolAsync(context.Background(), &meta.TenantPool{
		PoolID:         "pool-1",
		OrganizationID: "org-1",
		Size:           10,
	}, tenant.CredentialProvisionRequest{})

	if !job.rerun.Load() {
		t.Fatal("active metadata resume job did not receive a rerun request")
	}
}

func TestAdminTenantPoolReplenishBatchesBelowFreeWatermark(t *testing.T) {
	rt, _ := newAdminTenantPoolRuntime(t)
	ctx := context.Background()
	now := time.Now().UTC()
	pool := &meta.TenantPool{
		PoolID:         "pool-watermark-refill",
		OrganizationID: "org-1",
		Size:           10,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}
	if err := rt.meta.CreateTenantPool(ctx, pool); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	for i := 1; i <= 7; i++ {
		insertAdminPoolFreeTenant(t, rt, pool.PoolID, pool.OrganizationID, i)
	}

	rt.server.replenishTenantPoolAsync(ctx, pool, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})
	rt.server.forkWorkerWG.Wait()

	if got := rt.prov.batchPoolCalls.Load(); got != 1 {
		t.Fatalf("batch pool calls = %d, want 1", got)
	}
	deadline := time.Now().Add(5 * time.Second)
	for {
		free, err := rt.meta.CountFreeTenantPoolBindings(ctx, pool.OrganizationID)
		if err != nil {
			t.Fatalf("count free: %v", err)
		}
		if free == 10 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("free size = %d, want 10 after refill", free)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestAdminTenantPoolReplenishSubmitsLargeRefillAsSingleWave(t *testing.T) {
	oldRetryWindow := schemaInitRetryWindow
	schemaInitRetryWindow = 100 * time.Millisecond
	t.Cleanup(func() { schemaInitRetryWindow = oldRetryWindow })
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 100,
		TokenSecret: make([]byte, 32), Leader: newFollowerLeaderManager(t, metaStore)})
	defer srv.Close()
	// Make the complete three-pool refill wave visible as one Cloud batch.
	srv.managedSharedDBCloudBatchSize = 250
	ctx := context.Background()
	now := time.Now().UTC()
	pool := &meta.TenantPool{
		PoolID:         "pool-large-refill",
		OrganizationID: "org-1",
		Size:           250,
		Status:         meta.TenantPoolActive,
		CreatedAt:      now,
		UpdatedAt:      now,
	}
	if err := metaStore.CreateTenantPool(ctx, pool); err != nil {
		t.Fatalf("create pool: %v", err)
	}

	lockHeld := make(chan struct{})
	releaseLock := make(chan struct{})
	lockDone := make(chan error, 1)
	go func() {
		lockDone <- metaStore.WithTenantPoolLock(ctx, pool.PoolID, func(context.Context) error {
			close(lockHeld)
			<-releaseLock
			return nil
		})
	}()
	select {
	case <-lockHeld:
	case err := <-lockDone:
		t.Fatalf("hold tenant pool lock: %v", err)
	case <-time.After(time.Second):
		t.Fatal("timed out acquiring tenant pool lock for test")
	}
	t.Cleanup(func() {
		close(releaseLock)
		if err := <-lockDone; err != nil {
			t.Errorf("release tenant pool lock: %v", err)
		}
	})

	srv.replenishTenantPoolAsync(ctx, pool, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})
	deadline := time.Now().Add(10 * time.Second)
	for {
		if got := prov.sharedPoolBatchCalls.Load(); got >= 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("shared cloud batch calls = %d, want one complete refill wave", prov.sharedPoolBatchCalls.Load())
		}
		time.Sleep(10 * time.Millisecond)
	}
	deadline = time.Now().Add(10 * time.Second)
	for {
		free, err := metaStore.CountTenantPoolFreeSlots(ctx, pool.OrganizationID)
		if err != nil {
			t.Fatalf("count free: %v", err)
		}
		if free == pool.Size {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("free size = %d, want %d after refill", free, pool.Size)
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 1 {
		t.Fatalf("shared cloud batch calls = %d, want 1", got)
	}
	if got := prov.sharedPoolBatchMembers.Load(); got != 3 {
		t.Fatalf("shared cloud batch members = %d, want 3 physical pools", got)
	}
}

func TestManagedSharedDBCloudBatchRequestsAreGloballyBounded(t *testing.T) {
	const wantMaxConcurrentBatches = 5
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(poolManager.Close)
	poolManager.SetMetaStore(metaStore)
	batchStarted := make(chan struct{}, 7)
	batchRelease := make(chan struct{}, 7)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1",
		sharedPoolBatchStarted: batchStarted, sharedPoolBatchRelease: batchRelease}
	srv := &Server{meta: metaStore, pool: poolManager, provisioner: prov, sharedDBMaxTenants: 1,
		managedSharedDBCloudBatchSize: 1}
	ctx := context.Background()
	cred := tenant.CredentialProvisionRequest{PublicKey: "public", PrivateKey: "private"}
	dbIDs := make([]int64, 0, 7)
	for range 7 {
		row, err := srv.createManagedSharedDBPlan(ctx, "org-cloud-batch-bound", sharedDBProvisioningKey(cred))
		if err != nil {
			t.Fatal(err)
		}
		dbIDs = append(dbIDs, row.ID)
	}
	done := make(chan error, 1)
	go func() {
		_, createErr := srv.provisionManagedSharedDBPoolsBatchWithCredentials(ctx, dbIDs, cred)
		done <- createErr
	}()
	for i := 0; i < wantMaxConcurrentBatches; i++ {
		select {
		case <-batchStarted:
		case <-time.After(2 * time.Second):
			t.Fatalf("Cloud batch %d did not start", i+1)
		}
	}
	select {
	case <-batchStarted:
		t.Fatalf("more than %d Cloud batch requests started concurrently", wantMaxConcurrentBatches)
	case <-time.After(150 * time.Millisecond):
	}
	batchRelease <- struct{}{}
	select {
	case <-batchStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("next Cloud batch did not start after a request slot was released")
	}
	close(batchRelease)
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("provision managed shared DB batches: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Cloud batch provisioning did not finish")
	}
}

func TestAdminTenantPoolReplenishContinuesPastWatermarkToFillSlots(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 200,
		TokenSecret: make([]byte, 32), Leader: newFollowerLeaderManager(t, metaStore)})
	defer srv.Close()
	ctx := context.Background()
	now := time.Now().UTC()
	pool := &meta.TenantPool{PoolID: "pool-watermark-shared", OrganizationID: "org-watermark-shared",
		Size: 250, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(ctx, pool); err != nil {
		t.Fatalf("create pool: %v", err)
	}
	passwordCipher, err := poolManager.Encrypt(ctx, []byte("shared-pass"))
	if err != nil {
		t.Fatalf("encrypt shared password: %v", err)
	}
	if _, err := metaStore.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: pool.OrganizationID,
		Host:                    "shared.example.com",
		Port:                    4000,
		User:                    "root",
		PasswordCipher:          passwordCipher,
		Name:                    "tidbcloud_fs",
		MaxTenants:              200,
	}); err != nil {
		t.Fatalf("register shared db: %v", err)
	}
	if _, err := srv.createFreeSharedPoolTenants(ctx, pool.PoolID, 100,
		tenant.CredentialProvisionRequest{PublicKey: "public-1", PrivateKey: "private-1"}, nil); err != nil {
		t.Fatalf("seed shared pool: %v", err)
	}

	srv.replenishTenantPoolAsync(ctx, pool, tenant.CredentialProvisionRequest{
		PublicKey:  "public-1",
		PrivateKey: "private-1",
	})
	deadline := time.Now().Add(10 * time.Second)
	for {
		free, err := metaStore.CountTenantPoolFreeSlots(ctx, pool.OrganizationID)
		if err != nil {
			t.Fatalf("count free slots: %v", err)
		}
		if free == pool.Size {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("free slots = %d, want %d after refill", free, pool.Size)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestSharedTenantPoolLeaderReconcilerRefillsZeroInventory(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 100,
		TokenSecret: make([]byte, 32)})
	defer srv.Close()
	ctx := context.Background()
	now := time.Now().UTC()
	pool := &meta.TenantPool{PoolID: "pool-leader-reconcile", OrganizationID: "org-leader-reconcile",
		Size: 1, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(ctx, pool); err != nil {
		t.Fatalf("create pool: %v", err)
	}

	deadline := time.Now().Add(10 * time.Second)
	for {
		slots, countErr := metaStore.CountTenantPoolFreeSlots(ctx, pool.OrganizationID)
		if countErr != nil {
			t.Fatalf("count free slots: %v", countErr)
		}
		if slots == pool.Size {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("free slots = %d, want %d after leader reconcile", slots, pool.Size)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestSharedTenantPoolRefillDoesNotReplaceRecoverablePlannedCapacity(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	t.Cleanup(func() { poolManager.Close() })
	poolManager.SetMetaStore(metaStore)
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, cloudProvider: "aws", region: "us-east-1"}
	srv := NewWithConfig(Config{Meta: metaStore, Pool: poolManager, Provisioner: prov,
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared, SharedDBMaxTenants: 1,
		TokenSecret: make([]byte, 32),
		Leader:      newFollowerLeaderManager(t, metaStore)})
	t.Cleanup(func() { srv.Close() })
	ctx := context.Background()
	now := time.Now().UTC()
	logicalPool := &meta.TenantPool{PoolID: "pool-expired-plan", OrganizationID: "org-expired-plan",
		Size: 1, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	spendingLimit := int64(meta.MaxTiDBCloudSpendingLimit)
	dbID, err := metaStore.CreateManagedSharedDBPool(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: logicalPool.OrganizationID, ProvisioningKey: bytes.Repeat([]byte{1}, 32),
		CloudProvider: "aws", Region: "us-east-1", MaxTenants: 1, SpendingLimit: &spendingLimit,
	})
	if err != nil {
		t.Fatalf("CreateManagedSharedDBPool: %v", err)
	}
	tenantID := "tenant-expired-plan"
	if err := srv.insertPendingPoolTenant(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared, now); err != nil {
		t.Fatalf("insertPendingPoolTenant: %v", err)
	}
	fsID, err := metaStore.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatalf("EnsureFsID: %v", err)
	}
	if err := metaStore.CompleteSharedTenantPoolMember(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared, SchemaShape: meta.SchemaShapeShared},
		&meta.TenantPoolMembership{TenantID: tenantID, TiDBCloudOrganizationID: logicalPool.OrganizationID,
			PoolID: logicalPool.PoolID, PoolStatus: meta.TenantPoolBindingFree}); err != nil {
		t.Fatalf("CompleteSharedTenantPoolMember: %v", err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET updated_at = ? WHERE db_id = ?`, now.Add(-2*time.Minute), dbID); err != nil {
		t.Fatalf("age planned pool: %v", err)
	}

	srv.replenishTenantPool(ctx, logicalPool, tenant.CredentialProvisionRequest{})
	var physicalPools int
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalPools); err != nil {
		t.Fatalf("count pending physical pools: %v", err)
	}
	if physicalPools != 1 {
		t.Fatalf("pending recoverable physical pools = %d, want no replacement", physicalPools)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET status = ?, updated_at = ? WHERE db_id = ?`,
		meta.SharedDBStatusActive, now.Add(-20*time.Minute), dbID); err != nil {
		t.Fatalf("activate and age physical pool: %v", err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE tenants SET status = ? WHERE id = ?`, meta.TenantProvisioning, tenantID); err != nil {
		t.Fatalf("leave tenant activation pending: %v", err)
	}
	srv.replenishTenantPool(ctx, logicalPool, tenant.CredentialProvisionRequest{})
	if err := metaStore.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM db_pool WHERE org_id = ?`, logicalPool.OrganizationID).Scan(&physicalPools); err != nil {
		t.Fatalf("count active physical pools: %v", err)
	}
	if physicalPools != 1 {
		t.Fatalf("active pool with provisioning tenant produced %d physical pools, want no replacement", physicalPools)
	}
}

func TestLeaderTenantPoolReplenishmentDoesNotStartOnFollower(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	mgr := leader.NewManager(metaStore.DB())
	srv := &Server{
		leader: mgr, tenantPoolReconcileQueue: make(chan tenantPoolReconcileJob),
		leaderWorkersStarted: true, leaderWorkerCtx: context.Background(),
	}
	if srv.enqueueTenantPoolLeaderReplenishment(context.Background(), &meta.TenantPool{
		PoolID: "pool-follower", OrganizationID: "org-follower", Size: 1, Status: meta.TenantPoolActive,
	}, tenant.CredentialProvisionRequest{}) {
		t.Fatal("follower unexpectedly queued leader refill work")
	}
	if _, ok := srv.tenantPoolReplenishJobs.Load("pool-follower"); ok {
		t.Fatal("leader refill unexpectedly initialized the request-side per-pool gate")
	}
}

func TestTenantPoolLeaderReconcileWorkersQueueAndRestBetweenTasks(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rest := 40 * time.Millisecond
	srv := &Server{tenantPoolReconcileQueue: make(chan tenantPoolReconcileJob), tenantPoolReconcileWorkerRest: rest}
	enqueue := func(fn func(context.Context)) bool {
		select {
		case srv.tenantPoolReconcileQueue <- tenantPoolReconcileJob{run: fn}:
			return true
		case <-ctx.Done():
			return false
		}
	}
	workerDone := make(chan struct{})
	go func() {
		defer close(workerDone)
		srv.runTenantPoolReconcileWorker(ctx)
	}()
	firstStarted := make(chan struct{}, 1)
	release := make(chan struct{})
	if !enqueue(func(context.Context) {
		firstStarted <- struct{}{}
		<-release
	}) {
		t.Fatal("first leader reconcile task was not queued")
	}
	<-firstStarted
	secondStarted := make(chan struct{}, 1)
	secondQueued := make(chan bool, 1)
	go func() {
		secondQueued <- enqueue(func(context.Context) {
			secondStarted <- struct{}{}
		})
	}()
	select {
	case <-secondQueued:
		t.Fatal("second task left the queue while the only worker was occupied")
	case <-time.After(20 * time.Millisecond):
	}
	releasedAt := time.Now()
	close(release)
	select {
	case <-secondStarted:
		if elapsed := time.Since(releasedAt); elapsed < rest {
			t.Fatalf("queued task started after %s, want worker rest of at least %s", elapsed, rest)
		}
	case <-time.After(time.Second):
		t.Fatal("queued task did not start after the first worker rested")
	}
	if queued := <-secondQueued; !queued {
		t.Fatal("second leader reconcile task was not queued")
	}
	cancel()
	select {
	case <-workerDone:
	case <-time.After(time.Second):
		t.Fatal("tenant-pool reconcile worker did not stop")
	}
}

func TestTenantPoolLeaderReconcileCoalescesSamePoolIntoOneRerun(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	srv := &Server{
		tenantPoolReconcileQueue:      make(chan tenantPoolReconcileJob),
		tenantPoolReconcileWorkerRest: 20 * time.Millisecond,
	}
	workerDone := make(chan struct{})
	go func() {
		defer close(workerDone)
		srv.runTenantPoolReconcileWorker(ctx)
	}()

	started := make(chan int, 2)
	releaseFirst := make(chan struct{})
	var attempts atomic.Int32
	run := func(context.Context) {
		attempt := int(attempts.Add(1))
		started <- attempt
		if attempt == 1 {
			<-releaseFirst
		}
	}
	if !srv.enqueueTenantPoolLeaderReconcile(ctx, "pool-repeated-scan", run) {
		t.Fatal("initial leader reconcile task was not queued")
	}
	select {
	case attempt := <-started:
		if attempt != 1 {
			t.Fatalf("first attempt = %d, want 1", attempt)
		}
	case <-time.After(time.Second):
		t.Fatal("initial leader reconcile task did not start")
	}

	for i := 0; i < 3; i++ {
		if !srv.enqueueTenantPoolLeaderReconcile(ctx, "pool-repeated-scan", run) {
			t.Fatalf("coalesced trigger %d was rejected", i+1)
		}
	}
	select {
	case attempt := <-started:
		t.Fatalf("attempt %d started while the first attempt was still running", attempt)
	case <-time.After(30 * time.Millisecond):
	}
	close(releaseFirst)
	select {
	case attempt := <-started:
		if attempt != 2 {
			t.Fatalf("rerun attempt = %d, want 2", attempt)
		}
	case <-time.After(time.Second):
		t.Fatal("coalesced rerun did not start")
	}
	select {
	case attempt := <-started:
		t.Fatalf("unexpected extra coalesced attempt %d", attempt)
	case <-time.After(50 * time.Millisecond):
	}
	if got := attempts.Load(); got != 2 {
		t.Fatalf("leader reconcile attempts = %d, want 2", got)
	}
	if _, ok := srv.tenantPoolReplenishJobs.Load("pool-repeated-scan"); ok {
		t.Fatal("leader refill unexpectedly used the request-side per-pool gate")
	}

	cancel()
	select {
	case <-workerDone:
	case <-time.After(time.Second):
		t.Fatal("tenant-pool reconcile worker did not stop")
	}
}

func TestTenantPoolLeaderReconcileDoesNotRerunForDuplicateQueuedScan(t *testing.T) {
	state := &tenantPoolLeaderReconcileState{}
	if !state.request() {
		t.Fatal("initial scan was not accepted")
	}
	if state.request() {
		t.Fatal("duplicate queued scan was accepted as a second job")
	}
	state.start()
	if state.next() {
		t.Fatal("duplicate scan received before work started requested an unnecessary rerun")
	}
}

func TestManagedSharedDBStuckReconcilerFailsPoolAndPlacedTenant(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	ctx := context.Background()
	now := time.Now().UTC()
	logicalPool := &meta.TenantPool{PoolID: "logical-stuck", OrganizationID: "org-stuck", Size: 1,
		Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(ctx, logicalPool); err != nil {
		t.Fatalf("CreateTenantPool: %v", err)
	}
	spendingLimit := meta.MaxTiDBCloudSpendingLimit
	dbID, err := metaStore.CreateManagedSharedDBPool(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-stuck", ProvisioningKey: bytes.Repeat([]byte{1}, 32),
		CloudProvider: "aws", Region: "us-east-1", MaxTenants: 100, SpendingLimit: &spendingLimit,
	})
	if err != nil {
		t.Fatalf("CreateManagedSharedDBPool: %v", err)
	}
	tenantID := "tenant-stuck-reconcile"
	if err := metaStore.InsertTenant(ctx, &meta.Tenant{ID: tenantID, Status: meta.TenantPending,
		Provider: tenant.ProviderTiDBCloudNative, DBPasswordCipher: []byte{}, SchemaVersion: 1,
		CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatalf("InsertTenant: %v", err)
	}
	fsID, err := metaStore.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatalf("EnsureFsID: %v", err)
	}
	if err := metaStore.CompleteSharedTenantPoolMember(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared, SchemaShape: meta.SchemaShapeShared},
		&meta.TenantPoolMembership{TenantID: tenantID, TiDBCloudOrganizationID: "org-stuck",
			PoolID: logicalPool.PoolID, PoolStatus: meta.TenantPoolBindingFree}); err != nil {
		t.Fatalf("CompleteSharedTenantPoolMember: %v", err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET updated_at = ? WHERE db_id = ?`, now.Add(-20*time.Minute), dbID); err != nil {
		t.Fatalf("age db pool: %v", err)
	}

	srv := &Server{meta: metaStore, managedSharedDBStuckTimeout: 15 * time.Minute}
	srv.reconcileStuckManagedSharedDBPoolsWithCtx(ctx)
	pool, err := metaStore.GetSharedDB(ctx, dbID)
	if err != nil || pool.Status != meta.SharedDBStatusFailed {
		t.Fatalf("pool after stuck reconcile = %+v, %v", pool, err)
	}
	gotTenant, err := metaStore.GetTenant(ctx, tenantID)
	if err != nil || gotTenant.Status != meta.TenantFailed {
		t.Fatalf("tenant after stuck reconcile = %+v, %v", gotTenant, err)
	}
}

func TestManagedSharedDBActiveTenantReconcilerFinalizesStrandedTenant(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	ctx := context.Background()
	now := time.Now().UTC()
	spendingLimit := meta.MaxTiDBCloudSpendingLimit
	dbID, err := metaStore.CreateManagedSharedDBPool(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-active-recovery", ProvisioningKey: bytes.Repeat([]byte{1}, 32),
		CloudProvider: "aws", Region: "us-east-1", MaxTenants: 100, SpendingLimit: &spendingLimit,
	})
	if err != nil {
		t.Fatal(err)
	}
	tenantID := "tenant-active-recovery"
	if err := metaStore.InsertTenant(ctx, &meta.Tenant{ID: tenantID, Status: meta.TenantPending,
		Provider: tenant.ProviderTiDBCloudNativeShared, DBPasswordCipher: []byte{}, SchemaVersion: 1,
		CreatedAt: now, UpdatedAt: now}); err != nil {
		t.Fatal(err)
	}
	fsID, err := metaStore.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if err := metaStore.CompleteSharedTenantPoolMember(ctx, tenantID, tenant.ProviderTiDBCloudNativeShared,
		&meta.TenantPlacement{FsID: fsID, DbID: dbID, Placement: meta.PlacementShared, SchemaShape: meta.SchemaShapeShared},
		&meta.TenantPoolMembership{TenantID: tenantID, TiDBCloudOrganizationID: "org-active-recovery",
			PoolID: "logical-active-recovery", PoolStatus: meta.TenantPoolBindingFree}); err != nil {
		t.Fatal(err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET status = ? WHERE db_id = ?`, meta.SharedDBStatusActive, dbID); err != nil {
		t.Fatal(err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE tenants SET status = ? WHERE id = ?`, meta.TenantProvisioning, tenantID); err != nil {
		t.Fatal(err)
	}

	srv := &Server{meta: metaStore}
	srv.resumeActiveManagedSharedDBTenantsWithCtx(ctx)
	got, err := metaStore.GetTenant(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if got.Status != meta.TenantActive {
		t.Fatalf("stranded tenant status = %q, want active", got.Status)
	}
}

func TestManagedSharedDBStuckReconcilerSkipsPoolWithActiveRecoveryClaim(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	ctx := context.Background()
	spendingLimit := meta.MaxTiDBCloudSpendingLimit
	dbID, err := metaStore.CreateManagedSharedDBPool(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-stuck-active", ProvisioningKey: bytes.Repeat([]byte{8}, 32),
		CloudProvider: "aws", Region: "us-east-1", MaxTenants: 100, SpendingLimit: &spendingLimit,
	})
	if err != nil {
		t.Fatalf("CreateManagedSharedDBPool: %v", err)
	}
	if _, err := metaStore.DB().ExecContext(ctx, `UPDATE db_pool SET updated_at = ? WHERE db_id = ?`, time.Now().UTC().Add(-20*time.Minute), dbID); err != nil {
		t.Fatalf("age db pool: %v", err)
	}

	srv := &Server{meta: metaStore, managedSharedDBStuckTimeout: 15 * time.Minute}
	holderStarted := make(chan struct{})
	releaseHolder := make(chan struct{})
	holderDone := make(chan error, 1)
	go func() {
		holderDone <- srv.withSharedDBPoolWorkLock(ctx, dbID, func(context.Context) error {
			close(holderStarted)
			<-releaseHolder
			return nil
		})
	}()
	<-holderStarted

	srv.reconcileStuckManagedSharedDBPoolsWithCtx(ctx)
	row, err := metaStore.GetSharedDB(ctx, dbID)
	if err != nil {
		t.Fatal(err)
	}
	if row.Status != meta.SharedDBStatusPending {
		t.Fatalf("busy recovery pool status = %q, want pending", row.Status)
	}

	close(releaseHolder)
	if err := <-holderDone; err != nil {
		t.Fatalf("recovery claim holder: %v", err)
	}
	srv.reconcileStuckManagedSharedDBPoolsWithCtx(ctx)
	row, err = metaStore.GetSharedDB(ctx, dbID)
	if err != nil {
		t.Fatal(err)
	}
	if row.Status != meta.SharedDBStatusFailed {
		t.Fatalf("idle stuck pool status = %q, want failed", row.Status)
	}
}

func TestManagedSharedDBRefillCapsOneWaveAtFiftyPhysicalPools(t *testing.T) {
	if got := managedSharedDBRefillTenantCount(10_000, 100, 50); got != 5_000 {
		t.Fatalf("refill tenant count = %d, want capacity of 50 physical pools", got)
	}
	if got := managedSharedDBRefillTenantCount(1_600, 100, 50); got != 1_600 {
		t.Fatalf("refill tenant count = %d, want all 1600 tenants across 16 physical pools", got)
	}
	if got := managedSharedDBRefillTenantCount(100, 1, 50); got != 50 {
		t.Fatalf("refill tenant count with one tenant per DB = %d, want 50", got)
	}
	if got := managedSharedDBRefillTenantCount(10_000, 100, 12); got != 1_200 {
		t.Fatalf("refill tenant count with configured 12-DB limit = %d, want 1200", got)
	}
}

func TestManagedSharedDBReplenishSubmitsWholePhysicalPoolWaveConcurrently(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = metaStore.Close() }()
	testtidb.ResetMetaDB(t, metaStore.DB())
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	poolManager := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost"}, enc)
	defer poolManager.Close()
	poolManager.SetMetaStore(metaStore)
	batchStarted := make(chan struct{}, 1)
	batchRelease := make(chan struct{})
	prov := &fakeProvisioner{
		provider:               tenant.ProviderTiDBCloudNative,
		cloudProvider:          "aws",
		region:                 "us-east-1",
		sharedPoolBatchStarted: batchStarted,
		sharedPoolBatchRelease: batchRelease,
	}
	workerCtx, cancel := context.WithCancel(context.Background())
	srv := &Server{
		meta: metaStore, pool: poolManager, provisioner: prov,
		defaultTenantProvider: tenant.ProviderTiDBCloudNativeShared,
		sharedDBMaxTenants:    1, managedSharedDBCloudBatchSize: 50,
		tenantPoolRefillFreeRatio: DefaultTenantPoolRefillFreeRatio,
		forkWorkerCtx:             workerCtx, forkWorkerCancel: cancel,
	}
	defer srv.Close()
	now := time.Now().UTC()
	logicalPool := &meta.TenantPool{PoolID: "pool-whole-wave", OrganizationID: "org-whole-wave",
		Size: 50, Status: meta.TenantPoolActive, CreatedAt: now, UpdatedAt: now}
	if err := metaStore.CreateTenantPool(context.Background(), logicalPool); err != nil {
		t.Fatal(err)
	}

	workerDone := make(chan struct{})
	workStarter := func(ctx context.Context, fn func(context.Context)) bool {
		go func() {
			defer close(workerDone)
			fn(ctx)
		}()
		return true
	}
	var unexpectedTimer atomic.Bool
	timerStarter := func(context.Context, func(context.Context)) bool {
		unexpectedTimer.Store(true)
		return false
	}
	var releaseOnce sync.Once
	releaseBatches := func() { releaseOnce.Do(func() { close(batchRelease) }) }
	defer func() {
		releaseBatches()
		select {
		case <-workerDone:
		case <-time.After(10 * time.Second):
			t.Error("replenish worker did not stop during test cleanup")
		}
	}()

	srv.replenishTenantPoolAsyncWithStarters(context.Background(), logicalPool,
		tenant.CredentialProvisionRequest{}, workStarter, timerStarter)
	select {
	case <-batchStarted:
	case <-time.After(3 * time.Second):
		t.Fatal("Cloud batch did not start")
	}
	releaseBatches()
	select {
	case <-workerDone:
	case <-time.After(10 * time.Second):
		t.Fatal("whole-wave replenish did not finish")
	}
	if got := prov.sharedPoolBatchCalls.Load(); got != 1 {
		t.Fatalf("Cloud batch calls = %d, want 1", got)
	}
	if got := prov.sharedPoolBatchMembers.Load(); got != 50 {
		t.Fatalf("Cloud batch members = %d, want 50", got)
	}
	if unexpectedTimer.Load() {
		t.Fatal("unexpected replenish rerun timer")
	}
}

func TestTenantPoolEffectiveRefillRatioRejectsNaN(t *testing.T) {
	s := &Server{tenantPoolRefillFreeRatio: math.NaN()}
	if got := s.effectiveTenantPoolRefillFreeRatio(); got != DefaultTenantPoolRefillFreeRatio {
		t.Fatalf("effective refill ratio = %f, want %f", got, DefaultTenantPoolRefillFreeRatio)
	}
}

type adminTenantPoolSchemaInitRecorder struct {
	*quotaTestProvisioner

	schemaInitCalls   atomic.Int32
	mu                sync.Mutex
	lastSchemaInitDSN string
}

func newAdminTenantPoolRuntime(t *testing.T) (*quotaRuntime, *adminTenantPoolSchemaInitRecorder) {
	return newAdminTenantPoolRuntimeWithProvisioner(t, nil)
}

func newAdminTenantPoolRuntimeWithProvisioner(t *testing.T, wrap func(*adminTenantPoolSchemaInitRecorder) tenant.Provisioner) (*quotaRuntime, *adminTenantPoolSchemaInitRecorder) {
	t.Helper()
	var recorder *adminTenantPoolSchemaInitRecorder
	rt := newQuotaRuntimeWithOptions(t, tenant.ProviderTiDBCloudNative, quotaRuntimeOptions{provisioner: func(prov *quotaTestProvisioner) tenant.Provisioner {
		recorder = &adminTenantPoolSchemaInitRecorder{quotaTestProvisioner: prov}
		if wrap != nil {
			return wrap(recorder)
		}
		return recorder
	}})
	return rt, recorder
}

func insertAdminPoolFreeTenant(t *testing.T, rt *quotaRuntime, poolID, organizationID string, index int) string {
	t.Helper()
	ctx := context.Background()
	now := time.Now().UTC().Add(time.Duration(index) * time.Second)
	tenantID := fmt.Sprintf("%s-free-%d", poolID, index)
	clusterID := fmt.Sprintf("%s-cluster-%d", poolID, index)
	passCipher, err := rt.server.pool.Encrypt(ctx, []byte("pool-pass"))
	if err != nil {
		t.Fatalf("encrypt password: %v", err)
	}
	if err := rt.meta.InsertTenant(ctx, &meta.Tenant{
		ID:               tenantID,
		Status:           meta.TenantActive,
		DBHost:           "db.example.com",
		DBPort:           4000,
		DBUser:           "u.root",
		DBPasswordCipher: passCipher,
		DBName:           "tidbcloud_fs",
		DBTLS:            true,
		Provider:         tenant.ProviderTiDBCloudNative,
		ClusterID:        clusterID,
		SchemaVersion:    1,
		CreatedAt:        now,
		UpdatedAt:        now,
	}); err != nil {
		t.Fatalf("insert tenant %s: %v", tenantID, err)
	}
	if err := rt.meta.UpsertTenantTiDBCloudOrgBinding(ctx, &meta.TenantTiDBCloudOrgBinding{
		TenantID:       tenantID,
		OrganizationID: organizationID,
		ClusterID:      clusterID,
		PoolID:         poolID,
		PoolStatus:     meta.TenantPoolBindingFree,
		CreatedAt:      now,
		UpdatedAt:      now,
	}); err != nil {
		t.Fatalf("upsert binding %s: %v", tenantID, err)
	}
	return tenantID
}

type deadlineMetadataResumeProvisioner struct {
	*adminTenantPoolSchemaInitRecorder

	waitStarted     chan struct{}
	waitStartedOnce sync.Once
}

func (p *deadlineMetadataResumeProvisioner) WaitForPoolClustersMetadata(ctx context.Context, clusters []*tenant.ClusterInfo, req tenant.CredentialProvisionRequest) ([]*tenant.ClusterInfo, error) {
	p.waitStartedOnce.Do(func() { close(p.waitStarted) })
	<-ctx.Done()
	return p.quotaTestProvisioner.WaitForPoolClustersMetadata(context.Background(), clusters, req)
}

type groupStreamingMetadataResumeProvisioner struct {
	*adminTenantPoolSchemaInitRecorder

	slowTenantID string
	slowStarted  chan struct{}
	releaseSlow  chan struct{}
	slowOnce     sync.Once
}

func (p *groupStreamingMetadataResumeProvisioner) WaitForPoolClustersMetadata(ctx context.Context, clusters []*tenant.ClusterInfo, req tenant.CredentialProvisionRequest) ([]*tenant.ClusterInfo, error) {
	for _, cluster := range clusters {
		if cluster != nil && cluster.TenantID == p.slowTenantID {
			p.slowOnce.Do(func() { close(p.slowStarted) })
			select {
			case <-p.releaseSlow:
			case <-ctx.Done():
				return nil, ctx.Err()
			}
			break
		}
	}
	return p.quotaTestProvisioner.WaitForPoolClustersMetadata(ctx, clusters, req)
}

func (p *adminTenantPoolSchemaInitRecorder) InitSchema(_ context.Context, dsn string) error {
	p.mu.Lock()
	p.lastSchemaInitDSN = dsn
	p.mu.Unlock()
	p.schemaInitCalls.Add(1)
	return nil
}

func (p *adminTenantPoolSchemaInitRecorder) lastSchemaInitDSNSnapshot() string {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.lastSchemaInitDSN
}

func assertTenantUsesPrivateEndpointTLS(t *testing.T, got meta.Tenant) {
	t.Helper()
	dsn := tenantDSN("u.root", "pass", got.DBHost, got.DBPort, got.DBName, got.DBTLS, got.Provider)
	if !strings.Contains(dsn, "tls=skip-verify") {
		t.Fatalf("tenant DSN = %q, want tls=skip-verify", dsn)
	}
	if strings.Contains(dsn, "tls=true") {
		t.Fatalf("tenant DSN = %q, should not use tls=true for private endpoint", dsn)
	}
}

func assertSchemaInitUsesPrivateEndpointTLS(t *testing.T, dsn string) {
	t.Helper()
	if !strings.Contains(dsn, "tls=skip-verify") {
		t.Fatalf("schema init DSN = %q, want tls=skip-verify", dsn)
	}
	if strings.Contains(dsn, "tls=true") {
		t.Fatalf("schema init DSN = %q, should not use tls=true for private endpoint", dsn)
	}
}