front_test.go

package domainfront

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestFrontPool_TakeReturn(t *testing.T) {
	pool := newFrontPool(0)
	f := newFront(&Masquerade{Domain: "cdn.example.com", IpAddress: "1.2.3.4"}, "test")
	f.markSucceeded()

	pool.Replace([]*front{f})
	pool.addReady(f)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	got, err := pool.Take(ctx)
	require.NoError(t, err)
	assert.Equal(t, "1.2.3.4", got.IpAddress)

	pool.ReturnSuccess(got)
	assert.Equal(t, 1, pool.readyCount())
	assert.True(t, got.isSucceeding(), "ReturnSuccess should mark as succeeded")
}

func TestFrontPool_ReturnRequeue(t *testing.T) {
	pool := newFrontPool(0)
	f := newFront(&Masquerade{Domain: "cdn.example.com", IpAddress: "1.2.3.4"}, "test")
	pool.addReady(f)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	got, _ := pool.Take(ctx)
	// Return with requeue=true should not update LastSucceeded
	pool.Return(got, true)
	assert.False(t, got.isSucceeding(), "Return(requeue=true) should not mark as succeeded")
	assert.Equal(t, 1, pool.readyCount())
}

func TestFrontPool_TakeBlocksUntilReady(t *testing.T) {
	pool := newFrontPool(0)
	f := newFront(&Masquerade{Domain: "cdn.example.com", IpAddress: "1.2.3.4"}, "test")

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	// Should timeout since no fronts are ready
	_, err := pool.Take(ctx)
	assert.ErrorIs(t, err, context.DeadlineExceeded)

	// Now add a ready front
	ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
	defer cancel2()

	go func() {
		time.Sleep(50 * time.Millisecond)
		pool.addReady(f)
	}()

	got, err := pool.Take(ctx2)
	require.NoError(t, err)
	assert.Equal(t, "1.2.3.4", got.IpAddress)
}

func TestFrontPool_Replace_PreservesState(t *testing.T) {
	pool := newFrontPool(0)
	f1 := newFront(&Masquerade{Domain: "cdn.example.com", IpAddress: "1.2.3.4"}, "test")
	f1.markSucceeded()
	pool.Replace([]*front{f1})

	// Replace with same front — state should be preserved
	f2 := newFront(&Masquerade{Domain: "cdn.example.com", IpAddress: "1.2.3.4"}, "test")
	pool.Replace([]*front{f2})

	assert.True(t, f2.isSucceeding(), "state should be preserved after Replace")
}

func TestFrontPool_Close(t *testing.T) {
	pool := newFrontPool(0)

	ctx := context.Background()
	go func() {
		time.Sleep(50 * time.Millisecond)
		pool.Close()
	}()

	_, err := pool.Take(ctx)
	assert.Error(t, err)
}

func TestFrontPool_TakePreferringUntried(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	t.Run("prefers a front from an untried provider", func(t *testing.T) {
		pool := newFrontPool(0)
		aliyun := newFront(&Masquerade{Domain: "a.com", IpAddress: "1.1.1.1"}, "aliyun")
		cloudfront := newFront(&Masquerade{Domain: "c.com", IpAddress: "2.2.2.2"}, "cloudfront")
		// aliyun ahead of cloudfront in the ready queue: a plain Take would pick
		// aliyun again even though we've already tried it.
		pool.addReady(aliyun)
		pool.addReady(cloudfront)

		got, err := pool.TakePreferringUntried(ctx, map[string]struct{}{"aliyun": {}})
		require.NoError(t, err)
		assert.Equal(t, "cloudfront", got.ProviderID, "should skip the already-tried provider")
		// The skipped aliyun front must be returned to the queue, not dropped.
		assert.Equal(t, 1, pool.readyCount())
	})

	t.Run("falls back to first when every ready provider was tried", func(t *testing.T) {
		pool := newFrontPool(0)
		f1 := newFront(&Masquerade{Domain: "a.com", IpAddress: "1.1.1.1"}, "aliyun")
		f2 := newFront(&Masquerade{Domain: "b.com", IpAddress: "3.3.3.3"}, "aliyun")
		pool.addReady(f1)
		pool.addReady(f2)

		got, err := pool.TakePreferringUntried(ctx, map[string]struct{}{"aliyun": {}})
		require.NoError(t, err)
		assert.Equal(t, "aliyun", got.ProviderID)
		// The other tried front is requeued; only the selected one is consumed.
		assert.Equal(t, 1, pool.readyCount())
	})

	t.Run("empty tried map behaves like Take", func(t *testing.T) {
		pool := newFrontPool(0)
		f := newFront(&Masquerade{Domain: "a.com", IpAddress: "1.1.1.1"}, "aliyun")
		pool.addReady(f)

		got, err := pool.TakePreferringUntried(ctx, nil)
		require.NoError(t, err)
		assert.Equal(t, "1.1.1.1", got.IpAddress)
		assert.Equal(t, 0, pool.readyCount())
	})

	t.Run("blocks until a front is available, then honors ctx cancel", func(t *testing.T) {
		pool := newFrontPool(0)
		deadline, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
		defer cancel()
		_, err := pool.TakePreferringUntried(deadline, map[string]struct{}{"aliyun": {}})
		assert.ErrorIs(t, err, context.DeadlineExceeded)
	})
}

func TestFrontPool_Candidates(t *testing.T) {
	pool := newFrontPool(0)

	f1 := newFront(&Masquerade{Domain: "a.com", IpAddress: "1.1.1.1"}, "p1")
	f2 := newFront(&Masquerade{Domain: "b.com", IpAddress: "2.2.2.2"}, "p2")
	f2.markSucceeded()

	pool.Replace([]*front{f1, f2})

	candidates := pool.candidates()
	require.Len(t, candidates, 2)
	// Succeeded front should be first
	assert.Equal(t, "2.2.2.2", candidates[0].IpAddress)
}