使用KV数据库来管理IP名单

internal/iplibrary/ip_list_kv_test.go

@@ -0,0 +1,221 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary_test
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
+	"github.com/TeaOSLab/EdgeNode/internal/zero"
+	"testing"
+	"time"
+)
+
+func TestIPListKV_AddItem(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        1,
+			IpFrom:    "192.168.1.101",
+			IpTo:      "",
+			Version:   1,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        2,
+			IpFrom:    "192.168.1.102",
+			IpTo:      "",
+			Version:   2,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        3,
+			IpFrom:    "192.168.1.103",
+			IpTo:      "",
+			Version:   3,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestIPListKV_AddItems_Many(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	var count = 2
+	var from = 1
+	if testutils.IsSingleTesting() {
+		count = 2_000_000
+	}
+
+	var before = time.Now()
+	defer func() {
+		t.Logf("cost: %.2f s", time.Since(before).Seconds())
+	}()
+
+	for i := from; i <= from+count; i++ {
+		err = kv.AddItem(&pb.IPItem{
+			Id:        int64(i),
+			IpFrom:    testutils.RandIP(),
+			IpTo:      "",
+			Version:   int64(i),
+			ExpiredAt: fasttime.NewFastTime().Unix() + 86400,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestIPListKV_DeleteExpiredItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	err = kv.DeleteExpiredItems()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestIPListKV_UpdateMaxVersion(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	err = kv.UpdateMaxVersion(101)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maxVersion, err := kv.ReadMaxVersion()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log("version:", maxVersion)
+}
+
+func TestIPListKV_ReadMaxVersion(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maxVersion, err := kv.ReadMaxVersion()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log("version:", maxVersion)
+}
+
+func TestIPListKV_ReadItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for {
+		items, goNext, readErr := kv.ReadItems(0, 2)
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+		t.Log("====")
+		for _, item := range items {
+			t.Log(item.Id)
+		}
+
+		if !goNext {
+			break
+		}
+	}
+}
+
+func TestIPListKV_CountItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var count int
+	var m = map[int64]zero.Zero{}
+	for {
+		items, goNext, readErr := kv.ReadItems(0, 1000)
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+		for _, item := range items {
+			count++
+			m[item.Id] = zero.Zero{}
+		}
+
+		if !goNext {
+			break
+		}
+	}
+	t.Log("count:", count, "len:", len(m))
+}
+
+func TestIPListKV_Inspect(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = kv.TestInspect(t)
+	if err != nil {
+		t.Fatal(err)
+	}
+}