From 327e933faf262c3b5c124fd576d243d451611252 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 27 Aug 2020 11:37:35 +0200
Subject: [PATCH] intercept: tests++

---
 mitmproxy/addons/intercept.py           |  3 ++-
 test/mitmproxy/addons/test_intercept.py | 25 ++++++++++++++++++++++---
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/mitmproxy/addons/intercept.py b/mitmproxy/addons/intercept.py
index e36eff260..5eebd6dc0 100644
--- a/mitmproxy/addons/intercept.py
+++ b/mitmproxy/addons/intercept.py
@@ -30,7 +30,7 @@ class Intercept:
                 ctx.options.intercept_active = False
 
     def should_intercept(self, f: flow.Flow) -> bool:
-        return (
+        return bool(
                 ctx.options.intercept_active
                 and self.filt
                 and self.filt(f)
@@ -39,6 +39,7 @@ class Intercept:
 
     def process_flow(self, f: flow.Flow) -> None:
         if self.should_intercept(f):
+            assert f.reply
             if f.reply.state != "start":
                 return ctx.log.debug("Cannot intercept request that is already taken by another addon.")
             f.intercept()
diff --git a/test/mitmproxy/addons/test_intercept.py b/test/mitmproxy/addons/test_intercept.py
index 1d10fbc90..c3252db4a 100644
--- a/test/mitmproxy/addons/test_intercept.py
+++ b/test/mitmproxy/addons/test_intercept.py
@@ -43,12 +43,31 @@ def test_simple():
         tctx.cycle(r, f)
         assert f.intercepted
 
+
+def test_tcp():
+    r = intercept.Intercept()
+    with taddons.context(r) as tctx:
+        tctx.configure(r, intercept="~tcp")
+        f = tflow.ttcpflow()
+        tctx.cycle(r, f)
+        assert f.intercepted
+
         tctx.configure(r, intercept_active=False)
         f = tflow.ttcpflow()
         tctx.cycle(r, f)
         assert not f.intercepted
 
-        tctx.configure(r, intercept_active=True)
-        f = tflow.ttcpflow()
-        tctx.cycle(r, f)
+
+def test_already_taken():
+    r = intercept.Intercept()
+    with taddons.context(r) as tctx:
+        tctx.configure(r, intercept="~q")
+
+        f = tflow.tflow()
+        tctx.invoke(r, "request", f)
         assert f.intercepted
+
+        f = tflow.tflow()
+        f.reply.take()
+        tctx.invoke(r, "request", f)
+        assert not f.intercepted