From 84733731fe96d48f03db30edb3e2c8a308f4c79a Mon Sep 17 00:00:00 2001
From: Salad Dais <SaladDais@users.noreply.github.com>
Date: Fri, 4 Jun 2021 09:05:42 +0000
Subject: [PATCH] Add distinct tests for CapsClient and ProxyCapsClient

---
 tests/{proxy => base}/test_capsclient.py | 15 +++++++--------
 tests/proxy/integration/test_http.py     | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 8 deletions(-)
 rename tests/{proxy => base}/test_capsclient.py (85%)

diff --git a/tests/proxy/test_capsclient.py b/tests/base/test_capsclient.py
similarity index 85%
rename from tests/proxy/test_capsclient.py
rename to tests/base/test_capsclient.py
index 88493d3..726c3e9 100644
--- a/tests/proxy/test_capsclient.py
+++ b/tests/base/test_capsclient.py
@@ -1,17 +1,16 @@
+import unittest
+
 import aiohttp
 import aioresponses
 from yarl import URL
 
-from hippolyzer.lib.proxy.region import ProxiedRegion
-
-from . import BaseProxyTest
+from hippolyzer.lib.base.network.caps_client import CapsClient
 
 
-class TestCapsClient(BaseProxyTest):
+class TestCapsClient(unittest.IsolatedAsyncioTestCase):
     def setUp(self) -> None:
-        super().setUp()
-        self.region = ProxiedRegion(("127.0.0.1", 1), "", self.session)
-        self.caps_client = self.region.caps_client
+        self.caps = {}
+        self.caps_client = CapsClient(self.caps)
 
     async def test_bare_url_works(self):
         with aioresponses.aioresponses() as m:
@@ -33,7 +32,7 @@ class TestCapsClient(BaseProxyTest):
                 self.assertEqual(await resp.read_llsd(), 2)
 
     async def test_caps(self):
-        self.region.update_caps({"Foobar": "https://example.com/"})
+        self.caps.update({"Foobar": "https://example.com/"})
         with aioresponses.aioresponses() as m:
             m.post("https://example.com/baz", body=b"ok")
             data = {"hi": "hello"}
diff --git a/tests/proxy/integration/test_http.py b/tests/proxy/integration/test_http.py
index 721055a..0efed0c 100644
--- a/tests/proxy/integration/test_http.py
+++ b/tests/proxy/integration/test_http.py
@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import asyncio
 
+import aioresponses
 from mitmproxy.test import tflow, tutils
 from mitmproxy.http import HTTPFlow
+from yarl import URL
 
 from hippolyzer.lib.proxy.addon_utils import BaseAddon
 from hippolyzer.lib.proxy.addons import AddonManager
@@ -70,3 +72,21 @@ class LLUDPIntegrationTests(BaseProxyTest):
         mitm_flow: HTTPFlow = HTTPFlow.from_state(flow_state)
         # The response sent back to mitmproxy should have been our modified version
         self.assertEqual(True, mitm_flow.metadata["touched_addon"])
+
+
+class TestCapsClient(BaseProxyTest):
+    def setUp(self) -> None:
+        super().setUp()
+        self._setup_default_circuit()
+        self.caps = {}
+        self.caps_client = self.session.main_region.caps_client
+
+    async def test_requests_proxied_by_default(self):
+        with aioresponses.aioresponses() as m:
+            m.get("http://example.com/", body=b"foo")
+            async with self.caps_client.get("http://example.com/") as resp:
+                self.assertEqual(await resp.read(), b"foo")
+            kwargs = m.requests[("GET", URL("http://example.com/"))][0].kwargs
+        # Request should have been proxied, with a header marking it
+        self.assertEqual(kwargs['headers']["X-Hippo-Injected"], "1")
+        self.assertEqual(kwargs['proxy'], "http://127.0.0.1:9062")