add tests for chunking

Author: Dan Hertz

nightfall/api.py

@@ -4,11 +4,9 @@
 ~~~~~~~~~~~~~
     This module provides a class which abstracts the Nightfall REST API.
 """
-import time
 from datetime import datetime, timedelta
 import hmac
 import hashlib
-import json
 import logging
 import os
 from typing import Dict, List, Tuple, Optional
@@ -103,7 +101,7 @@ def scan_text(self, texts: List[str], detection_rules: Optional[List[DetectionRu
         return findings, parsed_response.get("redactedPayload")
 
     def _scan_text_v3(self, data: dict):
-        response = self.session.post(url=self.TEXT_SCAN_ENDPOINT_V3, data=json.dumps(data))
+        response = self.session.post(url=self.TEXT_SCAN_ENDPOINT_V3, json=data)
 
         self.logger.debug(f"HTTP Request URL: {response.request.url}")
         self.logger.debug(f"HTTP Request Body: {response.request.body}")
@@ -166,7 +164,7 @@ def _file_scan_initialize(self, location: str):
         data = {
             "fileSizeBytes": os.path.getsize(location)
         }
-        response = self.session.post(url=self.FILE_SCAN_INITIALIZE_ENDPOINT, data=json.dumps(data))
+        response = self.session.post(url=self.FILE_SCAN_INITIALIZE_ENDPOINT, json=data)
 
         return response
 
@@ -216,7 +214,7 @@ def _file_scan_scan(self, session_id: str, detection_rules: Optional[List[Detect
         if request_metadata:
             data["requestMetadata"] = request_metadata
 
-        response = self.session.post(url=self.FILE_SCAN_SCAN_ENDPOINT.format(session_id), data=json.dumps(data))
+        response = self.session.post(url=self.FILE_SCAN_SCAN_ENDPOINT.format(session_id), json=data)
         return response
 
     def validate_webhook(self, request_signature: str, request_timestamp: str, request_data: str) -> bool:

tests/test_api.py

@@ -269,17 +269,49 @@ def test_scan_file(tmpdir):
     for call in responses.calls:
         assert call.request.headers.get("Authorization") == "Bearer NF-NOT_REAL"
 
-    assert responses.calls[0].request.body == '{"fileSizeBytes": 44}'
+    assert responses.calls[0].request.body == b'{"fileSizeBytes": 44}'
     assert responses.calls[1].request.body == b"4916-6734-7572-5015 is"
     assert responses.calls[1].request.headers.get("X-UPLOAD-OFFSET") == '0'
     assert responses.calls[2].request.body == b" my credit card number"
     assert responses.calls[2].request.headers.get("X-UPLOAD-OFFSET") == '22'
-    assert responses.calls[4].request.body == '{"policy": {"webhookURL": "https://my-website.example/callback", ' \
-                                              '"detectionRuleUUIDs": ["a_uuid"]}, "requestMetadata": "some test data"}'
+    assert responses.calls[4].request.body == b'{"policy": {"webhookURL": "https://my-website.example/callback", ' \
+                                              b'"detectionRuleUUIDs": ["a_uuid"]}, "requestMetadata": "some test data"}'
     assert id == 1
     assert message == "scan_started"
 
 
+@responses.activate
+def test_file_scan_upload_short(tmpdir):
+    file = tmpdir.mkdir("test_data").join("file.txt")
+
+    file.write("4916-6734-7572-5015 is my credit card number")
+
+    nightfall = Nightfall("NF-NOT_REAL")
+
+    responses.add(responses.PATCH, 'https://api.nightfall.ai/v3/upload/1', status=204)
+
+    assert nightfall._file_scan_upload(1, file, 200)
+    assert len(responses.calls) == 1
+
+
+@responses.activate
+def test_file_scan_upload_long(tmpdir):
+    file = tmpdir.mkdir("test_data").join("file.txt")
+    test_str = b"4916-6734-7572-5015 is my credit card number"
+    file.write_binary(test_str)
+
+    responses.add(responses.PATCH, 'https://api.nightfall.ai/v3/upload/1', status=204)
+
+    nightfall = Nightfall("NF-NOT_REAL")
+
+    assert nightfall._file_scan_upload(1, file, 1)
+    assert len(responses.calls) == 44
+    for i, call in enumerate(responses.calls):
+        assert call.request.headers.get("Authorization") == "Bearer NF-NOT_REAL"
+        assert call.request.body.decode('utf-8') == test_str.decode('utf-8')[i]
+        assert call.request.headers.get("X-UPLOAD-OFFSET") == str(i)
+        
+
 @freeze_time("2021-10-04T17:30:50Z")
 def test_validate_webhook(nightfall):
     nightfall.signing_secret = "super-secret-shhhh"