py3++: test_flow_export

4 files changed, 34 insertions, 19 deletions

diff --git a/mitmproxy/flow/export.py b/mitmproxy/flow/export.py
index bffe5291..2b0f8984 100644
--- a/mitmproxy/flow/export.py
+++ b/mitmproxy/flow/export.py
@@ -4,15 +4,26 @@ import json
 import re
 from textwrap import dedent
 
+import six
 from six.moves import urllib
 
 import netlib.http
 
 
+def _native(s):
+    if six.PY2:
+        if isinstance(s, six.text_type):
+            return s.encode()
+    else:
+        if isinstance(s, six.binary_type):
+            return s.decode()
+    return s
+
+
 def dictstr(items, indent):
     lines = []
     for k, v in items:
-        lines.append(indent + "%s: %s,\n" % (repr(k), repr(v)))
+        lines.append(indent + "%s: %s,\n" % (repr(_native(k)), repr(_native(v))))
     return "{\n%s}\n" % "".join(lines)
 
 
@@ -20,7 +31,7 @@ def curl_command(flow):
     data = "curl "
 
     for k, v in flow.request.headers.fields:
-        data += "-H '%s:%s' " % (k, v)
+        data += "-H '%s:%s' " % (_native(k), _native(v))
 
     if flow.request.method != "GET":
         data += "-X %s " % flow.request.method
@@ -29,7 +40,7 @@ def curl_command(flow):
     data += "'%s'" % full_url
 
     if flow.request.content:
-        data += " --data-binary '%s'" % flow.request.content
+        data += " --data-binary '%s'" % _native(flow.request.content)
 
     return data
 
@@ -69,7 +80,7 @@ def python_code(flow):
             data = "\njson = %s\n" % dictstr(sorted(json_obj.items()), "    ")
             args += "\n    json=json,"
         else:
-            data = "\ndata = '''%s'''\n" % flow.request.body
+            data = "\ndata = '''%s'''\n" % _native(flow.request.content)
             args += "\n    data=data,"
 
     code = code.format(
@@ -85,7 +96,7 @@ def python_code(flow):
 
 def raw_request(flow):
     data = netlib.http.http1.assemble_request(flow.request)
-    return data
+    return _native(data)
 
 
 def is_json(headers, content):
@@ -127,16 +138,20 @@ def locust_code(flow):
 """).strip()
 
     components = [urllib.parse.quote(c, safe="") for c in flow.request.path_components]
-    file_name = "_".join(components)
-    name = re.sub('\W|^(?=\d)', '_', file_name)
-    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
+    name = re.sub('\W|^(?=\d)', '_', "_".join(components))
     if name == "" or name is None:
         new_name = "_".join([str(flow.request.host), str(flow.request.timestamp_start)])
         name = re.sub('\W|^(?=\d)', '_', new_name)
+
+    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
+
     args = ""
     headers = ""
     if flow.request.headers:
-        lines = [(k, v) for k, v in flow.request.headers.fields if k.lower() not in ["host", "cookie"]]
+        lines = [
+            (_native(k), _native(v)) for k, v in flow.request.headers.fields
+            if _native(k).lower() not in ["host", "cookie"]
+        ]
         lines = ["            '%s': '%s',\n" % (k, v) for k, v in lines]
         headers += "\n        headers = {\n%s        }\n" % "".join(lines)
         args += "\n            headers=headers,"
@@ -148,8 +163,8 @@ def locust_code(flow):
         args += "\n            params=params,"
 
     data = ""
-    if flow.request.body:
-        data = "\n        data = '''%s'''\n" % flow.request.body
+    if flow.request.content:
+        data = "\n        data = '''%s'''\n" % _native(flow.request.content)
         args += "\n            data=data,"
 
     code = code.format(
@@ -164,8 +179,8 @@ def locust_code(flow):
 
     host = flow.request.scheme + "://" + flow.request.host
     code = code.replace(host, "' + self.locust.host + '")
-    code = code.replace(quote_plus(host), "' + quote_plus(self.locust.host) + '")
-    code = code.replace(quote(host), "' + quote(self.locust.host) + '")
+    code = code.replace(urllib.parse.quote_plus(host), "' + quote_plus(self.locust.host) + '")
+    code = code.replace(urllib.parse.quote(host), "' + quote(self.locust.host) + '")
     code = code.replace("'' + ", "")
 
     return code
diff --git a/test/mitmproxy/data/test_flow_export/python_post_json.py b/test/mitmproxy/data/test_flow_export/python_post_json.py
index 6c1b9740..5ef110f3 100644
--- a/test/mitmproxy/data/test_flow_export/python_post_json.py
+++ b/test/mitmproxy/data/test_flow_export/python_post_json.py
@@ -8,8 +8,8 @@ headers = {
 
 
 json = {
-    u'email': u'example@example.com',
-    u'name': u'example',
+    'email': 'example@example.com',
+    'name': 'example',
 }
 
 
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index 9a263b1b..33c5137a 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -21,15 +21,15 @@ def python_equals(testdata, text):
 
 
 def req_get():
-    return netlib.tutils.treq(method='GET', content='', path=b"/path?a=foo&a=bar&b=baz")
+    return netlib.tutils.treq(method=b'GET', content=b'', path=b"/path?a=foo&a=bar&b=baz")
 
 
 def req_post():
-    return netlib.tutils.treq(method='POST', headers=())
+    return netlib.tutils.treq(method=b'POST', headers=())
 
 
 def req_patch():
-    return netlib.tutils.treq(method='PATCH', path=b"/path?query=param")
+    return netlib.tutils.treq(method=b'PATCH', path=b"/path?query=param")
 
 
 class TestExportCurlCommand():
diff --git a/tox.ini b/tox.ini
index 899fffc6..508f0066 100644
--- a/tox.ini
+++ b/tox.ini
@@ -16,7 +16,7 @@ commands =
 
 [testenv:py35]
 setenv =
-  TESTS = test/netlib test/pathod/ test/mitmproxy/script test/mitmproxy/test_contentview.py test/mitmproxy/test_custom_contentview.py test/mitmproxy/test_app.py test/mitmproxy/test_controller.py test/mitmproxy/test_fuzzing.py test/mitmproxy/test_script.py test/mitmproxy/test_web_app.py test/mitmproxy/test_utils.py test/mitmproxy/test_stateobject.py test/mitmproxy/test_cmdline.py test/mitmproxy/test_contrib_tnetstring.py test/mitmproxy/test_proxy.py test/mitmproxy/test_protocol_http1.py test/mitmproxy/test_platform_pf.py test/mitmproxy/test_server.py test/mitmproxy/test_filt.py
+  TESTS = test/netlib test/pathod/ test/mitmproxy/script test/mitmproxy/test_contentview.py test/mitmproxy/test_custom_contentview.py test/mitmproxy/test_app.py test/mitmproxy/test_controller.py test/mitmproxy/test_fuzzing.py test/mitmproxy/test_script.py test/mitmproxy/test_web_app.py test/mitmproxy/test_utils.py test/mitmproxy/test_stateobject.py test/mitmproxy/test_cmdline.py test/mitmproxy/test_contrib_tnetstring.py test/mitmproxy/test_proxy.py test/mitmproxy/test_protocol_http1.py test/mitmproxy/test_platform_pf.py test/mitmproxy/test_server.py test/mitmproxy/test_filt.py test/mitmproxy/test_flow_export.py
   HOME = {envtmpdir}
 
 [testenv:docs]