Merge 761853d into d36c1fd

MasoniteFramework · Sep 21, 2018 · 66d9b8d · 66d9b8d
2 parents d36c1fd + 761853d
commit 66d9b8d
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 11 deletions.
diff --git a/masonite/request.py b/masonite/request.py
@@ -315,20 +315,26 @@ def get_request_method(self):
 
         return self.environ['REQUEST_METHOD']
 
-    def header(self, key, value=None, http_prefix=True):
-        """Sets or gets a header depending on if value is passed in or not.
+    def header(self, key, value=None, http_prefix=None):
+        """Sets or gets a header depending on if "value" is passed in or not.
 
         Arguments:
-            key {string} -- The header you want to set or get.
+            key {string|dict} -- The header you want to set or get. If the key is a dictionary, loop through each key pair 
+                                    and add them to the headers.
 
         Keyword Arguments:
             value {string} -- The value you want to set (default: {None})
             http_prefix {bool} -- Whether it should have `HTTP_` prefixed to the value being set. (default: {True})
 
         Returns:
-            string|True|None -- [description]
+            string|None|True -- Either return the value if getting a header, None if it doesn't exist or True if setting the headers.
         """
 
+        if isinstance(key, dict):
+            for key, value in key.items():
+                self._set_header(key, value, http_prefix)
+            return True
+
         # Get Headers
         if value is None:
             if 'HTTP_{0}'.format(key) in self.environ:
@@ -338,14 +344,19 @@ def header(self, key, value=None, http_prefix=True):
             else:
                 return None
 
+        self._set_header(key, value, http_prefix)
+
+        return True
+
+    def _set_header(self, key, value, http_prefix):
         # Set Headers
         if http_prefix:
+            print('http_prefix', http_prefix)
             self.environ['HTTP_{0}'.format(key)] = str(value)
             self._headers.append(('HTTP_{0}'.format(key), str(value)))
         else:
             self.environ[key] = str(value)
             self._headers.append((key, str(value)))
-        return True
 
     def get_headers(self):
         """Returns all current headers to be set.

diff --git a/tests/test_requests.py b/tests/test_requests.py
@@ -235,11 +235,11 @@ def test_request_has_subdomain_returns_bool(self):
         request.environ['HTTP_HOST'] = 'test.localhost.com'
 
         request.header('TEST', 'set_this')
-        assert request.header('HTTP_TEST') == 'set_this'
-
-        request.header('TEST', 'set_this', http_prefix = None)
         assert request.header('TEST') == 'set_this'
 
+        request.header('TEST', 'set_this', http_prefix = True)
+        assert request.header('HTTP_TEST') == 'set_this'
+
 
     def test_redirect_compiles_url(self):
         app = App()
@@ -365,11 +365,31 @@ def test_request_sets_correct_header(self):
         request = app.make('Request').load_app(app)
 
         request.header('TEST', 'set_this')
+        assert request.header('TEST') == 'set_this'
+
+        request.header('TEST', 'set_this', http_prefix = True)
         assert request.header('HTTP_TEST') == 'set_this'
 
-        request.header('TEST', 'set_this', http_prefix = None)
-        assert request.header('TEST') == 'set_this'
+    def test_request_sets_headers_with_dictionary(self):
+        app = App()
+        app.bind('Request', self.request)
+        request = app.make('Request').load_app(app)
+
+        request.header({
+            'test_dict': 'test_value',
+            'test_dict1': 'test_value1'
+        })
+
+        assert request.header('test_dict') == 'test_value'
+        assert request.header('test_dict1') == 'test_value1'
+
+        request.header({
+            'test_dict': 'test_value',
+            'test_dict1': 'test_value1'
+        }, http_prefix = True)
 
+        assert request.header('HTTP_test_dict') == 'test_value'
+        assert request.header('HTTP_test_dict1') == 'test_value1'
 
     def test_request_gets_all_headers(self):
         app = App()
@@ -378,7 +398,7 @@ def test_request_gets_all_headers(self):
 
         request.header('TEST1', 'set_this_item')
         request.header('TEST2', 'set_this_item', http_prefix = None)
-        assert request.get_headers() == [('HTTP_TEST1', 'set_this_item'), ('TEST2', 'set_this_item')]
+        assert request.get_headers() == [('TEST1', 'set_this_item'), ('TEST2', 'set_this_item')]
 
     def test_request_sets_status_code(self):
         app = App()