From e58e87befe0c7bcd195321143a740f331bcd4afe Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Thu, 11 Apr 2019 12:25:20 +0800
Subject: [PATCH] refactor upload files with requests-toolbelt:

1, Simplify usage syntax;
2, support upload multiple fields.
---
 Pipfile                  |  1 +
 httprunner/__about__.py  |  2 +-
 httprunner/built_in.py   | 58 +++++++++++++++++++++++++++++++++-------
 httprunner/loader.py     |  5 ++--
 setup.py                 |  3 ++-
 tests/httpbin/upload.yml | 10 +++----
 6 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/Pipfile b/Pipfile
index 83a9df29..5e4c9ad5 100644
--- a/Pipfile
+++ b/Pipfile
@@ -11,6 +11,7 @@ har2case = "*"
 colorama = "*"
 colorlog = "*"
 requests-toolbelt = "*"
+filetype = "*"
 
 [dev-packages]
 Flask = "<1.0.0"
diff --git a/httprunner/__about__.py b/httprunner/__about__.py
index 01d4dc4f..cc08e19b 100644
--- a/httprunner/__about__.py
+++ b/httprunner/__about__.py
@@ -1,7 +1,7 @@
 __title__ = 'HttpRunner'
 __description__ = 'One-stop solution for HTTP(S) testing.'
 __url__ = 'https://github.com/HttpRunner/HttpRunner'
-__version__ = '2.1.0'
+__version__ = '2.1.1'
 __author__ = 'debugtalk'
 __author_email__ = 'mail@debugtalk.com'
 __license__ = 'Apache-2.0'
diff --git a/httprunner/built_in.py b/httprunner/built_in.py
index e3e0d1bd..86825bd0 100644
--- a/httprunner/built_in.py
+++ b/httprunner/built_in.py
@@ -12,10 +12,13 @@ import re
 import string
 import time
 
+import filetype
 from httprunner.compat import basestring, builtin_str, integer_types, str
 from httprunner.exceptions import ParamsError
 from requests_toolbelt import MultipartEncoder
 
+PWD = os.getcwd()
+
 
 """ built-in functions
 """
@@ -38,17 +41,54 @@ def get_current_date(fmt="%Y-%m-%d"):
     """
     return datetime.datetime.now().strftime(fmt)
 
-def multipart_encoder(field_name, file_path, file_type=None, file_headers=None):
-    if not os.path.isabs(file_path):
-        file_path = os.path.join(os.getcwd(), file_path)
 
-    filename = os.path.basename(file_path)
-    with open(file_path, 'rb') as f:
-        fields = {
-            field_name: (filename, f.read(), file_type)
-        }
+def multipart_encoder(**kwargs):
+    """ upload files with requests-toolbelt
+
+        - test:
+            name: upload file
+            variables:
+                file_path: "data/test.env"
+                multipart_encoder: ${multipart_encoder(file=$file_path)}
+            request:
+                url: /post
+                method: POST
+                headers:
+                    Content-Type: ${multipart_content_type($multipart_encoder)}
+                data: $multipart_encoder
+            validate:
+                - eq: ["status_code", 200]
+                - startswith: ["content.files.file", "UserName=test"]
+
+    """
+    def get_filetype(file_path):
+        file_type = filetype.guess(file_path)
+        if file_type:
+            return file_type.mime
+        else:
+            return "text/html"
+
+    fields_dict = {}
+    for key, value in kwargs.items():
+
+        if os.path.isabs(value):
+            _file_path = value
+            is_file = True
+        else:
+            global PWD
+            _file_path = os.path.join(PWD, value)
+            is_file = os.path.isfile(_file_path)
+
+        if is_file:
+            filename = os.path.basename(_file_path)
+            with open(_file_path, 'rb') as f:
+                mime_type = get_filetype(_file_path)
+                fields_dict[key] = (filename, f.read(), mime_type)
+        else:
+            fields_dict[key] = value
+
+    return MultipartEncoder(fields=fields_dict)
 
-    return MultipartEncoder(fields)
 
 def multipart_content_type(multipart_encoder):
     return multipart_encoder.content_type
diff --git a/httprunner/loader.py b/httprunner/loader.py
index ade04d6a..3fddb238 100644
--- a/httprunner/loader.py
+++ b/httprunner/loader.py
@@ -8,8 +8,7 @@ import os
 import sys
 
 import yaml
-from httprunner import exceptions, logger, parser, utils, validator
-
+from httprunner import built_in, exceptions, logger, parser, utils, validator
 
 ###############################################################################
 ##   file loader
@@ -263,7 +262,6 @@ def load_module_functions(module):
 def load_builtin_functions():
     """ load built_in module functions
     """
-    from httprunner import built_in
     return load_module_functions(built_in)
 
 
@@ -703,6 +701,7 @@ def load_project_tests(test_path, dot_env_path=None):
     # locate PWD and load debugtalk.py functions
 
     project_mapping["PWD"] = project_working_directory
+    built_in.PWD = project_working_directory
     project_mapping["functions"] = debugtalk_functions
 
     # load api
diff --git a/setup.py b/setup.py
index 3c8effee..9e118d2b 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,8 @@ install_requires = [
     "har2case",
     "colorama",
     "colorlog",
-    "requests_toolbelt"
+    "requests_toolbelt",
+    "filetype"
 ]
 
 class UploadCommand(Command):
diff --git a/tests/httpbin/upload.yml b/tests/httpbin/upload.yml
index 3abd8c76..344b90bd 100644
--- a/tests/httpbin/upload.yml
+++ b/tests/httpbin/upload.yml
@@ -5,17 +5,15 @@
 - test:
     name: upload file
     variables:
-        field_name: "file"
-        file_path: "LICENSE"
-        file_type: "text/html"
-        multipart_encoder: ${multipart_encoder($field_name, $file_path, $file_type)}
+        file_path: "data/test.env"
+        multipart_encoder: ${multipart_encoder(file=$file_path)}
     request:
         url: /post
         method: POST
         headers:
             Content-Type: ${multipart_content_type($multipart_encoder)}
         data: $multipart_encoder
-    validators:
+    validate:
         - eq: ["status_code", 200]
-        - startswith: ["content.files.file", "MIT License"]
+        - startswith: ["content.files.file", "UserName=test"]