From 39cfc982b5ef64d223a24de45105cf5ed973f0f8 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sat, 17 Feb 2018 23:37:09 +0100
Subject: [PATCH] test case for policy

---
 app/policy/config.go       | 18 ------------------
 app/policy/manager.go      | 14 +++++---------
 app/policy/manager_test.go | 37 +++++++++++++++++++++++++++++++++++++
 policy.go                  |  2 +-
 4 files changed, 43 insertions(+), 28 deletions(-)
 create mode 100644 app/policy/manager_test.go

diff --git a/app/policy/config.go b/app/policy/config.go
index 4d2f9452..589a2c48 100644
--- a/app/policy/config.go
+++ b/app/policy/config.go
@@ -14,24 +14,6 @@ func (s *Second) Duration() time.Duration {
 	return time.Second * time.Duration(s.Value)
 }
 
-// OverrideWith overrides current Policy with another one.
-func (p *Policy) OverrideWith(another *Policy) {
-	if another.Timeout != nil {
-		if another.Timeout.Handshake != nil {
-			p.Timeout.Handshake = another.Timeout.Handshake
-		}
-		if another.Timeout.ConnectionIdle != nil {
-			p.Timeout.ConnectionIdle = another.Timeout.ConnectionIdle
-		}
-		if another.Timeout.UplinkOnly != nil {
-			p.Timeout.UplinkOnly = another.Timeout.UplinkOnly
-		}
-		if another.Timeout.DownlinkOnly != nil {
-			p.Timeout.DownlinkOnly = another.Timeout.DownlinkOnly
-		}
-	}
-}
-
 func (p *Policy) ToCorePolicy() core.Policy {
 	var cp core.Policy
 	if p.Timeout != nil {
diff --git a/app/policy/manager.go b/app/policy/manager.go
index 14f7b7b9..14a1c4c7 100644
--- a/app/policy/manager.go
+++ b/app/policy/manager.go
@@ -19,19 +19,15 @@ func New(ctx context.Context, config *Config) (*Instance, error) {
 	}
 	if len(config.Level) > 0 {
 		for lv, p := range config.Level {
-			dp := core.DefaultPolicy()
-			dp.OverrideWith(p.ToCorePolicy())
-			m.levels[lv] = dp
+			m.levels[lv] = p.ToCorePolicy().OverrideWith(core.DefaultPolicy())
 		}
 	}
 
 	v := core.FromContext(ctx)
-	if v == nil {
-		return nil, newError("V is not in context.")
-	}
-
-	if err := v.RegisterFeature((*core.PolicyManager)(nil), m); err != nil {
-		return nil, newError("unable to register PolicyManager in core").Base(err).AtError()
+	if v != nil {
+		if err := v.RegisterFeature((*core.PolicyManager)(nil), m); err != nil {
+			return nil, newError("unable to register PolicyManager in core").Base(err).AtError()
+		}
 	}
 
 	return m, nil
diff --git a/app/policy/manager_test.go b/app/policy/manager_test.go
new file mode 100644
index 00000000..d56afb86
--- /dev/null
+++ b/app/policy/manager_test.go
@@ -0,0 +1,37 @@
+package policy_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"v2ray.com/core"
+	. "v2ray.com/core/app/policy"
+	. "v2ray.com/ext/assert"
+)
+
+func TestPolicy(t *testing.T) {
+	assert := With(t)
+
+	manager, err := New(context.Background(), &Config{
+		Level: map[uint32]*Policy{
+			0: &Policy{
+				Timeout: &Policy_Timeout{
+					Handshake: &Second{
+						Value: 2,
+					},
+				},
+			},
+		},
+	})
+	assert(err, IsNil)
+
+	pDefault := core.DefaultPolicy()
+
+	p0 := manager.ForLevel(0)
+	assert(p0.Timeouts.Handshake, Equals, 2*time.Second)
+	assert(p0.Timeouts.ConnectionIdle, Equals, pDefault.Timeouts.ConnectionIdle)
+
+	p1 := manager.ForLevel(1)
+	assert(p1.Timeouts.Handshake, Equals, pDefault.Timeouts.Handshake)
+}
diff --git a/policy.go b/policy.go
index a0717320..a333a84b 100644
--- a/policy.go
+++ b/policy.go
@@ -43,7 +43,7 @@ type Policy struct {
 
 // OverrideWith overrides the current Policy with another one. All values with default value will be overridden.
 func (p Policy) OverrideWith(another Policy) Policy {
-	p.Timeouts.OverrideWith(another.Timeouts)
+	p.Timeouts = p.Timeouts.OverrideWith(another.Timeouts)
 	return p
 }