Add support for specifying cookies in request (#274)

docs/source/http.md

@@ -109,6 +109,75 @@ stages:
 This test ensures that a cookie called `session-cookie` is returned from the
 'login' stage, and this cookie will be sent with all future stages of that test.
 
+#### Choosing cookies
+
+If you have multiple cookies for a domain, the `cookies` key
+can also be used in the request block to specify which one to send:
+
+```yaml
+---
+
+test_name: Test receiving and sending cookie
+
+includes:
+  - !include common.yaml
+
+stages:
+  - name: Expect multiple cookies returned
+    request:
+      url: "{host}/get_cookie"
+      method: POST
+    response:
+      status_code: 200
+      cookies:
+        - tavern-cookie-1
+        - tavern-cookie-2
+
+  - name: Only send one cookie
+    request:
+      url: "{host}/expect_cookie"
+      method: GET
+      cookies:
+        - tavern-cookie-1
+    response:
+      status_code: 200
+      body:
+        status: ok
+```
+
+Trying to specify a cookie which does not exist will fail the stage.
+
+To send _no_ cookies, simply use an empty array:
+
+```yaml
+---
+
+test_name: Test receiving and sending cookie
+
+includes:
+  - !include common.yaml
+
+stages:
+  - name: get cookie for domain
+    request:
+      url: "{host}/get_cookie"
+      method: POST
+    response:
+      status_code: 200
+      cookies:
+        - tavern-cookie-1
+
+  - name: Send no cookies
+    request:
+      url: "{host}/expect_cookie"
+      method: GET
+      cookies: []
+    response:
+      status_code: 403
+      body:
+        status: access denied
+```
+
 ### HTTP Basic Auth
 
 For a server that expects HTTP Basic Auth, the `auth` keyword can be used in the

tavern/_plugins/rest/request.py

@@ -1,6 +1,7 @@
+import contextlib
 import json
-import mimetypes
 import logging
+import mimetypes
 import os
 import warnings
 
@@ -12,6 +13,8 @@
 from contextlib2 import ExitStack
 from future.utils import raise_from
 import requests
+from requests.cookies import cookiejar_from_dict
+from requests.utils import dict_from_cookiejar
 from box import Box
 
 from tavern.util import exceptions
@@ -160,11 +163,34 @@ def add_request_args(keys, optional):
     return request_args
 
 
+@contextlib.contextmanager
+def _set_cookies_for_request(session, request_args):
+    """
+    Possibly reset session cookies for a single request then set them back.
+    If no cookies were present in the request arguments, do nothing.
+
+    This does not use try/finally because if it fails then we don't care about
+    the cookies anyway
+
+    Args:
+        session (requests.Session): Current session
+        request_args (dict): current request arguments
+    """
+    if "cookies" in request_args:
+        old_cookies = dict_from_cookiejar(session.cookies)
+        session.cookies = cookiejar_from_dict({})
+        yield
+        session.cookies = cookiejar_from_dict(old_cookies)
+    else:
+        yield
+
+
 class RestRequest(BaseRequest):
     def __init__(self, session, rspec, test_block_config):
         """Prepare request
 
         Args:
+            session (requests.Session): existing session
             rspec (dict): test spec
             test_block_config (dict): Any configuration for this the block of
                 tests
@@ -191,15 +217,31 @@ def __init__(self, session, rspec, test_block_config):
             "files",
             "stream",
             "timeout",
+            "cookies",
             "cert",
-            # "cookies",
             # "hooks",
         }
 
         check_expected_keys(expected, rspec)
 
         request_args = get_request_args(rspec, test_block_config)
 
+        # Need to do this down here - it is separate from getting request args as
+        # it depends on the state of the session
+        if "cookies" in rspec:
+            existing_cookies = session.cookies.get_dict()
+            missing = set(rspec["cookies"]) - set(existing_cookies.keys())
+            if missing:
+                logger.error("Missing cookies")
+                raise exceptions.MissingCookieError(
+                    "Tried to use cookies '{}' in request but only had '{}' available".format(
+                        rspec["cookies"], existing_cookies
+                    )
+                )
+            request_args["cookies"] = {
+                c: existing_cookies.get(c) for c in rspec["cookies"]
+            }
+
         logger.debug("Request args: %s", request_args)
 
         request_args.update(allow_redirects=False)
@@ -215,6 +257,7 @@ def prepared_request():
             # If there are open files, create a context manager around each so
             # they will be closed at the end of the request.
             with ExitStack() as stack:
+                stack.enter_context(_set_cookies_for_request(session, request_args))
                 self._request_args.update(self._get_file_arguments(stack))
                 return session.request(**self._request_args)
 

tavern/schemas/tests.schema.yaml

@@ -134,9 +134,12 @@ schema;stage:
           required: false
           sequence:
             - type: str
-              required: true
+
+        cookies:
+          type: seq
+          required: false
+          sequence:
             - type: str
-              required: true
 
         json:
           include: any_map_or_list_with_ext_function

tests/integration/server.py

@@ -1,3 +1,4 @@
+import base64
 import itertools
 import math
 import mimetypes
@@ -6,7 +7,6 @@
 
 from flask import Flask, request, jsonify, Response
 
-
 app = Flask(__name__)
 
 
@@ -217,3 +217,23 @@ def iter():
 def poll():
     response = {'status': next(statuses)}
     return jsonify(response)
+
+
+def _maybe_get_cookie_name():
+    return (request.get_json() or {}).get("cookie_name", "tavern-cookie")
+
+@app.route("/get_cookie", methods=["POST"])
+def give_cookie():
+    cookie_name = _maybe_get_cookie_name()
+    response = Response()
+    response.set_cookie(cookie_name, base64.b64encode(os.urandom(16)).decode("utf8"))
+    return response, 200
+
+
+@app.route("/expect_cookie", methods=["GET"])
+def expect_cookie():
+    cookie_name = _maybe_get_cookie_name()
+    if cookie_name not in request.cookies:
+        return jsonify({"error": "No cookie named {} in request".format(cookie_name)}), 400
+    else:
+        return jsonify({"status": "ok"}), 200