Feature: persistence fakeip (#1662)

component/fakeip/pool_test.go

@@ -2,38 +2,118 @@ package fakeip
 
 import (
 	"net"
+	"os"
 	"testing"
+	"time"
+
+	"github.com/Dreamacro/clash/component/profile/cachefile"
+	"github.com/Dreamacro/clash/component/trie"
 
 	"github.com/stretchr/testify/assert"
+	"go.etcd.io/bbolt"
 )
 
+func createPools(options Options) ([]*Pool, string, error) {
+	pool, err := New(options)
+	if err != nil {
+		return nil, "", err
+	}
+	filePool, tempfile, err := createCachefileStore(options)
+	if err != nil {
+		return nil, "", err
+	}
+
+	return []*Pool{pool, filePool}, tempfile, nil
+}
+
+func createCachefileStore(options Options) (*Pool, string, error) {
+	pool, err := New(options)
+	if err != nil {
+		return nil, "", err
+	}
+	f, err := os.CreateTemp("", "clash")
+	if err != nil {
+		return nil, "", err
+	}
+
+	db, err := bbolt.Open(f.Name(), 0o666, &bbolt.Options{Timeout: time.Second})
+	if err != nil {
+		return nil, "", err
+	}
+
+	pool.store = &cachefileStore{
+		cache: &cachefile.CacheFile{DB: db},
+	}
+	return pool, f.Name(), nil
+}
+
 func TestPool_Basic(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/29")
-	pool, _ := New(ipnet, 10, nil)
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
 
-	first := pool.Lookup("foo.com")
-	last := pool.Lookup("bar.com")
-	bar, exist := pool.LookBack(last)
+	for _, pool := range pools {
+		first := pool.Lookup("foo.com")
+		last := pool.Lookup("bar.com")
+		bar, exist := pool.LookBack(last)
 
-	assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
-	assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
-	assert.True(t, exist)
-	assert.Equal(t, bar, "bar.com")
+		assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
+		assert.Equal(t, pool.Lookup("foo.com"), net.IP{192, 168, 0, 2})
+		assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
+		assert.True(t, exist)
+		assert.Equal(t, bar, "bar.com")
+		assert.Equal(t, pool.Gateway(), net.IP{192, 168, 0, 1})
+		assert.Equal(t, pool.IPNet().String(), ipnet.String())
+		assert.True(t, pool.Exist(net.IP{192, 168, 0, 3}))
+		assert.False(t, pool.Exist(net.IP{192, 168, 0, 4}))
+		assert.False(t, pool.Exist(net.ParseIP("::1")))
+	}
 }
 
-func TestPool_Cycle(t *testing.T) {
+func TestPool_CycleUsed(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/30")
-	pool, _ := New(ipnet, 10, nil)
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
 
-	first := pool.Lookup("foo.com")
-	same := pool.Lookup("baz.com")
+	for _, pool := range pools {
+		first := pool.Lookup("foo.com")
+		same := pool.Lookup("baz.com")
+		assert.True(t, first.Equal(same))
+	}
+}
 
-	assert.True(t, first.Equal(same))
+func TestPool_Skip(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/30")
+	tree := trie.New()
+	tree.Insert("example.com", tree)
+	pools, tempfile, err := createPools(Options{
+		IPNet: ipnet,
+		Size:  10,
+		Host:  tree,
+	})
+	assert.Nil(t, err)
+	defer os.Remove(tempfile)
+
+	for _, pool := range pools {
+		assert.True(t, pool.ShouldSkipped("example.com"))
+		assert.False(t, pool.ShouldSkipped("foo.com"))
+	}
 }
 
 func TestPool_MaxCacheSize(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
-	pool, _ := New(ipnet, 2, nil)
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
 
 	first := pool.Lookup("foo.com")
 	pool.Lookup("bar.com")
@@ -45,7 +125,10 @@ func TestPool_MaxCacheSize(t *testing.T) {
 
 func TestPool_DoubleMapping(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
-	pool, _ := New(ipnet, 2, nil)
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
 
 	// fill cache
 	fooIP := pool.Lookup("foo.com")
@@ -70,9 +153,35 @@ func TestPool_DoubleMapping(t *testing.T) {
 	assert.False(t, bazIP.Equal(newBazIP))
 }
 
+func TestPool_Clone(t *testing.T) {
+	_, ipnet, _ := net.ParseCIDR("192.168.0.1/24")
+	pool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
+
+	first := pool.Lookup("foo.com")
+	last := pool.Lookup("bar.com")
+	assert.True(t, first.Equal(net.IP{192, 168, 0, 2}))
+	assert.True(t, last.Equal(net.IP{192, 168, 0, 3}))
+
+	newPool, _ := New(Options{
+		IPNet: ipnet,
+		Size:  2,
+	})
+	newPool.CloneFrom(pool)
+	_, firstExist := newPool.LookBack(first)
+	_, lastExist := newPool.LookBack(last)
+	assert.True(t, firstExist)
+	assert.True(t, lastExist)
+}
+
 func TestPool_Error(t *testing.T) {
 	_, ipnet, _ := net.ParseCIDR("192.168.0.1/31")
-	_, err := New(ipnet, 10, nil)
+	_, err := New(Options{
+		IPNet: ipnet,
+		Size:  10,
+	})
 
 	assert.Error(t, err)
 }