1 files changed, 193 insertions, 0 deletions
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index d98759bb..75a8090f 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -179,6 +179,199 @@ class TestRawRequest():
         """).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.path()
+
+    @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,
+        )
+
+    ### Additional tasks can go here ###
+
+
+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.path()
+
+    @task()
+    def path(self):
+        url = self.locust.host + '/path'
+        
+        data = '''content'''
+
+        self.response = self.client.request(
+            method='POST',
+            url=url,
+            data=data,
+        )
+
+    ### Additional tasks can go here ###
+
+
+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.path()
+
+    @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,
+        )
+
+    ### Additional tasks can go here ###
+
+
+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() + '\n'
+
+        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() + '\n'
+
+        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() + '\n'
+
+        assert flow_export.locust_task(flow) == result
+
+
 class TestIsJson():
 
     def test_empty(self):