Allow multi part file uploads (#870)

constraints.txt

@@ -269,6 +269,12 @@ twine==4.0.2
     # via tavern (pyproject.toml)
 types-pyyaml==6.0.12.2
     # via tavern (pyproject.toml)
+types-requests==2.31.0.0
+    # via tavern (pyproject.toml)
+types-setuptools==67.8.0.0
+    # via tavern (pyproject.toml)
+types-urllib3==1.26.25.13
+    # via types-requests
 typing-extensions==4.4.0
     # via mypy
 urllib3==1.26.13

pyproject.toml

@@ -80,6 +80,8 @@ dev = [
     "tox-travis",
     "twine",
     "wheel",
+    "types-setuptools",
+    "types-requests",
     # This has to be installed separately, otherwise you can't upload to pypi
     # "tbump@https://github.com/michaelboulton/tbump/archive/714ba8957a3c84b625608ceca39811ebe56229dc.zip",
 ]
@@ -132,6 +134,7 @@ addopts = [
     "--strict-markers",
     "-p", "no:logging",
     "--tb=short",
+    "--color=yes"
 ]
 norecursedirs = [
     ".git",

requirements.txt

@@ -856,6 +856,18 @@ types-pyyaml==6.0.12.2 \
     --hash=sha256:1e94e80aafee07a7e798addb2a320e32956a373f376655128ae20637adb2655b \
     --hash=sha256:6840819871c92deebe6a2067fb800c11b8a063632eb4e3e755914e7ab3604e83
     # via tavern (pyproject.toml)
+types-requests==2.31.0.0 \
+    --hash=sha256:7c5cea7940f8e92ec560bbc468f65bf684aa3dcf0554a6f8c4710f5f708dc598 \
+    --hash=sha256:c1c29d20ab8d84dff468d7febfe8e0cb0b4664543221b386605e14672b44ea25
+    # via tavern (pyproject.toml)
+types-setuptools==67.8.0.0 \
+    --hash=sha256:6df73340d96b238a4188b7b7668814b37e8018168aef1eef94a3b1872e3f60ff \
+    --hash=sha256:95c9ed61871d6c0e258433373a4e1753c0a7c3627a46f4d4058c7b5a08ab844f
+    # via tavern (pyproject.toml)
+types-urllib3==1.26.25.13 \
+    --hash=sha256:3300538c9dc11dad32eae4827ac313f5d986b8b21494801f1bf97a1ac6c03ae5 \
+    --hash=sha256:5dbd1d2bef14efee43f5318b5d36d805a489f6600252bb53626d4bfafd95e27c
+    # via types-requests
 typing-extensions==4.4.0 \
     --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \
     --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e

scripts/smoke.bash

@@ -2,7 +2,8 @@
 
 set -ex
 
-ruff tavern
+ruff tavern tests --fix
+black tavern tests
 
 # Separate as isort can interfere with other testenvs
 tox --parallel -c tox.ini        \

tavern/_core/schema/extensions.py

@@ -368,18 +368,25 @@ def validate_cert_tuple_or_str(value, rule_obj, path) -> bool:
 def validate_file_spec(value, rule_obj, path) -> bool:
     """Validate file upload arguments"""
 
+    logger = get_pykwalify_logger("tavern.schema.extensions")
+
     if not isinstance(value, dict):
         raise BadSchemaError(
             "File specification must be a mapping of file names to file specs, got {}".format(
                 value
             )
         )
 
+    if value.get("file_path"):
+        # If the file spec was a list, this function will be called for each item. Just call this
+        # function recursively to check each item.
+        return validate_file_spec({"file": value}, rule_obj, path)
+
     for _, filespec in value.items():
         if isinstance(filespec, str):
             file_path = filespec
         elif isinstance(filespec, dict):
-            valid = {"file_path", "content_type", "content_encoding"}
+            valid = {"file_path", "content_type", "content_encoding", "form_field_name"}
             extra = set(filespec.keys()) - valid
             if extra:
                 raise BadSchemaError(
@@ -399,7 +406,7 @@ def validate_file_spec(value, rule_obj, path) -> bool:
 
         if not os.path.exists(file_path):
             if re.search(".*{.+}.*", file_path):
-                get_pykwalify_logger("tavern.schemas.extensions").debug(
+                logger.debug(
                     "Could not find file path, but it might be a format variable, so continuing"
                 )
             else:

tavern/_core/schema/tests.jsonschema.yaml

@@ -162,7 +162,9 @@ definitions:
         description: Path to a file to upload as the request body
 
       files:
-        type: object
+        oneOf:
+          - type: object
+          - type: array
         description: Files to send as part of the request
 
       clear_session_cookies:

tavern/_plugins/rest/files.py

@@ -0,0 +1,189 @@
+import dataclasses
+import logging
+import mimetypes
+import os
+from contextlib import ExitStack
+from typing import Any, List, Optional, Tuple, Union
+
+from tavern._core import exceptions
+from tavern._core.dict_util import format_keys
+from tavern._core.pytest.config import TestConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class _Filespec:
+    """A description of a file for a file upload, possibly as part of a multi part upload"""
+
+    path: str
+    content_type: Optional[str] = None
+    content_encoding: Optional[str] = None
+    form_field_name: Optional[str] = None
+
+
+def _parse_filespec(filespec: Union[str, dict]) -> _Filespec:
+    """
+    Get configuration for uploading file
+
+    Args:
+        filespec: Can either be one of
+            - A path to a file
+            - A dict containing 'long' format, possibly including content type/encoding and the
+                multipart 'name'
+
+    Returns:
+        The parsed file spec
+
+    Raises:
+        exceptions.BadSchemaError: If the file spec was invalid
+    """
+    if isinstance(filespec, str):
+        return _Filespec(filespec)
+    elif isinstance(filespec, dict):
+        try:
+            # The one required key
+            path = filespec["file_path"]
+        except KeyError as e:
+            raise exceptions.BadSchemaError(
+                "File spec dict did not contain the required 'file_path' key"
+            ) from e
+
+        return _Filespec(
+            path,
+            filespec.get("content_type"),
+            filespec.get("content_encoding"),
+            filespec.get("form_field_name"),
+        )
+    else:
+        # Could remove, also done in schema check
+        raise exceptions.BadSchemaError(
+            "File specification must be a path or a dictionary"
+        )
+
+
+def guess_filespec(
+    filespec: Union[str, dict], stack: ExitStack, test_block_config: TestConfig
+) -> Tuple[List, Optional[str]]:
+    """tries to guess the content type and encoding from a file.
+
+    Args:
+        test_block_config: config for test/stage
+        stack: exit stack to add open files context to
+        filespec: a string path to a file or a dictionary of the file path, content type, and encoding.
+
+    Returns:
+        A tuple of either length 2 (filename and file object), 3 (as before, with content type),
+            or 4 (as before, with with content encoding). If a group name for the multipart upload
+            was specified, this is also returned.
+
+    Notes:
+        If a 4-tuple is returned, the last element is a dictionary of headers to send to requests,
+            _not_ the raw encoding value.
+    """
+    if not mimetypes.inited:
+        mimetypes.init()
+
+    parsed = _parse_filespec(filespec)
+
+    filepath = format_keys(parsed.path, test_block_config.variables)
+    filename = os.path.basename(filepath)
+
+    # a 2-tuple ('filename', fileobj)
+    file_spec = [
+        filename,
+        stack.enter_context(open(filepath, "rb")),
+    ]
+
+    # Try to guess as well, but don't override what the user specified
+    guessed_content_type, guessed_encoding = mimetypes.guess_type(filepath)
+    content_type = parsed.content_type or guessed_content_type
+    encoding = parsed.content_encoding or guessed_encoding
+
+    # If it doesn't have a mimetype, or can't guess it, don't
+    # send the content type for the file
+    if content_type:
+        # a 3-tuple ('filename', fileobj, 'content_type')
+        logger.debug("content_type for '%s' = '%s'", filename, content_type)
+        file_spec.append(content_type)
+        if encoding:
+            # or a 4-tuple ('filename', fileobj, 'content_type', custom_headers)
+            logger.debug("encoding for '%s' = '%s'", filename, encoding)
+            # encoding is None for no encoding or the name of the
+            # program used to encode (e.g. compress or gzip). The
+            # encoding is suitable for use as a Content-Encoding header.
+            file_spec.append({"Content-Encoding": encoding})
+
+    return file_spec, parsed.form_field_name
+
+
+def _parse_file_mapping(file_args, stack, test_block_config) -> dict:
+    """Parses a simple mapping of uploads where each key is mapped to one form field name which has one file"""
+    files_to_send = {}
+    for key, filespec in file_args.items():
+        file_spec, form_field_name = guess_filespec(filespec, stack, test_block_config)
+
+        # If it's a dict then the key is used as the name, at least to maintain backwards compatability
+        if form_field_name:
+            logger.warning(
+                f"Specified 'form_field_name' as '{form_field_name}' in file spec, but the file name was inferred to be '{key}' from the mapping - the form_field_name will be ignored"
+            )
+
+        files_to_send[key] = tuple(file_spec)
+    return files_to_send
+
+
+def _parse_file_list(file_args, stack, test_block_config) -> List:
+    """Parses a case where there may be multiple files uploaded as part of one form field"""
+    files_to_send: List[Any] = []
+    for filespec in file_args:
+        file_spec, form_field_name = guess_filespec(filespec, stack, test_block_config)
+
+        if not form_field_name:
+            raise exceptions.BadSchemaError(
+                "If specifying a list of files to upload for a multi part upload, the 'form_field_name' key must also be specified for each file to upload"
+            )
+
+        files_to_send.append(
+            (
+                form_field_name,
+                tuple(file_spec),
+            )
+        )
+
+    return files_to_send
+
+
+def get_file_arguments(
+    request_args: dict, stack: ExitStack, test_block_config: TestConfig
+) -> dict:
+    """Get correct arguments for anything that should be passed as a file to
+    requests
+
+    Args:
+        request_args: args passed to requests
+        test_block_config: config for test
+        stack: context stack to add file objects to so they're
+            closed correctly after use
+
+    Returns:
+        mapping of 'files' block to pass directly to requests
+    """
+
+    files_to_send: Optional[Union[dict, List]] = None
+
+    file_args = request_args.get("files")
+
+    if isinstance(file_args, dict):
+        files_to_send = _parse_file_mapping(file_args, stack, test_block_config)
+    elif isinstance(file_args, list):
+        files_to_send = _parse_file_list(file_args, stack, test_block_config)
+    elif file_args is not None:
+        raise exceptions.BadSchemaError(
+            f"'files' key in a HTTP request can only be a dict or a list but was {type(file_args)}"
+        )
+
+    if files_to_send:
+        return {"files": files_to_send}
+    else:
+        return {}