netlib.tutils -> mitmproxy.test.tutils

There's a LOT more to be done refactoring our different conflicting test utils.

1 files changed, 130 insertions, 0 deletions

diff --git a/mitmproxy/test/tutils.py b/mitmproxy/test/tutils.py
new file mode 100644
index 00000000..fc7c0eb9
--- /dev/null
+++ b/mitmproxy/test/tutils.py
@@ -0,0 +1,130 @@
+from io import BytesIO
+import tempfile
+import os
+import time
+import shutil
+from contextlib import contextmanager
+import sys
+
+from mitmproxy.utils import data
+from netlib import tcp
+from netlib import http
+
+
+def treader(bytes):
+    """
+        Construct a tcp.Read object from bytes.
+    """
+    fp = BytesIO(bytes)
+    return tcp.Reader(fp)
+
+
+@contextmanager
+def tmpdir(*args, **kwargs):
+    orig_workdir = os.getcwd()
+    temp_workdir = tempfile.mkdtemp(*args, **kwargs)
+    os.chdir(temp_workdir)
+
+    yield temp_workdir
+
+    os.chdir(orig_workdir)
+    shutil.rmtree(temp_workdir)
+
+
+def _check_exception(expected, actual, exc_tb):
+    if isinstance(expected, str):
+        if expected.lower() not in str(actual).lower():
+            raise AssertionError(
+                "Expected %s, but caught %s" % (
+                    repr(expected), repr(actual)
+                )
+            )
+    else:
+        if not isinstance(actual, expected):
+            raise AssertionError(
+                "Expected %s, but caught %s %s" % (
+                    expected.__name__, actual.__class__.__name__, repr(actual)
+                )
+            )
+
+
+def raises(expected_exception, obj=None, *args, **kwargs):
+    """
+        Assert that a callable raises a specified exception.
+
+        :exc An exception class or a string. If a class, assert that an
+        exception of this type is raised. If a string, assert that the string
+        occurs in the string representation of the exception, based on a
+        case-insenstivie match.
+
+        :obj A callable object.
+
+        :args Arguments to be passsed to the callable.
+
+        :kwargs Arguments to be passed to the callable.
+    """
+    if obj is None:
+        return RaisesContext(expected_exception)
+    else:
+        try:
+            ret = obj(*args, **kwargs)
+        except Exception as actual:
+            _check_exception(expected_exception, actual, sys.exc_info()[2])
+        else:
+            raise AssertionError("No exception raised. Return value: {}".format(ret))
+
+
+class RaisesContext:
+    def __init__(self, expected_exception):
+        self.expected_exception = expected_exception
+
+    def __enter__(self):
+        return
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not exc_type:
+            raise AssertionError("No exception raised.")
+        else:
+            _check_exception(self.expected_exception, exc_val, exc_tb)
+        return True
+
+
+test_data = data.Data(__name__).push("../../test/netlib")
+
+
+def treq(**kwargs):
+    """
+    Returns:
+        netlib.http.Request
+    """
+    default = dict(
+        first_line_format="relative",
+        method=b"GET",
+        scheme=b"http",
+        host=b"address",
+        port=22,
+        path=b"/path",
+        http_version=b"HTTP/1.1",
+        headers=http.Headers(((b"header", b"qvalue"), (b"content-length", b"7"))),
+        content=b"content"
+    )
+    default.update(kwargs)
+    return http.Request(**default)
+
+
+def tresp(**kwargs):
+    """
+    Returns:
+        netlib.http.Response
+    """
+    default = dict(
+        http_version=b"HTTP/1.1",
+        status_code=200,
+        reason=b"OK",
+        headers=http.Headers(((b"header-response", b"svalue"), (b"content-length", b"7"))),
+        content=b"message",
+        timestamp_start=time.time(),
+        timestamp_end=time.time(),
+    )
+    default.update(kwargs)
+    return http.Response(**default)