From 87eff1fa5358c79b4c91fbba68e069bebd92f326 Mon Sep 17 00:00:00 2001
From: Dain Nilsson <dain@yubico.com>
Date: Wed, 11 Mar 2015 14:38:35 +0100
Subject: [PATCH] Fixed failing test.

---
 test/test_client_api.py | 9 ++++++---
 yubiauth/client/rest.py | 9 ++++++++-
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/test/test_client_api.py b/test/test_client_api.py
index 0f5a1aa..2fa3457 100644
--- a/test/test_client_api.py
+++ b/test/test_client_api.py
@@ -37,16 +37,19 @@ def test_login_logout():
 
 
 def test_header_authentication():
-    app.post(
+    json_session_id = app.post(
         '/login',
         {'username': 'user1', 'password': 'pass1'}
-    )
+    ).json.get('session')
 
     session_id = app.cookies[SESSION_COOKIE]
-    #Clear the cookie
+    # Clear the cookie
     app.reset()
     app.get('/status', status=400)
 
+    # Make sure cookie matches json body
+    assert session_id == json_session_id
+
     headers = {SESSION_HEADER: session_id}
 
     username = app.get('/status', headers=headers).json['username']
diff --git a/yubiauth/client/rest.py b/yubiauth/client/rest.py
index 301343c..6b87ab1 100644
--- a/yubiauth/client/rest.py
+++ b/yubiauth/client/rest.py
@@ -112,7 +112,14 @@ def new_func(self, request, *args, **kwargs):
 
 
 def get_session_cookie(request):
-    cookie = request.environ['beaker.session']._headers['cookie_out']
+    log.info('beaker.session: %r', request.environ['beaker.session']._headers)
+    headers = request.environ['beaker.session']._headers
+    if SESSION_COOKIE in (headers.get('cookie_set') or ''):
+        cookie = headers['cookie_set']
+    elif SESSION_COOKIE in (headers.get('cookie') or ''):
+        cookie = headers['cookie']
+    else:
+        return None
     return cookie.split('=', 1)[1].split(';', 1)[0]