1 files changed, 187 insertions, 2 deletions
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index 62161d5d..a4264e10 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -142,7 +142,7 @@ class TestExportPythonCode():
         assert flow_export.python_code(flow) == result
 
 
-def TestRawRequest():
+class TestRawRequest():
 
     def test_get(self):
         flow = tutils.tflow(req=req_get)
@@ -159,9 +159,10 @@ def TestRawRequest():
         flow = tutils.tflow(req=req_post)
         result = dedent("""
             POST /path HTTP/1.1\r
+            content-type: application/json\r
             host: address:22\r
             \r
-            content
+            {"name": "example", "email": "example@example.com"}
         """).strip()
         assert flow_export.raw_request(flow) == result
 
@@ -176,3 +177,187 @@ def TestRawRequest():
             content
         """).strip()
         assert flow_export.raw_request(flow) == result
+
+class TestExportLocustCode():
+
+    def test_get(self):
+        flow = tutils.tflow(req=req_get)
+        result = """
+from locust import HttpLocust, TaskSet, task
+
+class UserBehavior(TaskSet):
+    def on_start(self):
+        ''' on_start is called when a Locust start before any task is scheduled '''
+        self.flow()
+
+    @task()
+    def flow(self):
+        url = '' + self.locust.host +'/path'
+        
+        headers = {
+            'header': 'qvalue',
+            'content-length': '7',
+        }
+
+        self.response = self.client.request(
+            method='GET',
+            url=url,
+            headers=headers,
+        )
+
+class WebsiteUser(HttpLocust):
+    task_set = UserBehavior
+    min_wait=1000
+    max_wait=3000
+        """.strip()
+
+        assert flow_export.locust_code(flow) == result
+
+    def test_post(self):
+        req_post.content = '''content'''
+        req_post.headers = ''
+        flow = tutils.tflow(req=req_post)
+        result = """
+from locust import HttpLocust, TaskSet, task
+
+class UserBehavior(TaskSet):
+    def on_start(self):
+        ''' on_start is called when a Locust start before any task is scheduled '''
+        self.flow()
+
+    @task()
+    def flow(self):
+        url = '' + self.locust.host +'/path'
+        
+        data = '''content'''
+
+        self.response = self.client.request(
+            method='POST',
+            url=url,
+            data=data,
+        )
+
+class WebsiteUser(HttpLocust):
+    task_set = UserBehavior
+    min_wait=1000
+    max_wait=3000
+
+        """.strip()
+
+        assert flow_export.locust_code(flow) == result
+
+
+    def test_patch(self):
+        flow = tutils.tflow(req=req_patch)
+        result = """
+from locust import HttpLocust, TaskSet, task
+
+class UserBehavior(TaskSet):
+    def on_start(self):
+        ''' on_start is called when a Locust start before any task is scheduled '''
+        self.flow()
+
+    @task()
+    def flow(self):
+        url = '' + self.locust.host +'/path'
+        
+        headers = {
+            'header': 'qvalue',
+            'content-length': '7',
+        }
+
+        params = {
+            'query': 'param',
+        }
+
+        data = '''content'''
+
+        self.response = self.client.request(
+            method='PATCH',
+            url=url,
+            headers=headers,
+            params=params,
+            data=data,
+        )
+
+class WebsiteUser(HttpLocust):
+    task_set = UserBehavior
+    min_wait=1000
+    max_wait=3000
+
+        """.strip()
+
+        assert flow_export.locust_code(flow) == result
+
+
+class TestExportLocustTask():
+
+    def test_get(self):
+        flow = tutils.tflow(req=req_get)
+        result = '    ' + """
+    @task()
+    def path(self):
+        url = '' + self.locust.host +'/path'
+        
+        headers = {
+            'header': 'qvalue',
+            'content-length': '7',
+        }
+    
+        self.response = self.client.request(
+            method='GET',
+            url=url,
+            headers=headers,
+        )
+        """.strip()
+
+        assert flow_export.locust_task(flow) == result
+
+    def test_post(self):
+        flow = tutils.tflow(req=req_post)
+        result = '    ' + """
+    @task()
+    def path(self):
+        url = '' + self.locust.host +'/path'
+        
+        data = '''content'''
+    
+        self.response = self.client.request(
+            method='POST',
+            url=url,
+            data=data,
+        )
+    
+        """.strip()
+
+        assert flow_export.locust_task(flow) == result
+
+
+    def test_patch(self):
+        flow = tutils.tflow(req=req_patch)
+        result = '    ' + """
+    @task()
+    def path(self):
+        url = '' + self.locust.host +'/path'
+        
+        headers = {
+            'header': 'qvalue',
+            'content-length': '7',
+        }
+    
+        params = {
+            'query': 'param',
+        }
+    
+        data = '''content'''
+    
+        self.response = self.client.request(
+            method='PATCH',
+            url=url,
+            headers=headers,
+            params=params,
+            data=data,
+        )
+        """.strip()
+
+        assert flow_export.locust_task(flow) == result