From 142abb08b3080b28c9d134b56be8bf528f20784a Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 11 Sep 2017 12:08:19 +0200
Subject: [PATCH] refactor drop tests, add test with dropped packets in both
 directions

---
 integrationtests/gquic/drop_test.go   | 54 ++++++++++++---------------
 integrationtests/tools/proxy/proxy.go | 22 +++++++++++
 2 files changed, 46 insertions(+), 30 deletions(-)

diff --git a/integrationtests/gquic/drop_test.go b/integrationtests/gquic/drop_test.go
index 2d4712613..1e0534b8d 100644
--- a/integrationtests/gquic/drop_test.go
+++ b/integrationtests/gquic/drop_test.go
@@ -16,6 +16,8 @@ import (
 	. "github.com/onsi/gomega/gexec"
 )
 
+var directions = []quicproxy.Direction{quicproxy.DirectionIncoming, quicproxy.DirectionOutgoing, quicproxy.DirectionBoth}
+
 var _ = Describe("Drop tests", func() {
 	var proxy *quicproxy.QuicProxy
 
@@ -47,41 +49,33 @@ var _ = Describe("Drop tests", func() {
 	})
 
 	Context("after the crypto handshake", func() {
-		for i := range protocol.SupportedVersions {
-			version := protocol.SupportedVersions[i]
-
-			Context(fmt.Sprintf("with quic version %d", version), func() {
-				dropTests("dropping every 4th packet", 4, 1, runDropTest, version)
-				dropTests("dropping 10 packets every 100 packets", 100, 10, runDropTest, version)
-			})
-		}
-	})
-})
-
-func dropTests(
-	context string,
-	interval uint64,
-	dropInARow uint64,
-	runDropTest func(dropCallback quicproxy.DropCallback, version protocol.VersionNumber),
-	version protocol.VersionNumber) {
-	Context(context, func() {
-		dropper := func(p uint64) bool {
+		deterministicDropper := func(p, interval, dropInARow uint64) bool {
 			if p <= 10 { // don't interfere with the crypto handshake
 				return false
 			}
 			return (p % interval) < dropInARow
 		}
 
-		It("gets a file when many outgoing packets are dropped", func() {
-			runDropTest(func(d quicproxy.Direction, p uint64) bool {
-				return d == quicproxy.DirectionOutgoing && dropper(p)
-			}, version)
-		})
+		for _, v := range protocol.SupportedVersions {
+			version := v
 
-		It("gets a file when many incoming packets are dropped", func() {
-			runDropTest(func(d quicproxy.Direction, p uint64) bool {
-				return d == quicproxy.DirectionIncoming && dropper(p)
-			}, version)
-		})
+			Context(fmt.Sprintf("with QUIC version %d", version), func() {
+				for _, d := range directions {
+					direction := d
+
+					It(fmt.Sprintf("downloads a file when every 5th packet is dropped in %s direction", d), func() {
+						runDropTest(func(d quicproxy.Direction, p uint64) bool {
+							return d.Is(direction) && deterministicDropper(p, 5, 1)
+						}, version)
+					})
+
+					It(fmt.Sprintf("downloads a file when 10 packets every 100 packet are dropped in %s direction", d), func() {
+						runDropTest(func(d quicproxy.Direction, p uint64) bool {
+							return d.Is(direction) && deterministicDropper(p, 100, 10)
+						}, version)
+					})
+				}
+			})
+		}
 	})
-}
+})
diff --git a/integrationtests/tools/proxy/proxy.go b/integrationtests/tools/proxy/proxy.go
index b8f2cf6b7..a9d0d4cf2 100644
--- a/integrationtests/tools/proxy/proxy.go
+++ b/integrationtests/tools/proxy/proxy.go
@@ -26,8 +26,30 @@ const (
 	DirectionIncoming Direction = iota
 	// DirectionOutgoing is the direction from the server to the client.
 	DirectionOutgoing
+	// DirectionBoth is both incoming and outgoing
+	DirectionBoth
 )
 
+func (d Direction) String() string {
+	switch d {
+	case DirectionIncoming:
+		return "incoming"
+	case DirectionOutgoing:
+		return "outgoing"
+	case DirectionBoth:
+		return "both"
+	default:
+		panic("unknown direction")
+	}
+}
+
+func (d Direction) Is(dir Direction) bool {
+	if d == DirectionBoth || dir == DirectionBoth {
+		return true
+	}
+	return d == dir
+}
+
 // DropCallback is a callback that determines which packet gets dropped.
 type DropCallback func(dir Direction, packetCount uint64) bool