Zoraxy/mod/statistic/statistic_test.go

package statistic_test

import (
	"net"
	"os"
	"testing"
	"time"

	"math/rand"

	"imuslab.com/zoraxy/mod/database"
	"imuslab.com/zoraxy/mod/database/dbinc"
	"imuslab.com/zoraxy/mod/geodb"
	"imuslab.com/zoraxy/mod/statistic"
)

const test_db_path = "test_db"

func getNewDatabase() *database.Database {
	db, err := database.NewDatabase(test_db_path, dbinc.BackendLevelDB)
	if err != nil {
		panic(err)
	}
	db.NewTable("stats")
	return db
}

func clearDatabase(db *database.Database) {
	db.Close()
	os.RemoveAll(test_db_path)
}

func TestNewStatisticCollector(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, err := statistic.NewStatisticCollector(option)
	if err != nil {
		t.Fatalf("Expected no error, got %v", err)
	}
	if collector == nil {
		t.Fatalf("Expected collector, got nil")
	}

}

func TestSaveSummaryOfDay(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	collector.SaveSummaryOfDay()
	// Add assertions to check if data is saved correctly
}

func TestLoadSummaryOfDay(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	year, month, day := time.Now().Date()
	summary := collector.LoadSummaryOfDay(year, month, day)
	if summary == nil {
		t.Fatalf("Expected summary, got nil")
	}
}

func TestResetSummaryOfDay(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	collector.ResetSummaryOfDay()
	if collector.DailySummary.TotalRequest != 0 {
		t.Fatalf("Expected TotalRequest to be 0, got %v", collector.DailySummary.TotalRequest)
	}
}

func TestGetCurrentRealtimeStatIntervalId(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	intervalId := collector.GetCurrentRealtimeStatIntervalId()
	if intervalId < 0 || intervalId > 287 {
		t.Fatalf("Expected intervalId to be between 0 and 287, got %v", intervalId)
	}
}

func TestRecordRequest(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	requestInfo := statistic.RequestInfo{
		IpAddr:                        "127.0.0.1",
		RequestOriginalCountryISOCode: "US",
		Succ:                          true,
		StatusCode:                    200,
		ForwardType:                   "type1",
		Referer:                       "http://example.com",
		UserAgent:                     "Mozilla/5.0",
		RequestURL:                    "/test",
		Target:                        "target1",
	}
	collector.RecordRequest(requestInfo)
	time.Sleep(1 * time.Second) // Wait for the goroutine to finish
	if collector.DailySummary.TotalRequest != 1 {
		t.Fatalf("Expected TotalRequest to be 1, got %v", collector.DailySummary.TotalRequest)
	}
}

func TestScheduleResetRealtimeStats(t *testing.T) {
	db := getNewDatabase()
	defer clearDatabase(db)
	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	stopChan := collector.ScheduleResetRealtimeStats()
	if stopChan == nil {
		t.Fatalf("Expected stopChan, got nil")
	}
	collector.Close()
}

func TestNewDailySummary(t *testing.T) {
	summary := statistic.NewDailySummary()
	if summary.TotalRequest != 0 {
		t.Fatalf("Expected TotalRequest to be 0, got %v", summary.TotalRequest)
	}
	if summary.ForwardTypes == nil {
		t.Fatalf("Expected ForwardTypes to be initialized, got nil")
	}
	if summary.RequestOrigin == nil {
		t.Fatalf("Expected RequestOrigin to be initialized, got nil")
	}
	if summary.RequestClientIp == nil {
		t.Fatalf("Expected RequestClientIp to be initialized, got nil")
	}
	if summary.Referer == nil {
		t.Fatalf("Expected Referer to be initialized, got nil")
	}
	if summary.UserAgent == nil {
		t.Fatalf("Expected UserAgent to be initialized, got nil")
	}
	if summary.RequestURL == nil {
		t.Fatalf("Expected RequestURL to be initialized, got nil")
	}
}

func generateTestRequestInfo(db *database.Database) statistic.RequestInfo {
	//Generate a random IPv4 address
	randomIpAddr := ""
	for {
		ip := net.IPv4(byte(rand.Intn(256)), byte(rand.Intn(256)), byte(rand.Intn(256)), byte(rand.Intn(256)))
		if !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsMulticast() && !ip.IsUnspecified() {
			randomIpAddr = ip.String()
			break
		}
	}

	//Resolve the country code for this IP
	ipLocation := "unknown"
	geoIpResolver, err := geodb.NewGeoDb(db, &geodb.StoreOptions{
		AllowSlowIpv4LookUp: false,
		AllowSlowIpv6Lookup: true, //Just to save some RAM
	})

	if err == nil {
		ipInfo, _ := geoIpResolver.ResolveCountryCodeFromIP(randomIpAddr)
		ipLocation = ipInfo.CountryIsoCode
	}

	forwardType := "host-http"
	//Generate a random forward type between "subdomain-http" and "host-https"
	if rand.Intn(2) == 1 {
		forwardType = "subdomain-http"
	}

	//Generate 5 random refers URL and pick from there
	referers := []string{"https://example.com", "https://example.org", "https://example.net", "https://example.io", "https://example.co"}
	referer := referers[rand.Intn(5)]

	return statistic.RequestInfo{
		IpAddr:                        randomIpAddr,
		RequestOriginalCountryISOCode: ipLocation,
		Succ:                          true,
		StatusCode:                    200,
		ForwardType:                   forwardType,
		Referer:                       referer,
		UserAgent:                     "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36",
		RequestURL:                    "/benchmark",
		Target:                        "test.imuslab.internal",
	}
}

func BenchmarkRecordRequest(b *testing.B) {
	db := getNewDatabase()
	defer clearDatabase(db)

	option := statistic.CollectorOption{Database: db}
	collector, _ := statistic.NewStatisticCollector(option)
	var requestInfo statistic.RequestInfo = generateTestRequestInfo(db)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		collector.RecordRequest(requestInfo)
		collector.SaveSummaryOfDay()
	}

	//Write the current in-memory summary to database file
	b.StopTimer()

	//Print the generated summary
	//testSummary := collector.GetCurrentDailySummary()
	//statistic.PrintDailySummary(testSummary)
}