domainfront/front_test.go at main · getlantern/domainfront · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
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_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)
}