From ee0ab9f1fabd801fa2c964b1e46fb8c82365bd87 Mon Sep 17 00:00:00 2001
From: JF Paradis <jparadis@appomni.com>
Date: Thu, 2 Feb 2023 20:57:57 -0800
Subject: [PATCH] added omit_headers

---
 drf_batch_requests/request.py     |  1 +
 drf_batch_requests/response.py    |  8 ++++++--
 drf_batch_requests/serializers.py |  2 ++
 drf_batch_requests/views.py       |  1 +
 tests/test_view.py                | 24 ++++++++++++++++++++++++
 5 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/drf_batch_requests/request.py b/drf_batch_requests/request.py
index 1b1e64b..ad8cf59 100644
--- a/drf_batch_requests/request.py
+++ b/drf_batch_requests/request.py
@@ -23,6 +23,7 @@ class BatchRequest(HttpRequest):
     def __init__(self, request, request_data):
         super(BatchRequest, self).__init__()
         self.name = request_data.get('name')
+        self.omit_headers = request_data.get('omit_headers', False)
         self.omit_response_on_success = request_data.get('omit_response_on_success', False)
 
         self._stream = BytesIO(request_data['_body'].encode('utf-8'))
diff --git a/drf_batch_requests/response.py b/drf_batch_requests/response.py
index be950cf..7d98c38 100644
--- a/drf_batch_requests/response.py
+++ b/drf_batch_requests/response.py
@@ -27,14 +27,18 @@ class BatchResponse:
     _return_body: bool = True
 
     def __init__(self, name: str, status_code: int, body: str, headers: Iterable[ResponseHeader] = None,
-                 omit_response_on_success: bool = False, status_text: str = None):
+                 omit_headers: bool = False, omit_response_on_success: bool = False, status_text: str = None):
         self.name = name
         self.status_code = status_code
         self.status_text = status_text
         self.body = body
-        self.headers = headers or []
         self.omit_response_on_success = omit_response_on_success
 
+        if omit_headers:
+            self.headers = []
+        else:
+            self.headers = headers or []
+
         if is_success(self.status_code):
             try:
                 self._data = json.loads(self.body)
diff --git a/drf_batch_requests/serializers.py b/drf_batch_requests/serializers.py
index 0b565ff..133ec22 100644
--- a/drf_batch_requests/serializers.py
+++ b/drf_batch_requests/serializers.py
@@ -21,6 +21,8 @@ class SingleRequestSerializer(serializers.Serializer):
     data = serializers.SerializerMethodField()
     files = serializers.SerializerMethodField()
 
+    omit_headers = serializers.BooleanField(required=False)
+
     def validate_headers(self, value):
         if isinstance(value, dict):
             return value
diff --git a/drf_batch_requests/views.py b/drf_batch_requests/views.py
index a04f7bb..9081f3c 100644
--- a/drf_batch_requests/views.py
+++ b/drf_batch_requests/views.py
@@ -74,6 +74,7 @@ def post(self, request, *args, **kwargs):
                         ResponseHeader(key, value)
                         for key, value in header_items
                     ],
+                    omit_headers=current_request.omit_headers,
                     omit_response_on_success=current_request.omit_response_on_success,
                     status_text=response.reason_phrase
                 )
diff --git a/tests/test_view.py b/tests/test_view.py
index b96d8e2..ebd1d2f 100644
--- a/tests/test_view.py
+++ b/tests/test_view.py
@@ -33,6 +33,30 @@ def test_json_batch(self):
             ','.join([str(o['id']) for o in responses_data[0]['data']])
         )
 
+    def test_omit_headers(self):
+        batch = [
+            {
+                "method": "GET",
+                "relative_url": "/tests/test/",
+                "omit_headers": True,
+            },
+            {
+                "method": "GET",
+                "relative_url": "/tests/test/",
+                "omit_headers": False,
+            },
+            {
+                "method": "GET",
+                "relative_url": "/tests/test/",
+            },
+        ]
+
+        responses = self.forced_auth_req('post', '/batch/', data={'batch': batch})
+
+        self.assertFalse(responses.data[0]['headers'])
+        self.assertTrue(responses.data[1]['headers'])
+        self.assertTrue(responses.data[2]['headers'])
+
     def test_multipart_simple_request(self):
         batch = [
             {