Fix more tests.

Jigsaw-Code · Mar 14, 2024 · e8d0fad · e8d0fad
1 parent db1be10
commit e8d0fad
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/cmd/outline-ss-server/metrics.go b/cmd/outline-ss-server/metrics.go
@@ -278,12 +278,12 @@ func (m *outlineMetrics) SetNumAccessKeys(numKeys int, ports int) {
 	m.ports.Set(float64(ports))
 }
 
-func (m *outlineMetrics) AddOpenTCPConnection(ip net.Addr) {
-	clientInfo, err := ipinfo.GetIPInfoFromAddr(m.IPInfoMap, ip)
+func (m *outlineMetrics) AddOpenTCPConnection(addr net.Addr) {
+	clientInfo, err := ipinfo.GetIPInfoFromAddr(m.IPInfoMap, addr)
 	if err != nil {
 		logger.Warningf("Failed client info lookup: %v", err)
 	}
-	logger.Debugf("Got info \"%#v\" for IP %v", clientInfo, ip.String())
+	logger.Debugf("Got info \"%#v\" for IP %v", clientInfo, addr.String())
 	m.tcpOpenConnections.WithLabelValues(clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN)).Inc()
 }
 

diff --git a/cmd/outline-ss-server/metrics_test.go b/cmd/outline-ss-server/metrics_test.go
@@ -24,6 +24,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type fakeAddr string
+
+func (a fakeAddr) String() string  { return string(a) }
+func (a fakeAddr) Network() string { return "" }
+
 func TestMethodsDontPanic(t *testing.T) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewPedanticRegistry())
 	proxyMetrics := metrics.ProxyMetrics{
@@ -34,12 +39,12 @@ func TestMethodsDontPanic(t *testing.T) {
 	}
 	ssMetrics.SetBuildInfo("0.0.0-test")
 	ssMetrics.SetNumAccessKeys(20, 2)
-	ssMetrics.AddOpenTCPConnection(ipinfo.IPInfo{CountryCode: "US", ASN: 100})
-	ssMetrics.AddClosedTCPConnection(ipinfo.IPInfo{CountryCode: "US", ASN: 100}, "1", "OK", proxyMetrics, 10*time.Millisecond)
+	ssMetrics.AddOpenTCPConnection(fakeAddr("127.0.0.1:9"))
+	ssMetrics.AddClosedTCPConnection(fakeAddr("127.0.0.1:9"), "1", "OK", proxyMetrics, 10*time.Millisecond)
 	ssMetrics.AddUDPPacketFromClient(ipinfo.IPInfo{CountryCode: "US", ASN: 100}, "2", "OK", 10, 20)
 	ssMetrics.AddUDPPacketFromTarget(ipinfo.IPInfo{CountryCode: "US", ASN: 100}, "3", "OK", 10, 20)
-	ssMetrics.AddUDPNatEntry()
-	ssMetrics.RemoveUDPNatEntry()
+	ssMetrics.AddUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-1")
+	ssMetrics.RemoveUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-1")
 	ssMetrics.AddTCPProbe("ERR_CIPHER", "eof", 443, proxyMetrics.ClientProxy)
 	ssMetrics.AddTCPCipherSearch(true, 10*time.Millisecond)
 	ssMetrics.AddUDPCipherSearch(true, 10*time.Millisecond)
@@ -54,21 +59,21 @@ func BenchmarkOpenTCP(b *testing.B) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewRegistry())
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		ssMetrics.AddOpenTCPConnection(ipinfo.IPInfo{CountryCode: "ZZ", ASN: 100})
+		ssMetrics.AddOpenTCPConnection(fakeAddr("127.0.0.1:9"))
 	}
 }
 
 func BenchmarkCloseTCP(b *testing.B) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewRegistry())
-	clientInfo := ipinfo.IPInfo{CountryCode: "ZZ", ASN: 100}
+	addr := fakeAddr("127.0.0.1:9")
 	accessKey := "key 1"
 	status := "OK"
 	data := metrics.ProxyMetrics{}
 	timeToCipher := time.Microsecond
 	duration := time.Minute
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		ssMetrics.AddClosedTCPConnection(clientInfo, accessKey, status, data, duration)
+		ssMetrics.AddClosedTCPConnection(addr, accessKey, status, data, duration)
 		ssMetrics.AddTCPCipherSearch(true, timeToCipher)
 	}
 }
@@ -115,7 +120,7 @@ func BenchmarkNAT(b *testing.B) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewRegistry())
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		ssMetrics.AddUDPNatEntry()
-		ssMetrics.RemoveUDPNatEntry()
+		ssMetrics.AddUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-0")
+		ssMetrics.RemoveUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-0")
 	}
 }
diff --git a/internal/integration_test/integration_test.go b/internal/integration_test/integration_test.go
@@ -167,7 +167,7 @@ type statusMetrics struct {
 	statuses []string
 }
 
-func (m *statusMetrics) AddClosedTCPConnection(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics, duration time.Duration) {
+func (m *statusMetrics) AddClosedTCPConnection(ip net.Addr, accessKey string, status string, data metrics.ProxyMetrics, duration time.Duration) {
 	m.Lock()
 	m.statuses = append(m.statuses, status)
 	m.Unlock()
@@ -247,10 +247,10 @@ func (m *fakeUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, access
 func (m *fakeUDPMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
 	m.down = append(m.down, udpRecord{clientInfo, accessKey, status, targetProxyBytes, proxyClientBytes})
 }
-func (m *fakeUDPMetrics) AddUDPNatEntry() {
+func (m *fakeUDPMetrics) AddUDPNatEntry(addr net.Addr, accessKey string) {
 	m.natAdded++
 }
-func (m *fakeUDPMetrics) RemoveUDPNatEntry() {
+func (m *fakeUDPMetrics) RemoveUDPNatEntry(addr net.Addr, accessKey string) {
 	// Not tested because it requires waiting for a long timeout.
 }
 func (m *fakeUDPMetrics) AddUDPCipherSearch(accessKeyFound bool, timeToCipher time.Duration) {}