From cef1836f5a0638242be288508e810117bc71f322 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sat, 27 Jun 2020 00:28:53 -0400
Subject: [PATCH] Add unit test for outbound handler

---
 app/proxyman/outbound/handler_test.go | 67 +++++++++++++++++++++++++++
 context.go                            |  5 +-
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/app/proxyman/outbound/handler_test.go b/app/proxyman/outbound/handler_test.go
index 25906f17a..c05183784 100644
--- a/app/proxyman/outbound/handler_test.go
+++ b/app/proxyman/outbound/handler_test.go
@@ -1,13 +1,80 @@
 package outbound_test
 
 import (
+	"context"
 	"testing"
 
+	"v2ray.com/core"
+	"v2ray.com/core/app/policy"
 	. "v2ray.com/core/app/proxyman/outbound"
+	"v2ray.com/core/app/stats"
+	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/serial"
 	"v2ray.com/core/features/outbound"
+	"v2ray.com/core/proxy/freedom"
+	"v2ray.com/core/transport/internet"
 )
 
 func TestInterfaces(t *testing.T) {
 	_ = (outbound.Handler)(new(Handler))
 	_ = (outbound.Manager)(new(Manager))
 }
+
+const v2rayKey core.V2rayKey = 1
+
+func TestOutboundWithoutStatCounter(t *testing.T) {
+	config := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&stats.Config{}),
+			serial.ToTypedMessage(&policy.Config{
+				System: &policy.SystemPolicy{
+					Stats: &policy.SystemPolicy_Stats{
+						InboundUplink: true,
+					},
+				},
+			}),
+		},
+	}
+
+	v, _ := core.New(config)
+	v.AddFeature((outbound.Manager)(new(Manager)))
+	ctx := context.WithValue(context.Background(), v2rayKey, v)
+	h, _ := NewHandler(ctx, &core.OutboundHandlerConfig{
+		Tag:           "tag",
+		ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+	})
+	conn, _ := h.(*Handler).Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
+	_, ok := conn.(*internet.StatCouterConnection)
+	if ok {
+		t.Errorf("Expected conn to not be StatCouterConnection")
+	}
+}
+
+func TestOutboundWithStatCounter(t *testing.T) {
+	config := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&stats.Config{}),
+			serial.ToTypedMessage(&policy.Config{
+				System: &policy.SystemPolicy{
+					Stats: &policy.SystemPolicy_Stats{
+						OutboundUplink: true,
+						OutboundDownlink: true,
+					},
+				},
+			}),
+		},
+	}
+
+	v, _ := core.New(config)
+	v.AddFeature((outbound.Manager)(new(Manager)))
+	ctx := context.WithValue(context.Background(), v2rayKey, v)
+	h, _ := NewHandler(ctx, &core.OutboundHandlerConfig{
+		Tag:           "tag",
+		ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+	})
+	conn, _ := h.(*Handler).Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
+	_, ok := conn.(*internet.StatCouterConnection)
+	if !ok {
+		t.Errorf("Expected conn to be StatCouterConnection")
+	}
+}
diff --git a/context.go b/context.go
index 5dd7aee7a..84cd2396b 100644
--- a/context.go
+++ b/context.go
@@ -6,9 +6,10 @@ import (
 	"context"
 )
 
-type key int
+// V2rayKey is the key type of Instance in Context, exported for test.
+type V2rayKey int
 
-const v2rayKey key = 1
+const v2rayKey V2rayKey = 1
 
 // FromContext returns an Instance from the given context, or nil if the context doesn't contain one.
 func FromContext(ctx context.Context) *Instance {