github.com/ipfans/trojan-go@v0.11.0/statistic/memory/memory_test.go (about)

     1  package memory
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"strconv"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/ipfans/trojan-go/common"
    11  	"github.com/ipfans/trojan-go/config"
    12  )
    13  
    14  func TestMemoryAuth(t *testing.T) {
    15  	cfg := &Config{
    16  		Passwords: nil,
    17  	}
    18  	ctx := config.WithConfig(context.Background(), Name, cfg)
    19  	auth, err := NewAuthenticator(ctx)
    20  	common.Must(err)
    21  	auth.AddUser("user1")
    22  	valid, user := auth.AuthUser("user1")
    23  	if !valid {
    24  		t.Fatal("add, auth")
    25  	}
    26  	if user.Hash() != "user1" {
    27  		t.Fatal("Hash")
    28  	}
    29  	user.AddTraffic(100, 200)
    30  	sent, recv := user.GetTraffic()
    31  	if sent != 100 || recv != 200 {
    32  		t.Fatal("traffic")
    33  	}
    34  	sent, recv = user.ResetTraffic()
    35  	if sent != 100 || recv != 200 {
    36  		t.Fatal("ResetTraffic")
    37  	}
    38  	sent, recv = user.GetTraffic()
    39  	if sent != 0 || recv != 0 {
    40  		t.Fatal("ResetTraffic")
    41  	}
    42  
    43  	user.AddIP("1234")
    44  	user.AddIP("5678")
    45  	if user.GetIP() != 0 {
    46  		t.Fatal("GetIP")
    47  	}
    48  
    49  	user.SetIPLimit(2)
    50  	user.AddIP("1234")
    51  	user.AddIP("5678")
    52  	user.DelIP("1234")
    53  	if user.GetIP() != 1 {
    54  		t.Fatal("DelIP")
    55  	}
    56  	user.DelIP("5678")
    57  
    58  	user.SetIPLimit(2)
    59  	if !user.AddIP("1") || !user.AddIP("2") {
    60  		t.Fatal("AddIP")
    61  	}
    62  	if user.AddIP("3") {
    63  		t.Fatal("AddIP")
    64  	}
    65  	if !user.AddIP("2") {
    66  		t.Fatal("AddIP")
    67  	}
    68  
    69  	user.SetTraffic(1234, 4321)
    70  	if a, b := user.GetTraffic(); a != 1234 || b != 4321 {
    71  		t.Fatal("SetTraffic")
    72  	}
    73  
    74  	user.ResetTraffic()
    75  	go func() {
    76  		for {
    77  			k := 100
    78  			time.Sleep(time.Second / time.Duration(k))
    79  			user.AddTraffic(2000/k, 1000/k)
    80  		}
    81  	}()
    82  	time.Sleep(time.Second * 4)
    83  	if sent, recv := user.GetSpeed(); sent > 3000 || sent < 1000 || recv > 1500 || recv < 500 {
    84  		t.Error("GetSpeed", sent, recv)
    85  	} else {
    86  		t.Log("GetSpeed", sent, recv)
    87  	}
    88  
    89  	user.SetSpeedLimit(30, 20)
    90  	time.Sleep(time.Second * 4)
    91  	if sent, recv := user.GetSpeed(); sent > 60 || recv > 40 {
    92  		t.Error("SetSpeedLimit", sent, recv)
    93  	} else {
    94  		t.Log("SetSpeedLimit", sent, recv)
    95  	}
    96  
    97  	user.SetSpeedLimit(0, 0)
    98  	time.Sleep(time.Second * 4)
    99  	if sent, recv := user.GetSpeed(); sent < 30 || recv < 20 {
   100  		t.Error("SetSpeedLimit", sent, recv)
   101  	} else {
   102  		t.Log("SetSpeedLimit", sent, recv)
   103  	}
   104  
   105  	auth.AddUser("user2")
   106  	valid, _ = auth.AuthUser("user2")
   107  	if !valid {
   108  		t.Fatal()
   109  	}
   110  	auth.DelUser("user2")
   111  	valid, _ = auth.AuthUser("user2")
   112  	if valid {
   113  		t.Fatal()
   114  	}
   115  	auth.AddUser("user3")
   116  	users := auth.ListUsers()
   117  	if len(users) != 2 {
   118  		t.Fatal()
   119  	}
   120  	user.Close()
   121  	auth.Close()
   122  }
   123  
   124  func BenchmarkMemoryUsage(b *testing.B) {
   125  	cfg := &Config{
   126  		Passwords: nil,
   127  	}
   128  	ctx := config.WithConfig(context.Background(), Name, cfg)
   129  	auth, err := NewAuthenticator(ctx)
   130  	common.Must(err)
   131  
   132  	m1 := runtime.MemStats{}
   133  	m2 := runtime.MemStats{}
   134  	runtime.ReadMemStats(&m1)
   135  	for i := 0; i < b.N; i++ {
   136  		common.Must(auth.AddUser(common.SHA224String("hash" + strconv.Itoa(i))))
   137  	}
   138  	runtime.ReadMemStats(&m2)
   139  
   140  	b.ReportMetric(float64(m2.Alloc-m1.Alloc)/1024/1024, "MiB(Alloc)")
   141  	b.ReportMetric(float64(m2.TotalAlloc-m1.TotalAlloc)/1024/1024, "MiB(TotalAlloc)")
   142  }