Combine cookies from original request and session file

httpie/sessions.py

@@ -4,6 +4,8 @@
 """
 import os
 import re
+
+from http.cookies import SimpleCookie
 from pathlib import Path
 from typing import Iterable, Optional, Union
 from urllib.parse import urlsplit
@@ -76,7 +78,13 @@ def update_headers(self, request_headers: RequestHeadersDict):
                 continue  # Ignore explicitly unset headers
 
             value = value.decode('utf8')
-            if name == 'User-Agent' and value.startswith('HTTPie/'):
+            if name.lower() == 'user-agent' and value.startswith('HTTPie/'):
+                continue
+
+            if name.lower() == 'cookie':
+                for cookie_name, morsel in SimpleCookie(value).items():
+                    self['cookies'][cookie_name] = {'value': morsel.value}
+                del request_headers['Cookie']
                 continue
 
             for prefix in SESSION_IGNORED_HEADER_PREFIXES:

tests/test_sessions.py

@@ -35,6 +35,27 @@ def env(self):
         return MockEnvironment(config_dir=self.config_dir)
 
 
+class CookieTestBase:
+    def setup_method(self, method):
+        self.config_dir = mk_config_dir()
+
+        orig_session = {
+            'cookies': {
+                'cookie1': {
+                    'value': 'foo',
+                },
+                'cookie2': {
+                    'value': 'foo',
+                }
+            }
+        }
+        self.session_path = self.config_dir / 'test-session.json'
+        self.session_path.write_text(json.dumps(orig_session))
+
+    def teardown_method(self, method):
+        shutil.rmtree(self.config_dir)
+
+
 class TestSessionFlow(SessionTestBase):
     """
     These tests start with an existing session created in `setup_method()`.
@@ -191,13 +212,7 @@ def test_download_in_session(self, httpbin):
             os.chdir(cwd)
 
 
-class TestExpiredCookies:
-
-    def setup_method(self, method):
-        self.config_dir = mk_config_dir()
-
-    def teardown_method(self, method):
-        shutil.rmtree(self.config_dir)
+class TestExpiredCookies(CookieTestBase):
 
     @pytest.mark.parametrize(
         argnames=['initial_cookie', 'expired_cookie'],
@@ -213,29 +228,16 @@ def test_removes_expired_cookies_from_session_obj(self, initial_cookie, expired_
         assert expired_cookie not in session.cookies
 
     def test_expired_cookies(self, httpbin):
-        orig_session = {
-            'cookies': {
-                'to_expire': {
-                    'value': 'foo'
-                },
-                'to_stay': {
-                    'value': 'foo'
-                },
-            }
-        }
-        session_path = self.config_dir / 'test-session.json'
-        session_path.write_text(json.dumps(orig_session))
-
         r = http(
-            '--session', str(session_path),
+            '--session', str(self.session_path),
             '--print=H',
-            httpbin.url + '/cookies/delete?to_expire',
+            httpbin.url + '/cookies/delete?cookie2',
         )
-        assert 'Cookie: to_expire=foo; to_stay=foo' in r
+        assert 'Cookie: cookie1=foo; cookie2=foo' in r
 
-        updated_session = json.loads(session_path.read_text())
-        assert 'to_stay' in updated_session['cookies']
-        assert 'to_expire' not in updated_session['cookies']
+        updated_session = json.loads(self.session_path.read_text())
+        assert 'cookie1' in updated_session['cookies']
+        assert 'cookie2' not in updated_session['cookies']
 
     @pytest.mark.parametrize(
         argnames=['headers', 'now', 'expected_expired'],
@@ -277,3 +279,90 @@ def test_expired_cookies(self, httpbin):
     )
     def test_get_expired_cookies_manages_multiple_cookie_headers(self, headers, now, expected_expired):
         assert get_expired_cookies(headers, now=now) == expected_expired
+
+
+class TestCookieStorage(CookieTestBase):
+
+    @pytest.mark.parametrize(
+        argnames=['new_cookies', 'new_cookies_dict', 'expected'],
+        argvalues=[(
+            'new=bar',
+            {'new': 'bar'},
+            'cookie1=foo; cookie2=foo; new=bar'
+        ),
+            (
+            'new=bar;chocolate=milk',
+            {'new': 'bar', 'chocolate': 'milk'},
+            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
+        ),
+            (
+            'new=bar; chocolate=milk',
+            {'new': 'bar', 'chocolate': 'milk'},
+            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
+        ),
+            (
+            'new=bar;; chocolate=milk;;;',
+            {'new': 'bar', 'chocolate': 'milk'},
+            'cookie1=foo; cookie2=foo; new=bar'
+        ),
+            (
+            'new=bar; chocolate=milk;;;',
+            {'new': 'bar', 'chocolate': 'milk'},
+            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
+        )
+        ]
+    )
+    def test_existing_and_new_cookies_sent_in_request(self, new_cookies, new_cookies_dict, expected, httpbin):
+        r = http(
+            '--session', str(self.session_path),
+            '--print=H',
+            httpbin.url,
+            'Cookie:' + new_cookies,
+        )
+        # Note: cookies in response are in alphabetical order
+        assert 'Cookie: ' + expected in r
+
+        updated_session = json.loads(self.session_path.read_text())
+        for name, value in new_cookies_dict.items():
+            assert name, value in updated_session['cookies']
+            assert 'Cookie' not in updated_session['headers']
+
+    @pytest.mark.parametrize(
+        argnames=['cli_cookie', 'set_cookie', 'expected'],
+        argvalues=[(
+            '',
+            '/cookies/set/cookie1/bar',
+            'bar'
+        ),
+            (
+            'cookie1=not_foo',
+            '/cookies/set/cookie1/bar',
+            'bar'
+        ),
+            (
+            'cookie1=not_foo',
+            '',
+            'not_foo'
+        ),
+            (
+            '',
+            '',
+            'foo'
+        )
+        ]
+    )
+    def test_cookie_storage_priority(self, cli_cookie, set_cookie, expected, httpbin):
+        """
+        Expected order of priority for cookie storage in session file:
+        1. set-cookie (from server)
+        2. command line arg
+        3. cookie already stored in session file
+        """
+        r = http(
+            '--session', str(self.session_path),
+            httpbin.url + set_cookie,
+            'Cookie:' + cli_cookie,
+        )
+        updated_session = json.loads(self.session_path.read_text())
+
+        assert updated_session['cookies']['cookie1']['value'] == expected