Upstream proxy auth to addon

3 files changed, 54 insertions, 25 deletions

diff --git a/test/mitmproxy/addons/test_upstream_proxy_auth.py b/test/mitmproxy/addons/test_upstream_proxy_auth.py
new file mode 100644
index 00000000..e9a7f4ef
--- /dev/null
+++ b/test/mitmproxy/addons/test_upstream_proxy_auth.py
@@ -0,0 +1,54 @@
+import base64
+
+from mitmproxy import exceptions
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
+from mitmproxy.addons import upstream_proxy_auth
+
+
+def test_configure():
+    up = upstream_proxy_auth.UpstreamProxyAuth()
+    with taddons.context() as tctx:
+        tctx.configure(up, upstream_auth="test:test")
+        assert up.auth == b"Basic" + b" " + base64.b64encode(b"test:test")
+
+        tctx.configure(up, upstream_auth="test:")
+        assert up.auth == b"Basic" + b" " + base64.b64encode(b"test:")
+
+        tctx.configure(up, upstream_auth=None)
+        assert not up.auth
+
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=""
+        )
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=":"
+        )
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=":test"
+        )
+
+
+def test_simple():
+    up = upstream_proxy_auth.UpstreamProxyAuth()
+    with taddons.context() as tctx:
+        tctx.configure(up, upstream_auth="foo:bar")
+
+        f = tflow.tflow()
+        f.mode = "upstream"
+        up.requestheaders(f)
+        assert "proxy-authorization" in f.request.headers
+
+        f = tflow.tflow()
+        up.requestheaders(f)
+        assert "proxy-authorization" not in f.request.headers
diff --git a/test/mitmproxy/test_proxy.py b/test/mitmproxy/test_proxy.py
index 7cadb6c2..8847c088 100644
--- a/test/mitmproxy/test_proxy.py
+++ b/test/mitmproxy/test_proxy.py
@@ -107,14 +107,10 @@ class TestProcessProxyOptions:
         self.assert_noerr("-T")
 
         self.assert_noerr("-U", "http://localhost")
-        self.assert_err("expected one argument", "-U")
         self.assert_err("Invalid server specification", "-U", "upstream")
 
         self.assert_noerr("--upstream-auth", "test:test")
         self.assert_err("expected one argument", "--upstream-auth")
-        self.assert_err(
-            "Invalid upstream auth specification", "--upstream-auth", "test"
-        )
         self.assert_err("mutually exclusive", "-R", "http://localhost", "-T")
 
     def test_socks_auth(self):
diff --git a/test/mitmproxy/test_proxy_config.py b/test/mitmproxy/test_proxy_config.py
index e012cb5e..e2c39846 100644
--- a/test/mitmproxy/test_proxy_config.py
+++ b/test/mitmproxy/test_proxy_config.py
@@ -1,5 +1,4 @@
 from mitmproxy.test import tutils
-import base64
 from mitmproxy.proxy import config
 
 
@@ -26,23 +25,3 @@ def test_parse_server_spec():
         config.parse_server_spec,
         "http://"
     )
-
-
-def test_parse_upstream_auth():
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ""
-    )
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ":"
-    )
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ":test"
-    )
-    assert config.parse_upstream_auth("test:test") == b"Basic" + b" " + base64.b64encode(b"test:test")
-    assert config.parse_upstream_auth("test:") == b"Basic" + b" " + base64.b64encode(b"test:")