From 6ddafdbd7317a13aee79e0c970c548603f041636 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 23 Oct 2017 23:47:20 +0200
Subject: [PATCH] net assert

---
 common/net/address_test.go   | 14 ++++++-----
 common/net/testing/assert.go | 46 ++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 6 deletions(-)
 create mode 100644 common/net/testing/assert.go

diff --git a/common/net/address_test.go b/common/net/address_test.go
index 424727314..58583d248 100644
--- a/common/net/address_test.go
+++ b/common/net/address_test.go
@@ -5,20 +5,22 @@ import (
 	"testing"
 
 	. "v2ray.com/core/common/net"
+	. "v2ray.com/core/common/net/testing"
 	"v2ray.com/core/testing/assert"
+	. "v2ray.com/ext/assert"
 )
 
 func TestIPv4Address(t *testing.T) {
-	assert := assert.On(t)
+	assert := With(t)
 
 	ip := []byte{byte(1), byte(2), byte(3), byte(4)}
 	addr := IPAddress(ip)
 
-	assert.Address(addr).IsIPv4()
-	assert.Address(addr).IsNotIPv6()
-	assert.Address(addr).IsNotDomain()
-	assert.Bytes(addr.IP()).Equals(ip)
-	assert.Address(addr).EqualsString("1.2.3.4")
+	assert(addr, IsIPv4)
+	assert(addr, Not(IsIPv6))
+	assert(addr, Not(IsDomain))
+	assert([]byte(addr.IP()), Equals, ip)
+	assert(addr.String(), Equals, "1.2.3.4")
 }
 
 func TestIPv6Address(t *testing.T) {
diff --git a/common/net/testing/assert.go b/common/net/testing/assert.go
new file mode 100644
index 000000000..9256d049e
--- /dev/null
+++ b/common/net/testing/assert.go
@@ -0,0 +1,46 @@
+package testing
+
+import (
+	"reflect"
+
+	"v2ray.com/core/common/net"
+	"v2ray.com/ext/assert"
+)
+
+var IsIPv4 = assert.CreateOperator(reflect.TypeOf(net.Address(nil)), reflect.ValueOf(func(a net.Address) bool {
+	return a.Family().IsIPv4()
+}), 1, "is IPv4")
+
+var IsIPv6 = assert.CreateOperator(reflect.TypeOf(net.Address(nil)), reflect.ValueOf(func(a net.Address) bool {
+	return a.Family().IsIPv6()
+}), 1, "is IPv6")
+
+var IsIP = assert.CreateOperator(reflect.TypeOf(net.Address(nil)), reflect.ValueOf(func(a net.Address) bool {
+	return a.Family().IsIPv4() || a.Family().IsIPv6()
+}), 1, "is IP")
+
+var IsTCP = assert.CreateOperator(reflect.TypeOf(net.Destination{}), reflect.ValueOf(func(a net.Destination) bool {
+	return a.Network == net.Network_TCP
+}), 1, "is TCP")
+
+var IsUDP = assert.CreateOperator(reflect.TypeOf(net.Destination{}), reflect.ValueOf(func(a net.Destination) bool {
+	return a.Network == net.Network_UDP
+}), 1, "is UDP")
+
+var IsDomain = assert.CreateOperator(reflect.TypeOf(net.Address(nil)), reflect.ValueOf(func(a net.Address) bool {
+	return a.Family().IsDomain()
+}), 1, "is Domain")
+
+func init() {
+	assert.RegisterEqualsOperator(reflect.TypeOf((*net.Address)(nil)).Elem(), reflect.ValueOf(func(a, b net.Address) bool {
+		return a == b
+	}))
+
+	assert.RegisterEqualsOperator(reflect.TypeOf(net.Destination{}), reflect.ValueOf(func(a, b net.Destination) bool {
+		return a == b
+	}))
+
+	assert.RegisterEqualsOperator(reflect.TypeOf(net.Port(0)), reflect.ValueOf(func(a, b net.Port) bool {
+		return a == b
+	}))
+}