Fixed failing test.

Yubico · Mar 11, 2015 · 87eff1f · 87eff1f
1 parent 9c0c293
commit 87eff1f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git 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
@@ -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]