update truncation

justinmerrell · justinmerrell · commit 300d73a0ec4b · 2024-02-12T17:20:35.000-05:00
diff --git a/runpod/serverless/modules/rp_logger.py b/runpod/serverless/modules/rp_logger.py
@@ -76,7 +76,9 @@ def log(self, message, message_level='INFO', job_id=None):
         # Truncate message over 10MB, remove chunk from the middle
         if len(message) > MAX_MESSAGE_LENGTH:
             half_max_length = MAX_MESSAGE_LENGTH // 2
-            message = message[:half_max_length] + '\n...TRUNCATED...\n' + message[-half_max_length:]
+            truncated_amount = len(message) - MAX_MESSAGE_LENGTH
+            truncation_note = f'\n...TRUNCATED {truncated_amount} CHARACTERS...\n'
+            message = message[:half_max_length] + truncation_note + message[-half_max_length:]
 
         if os.environ.get('RUNPOD_ENDPOINT_ID'):
             log_json = {
diff --git a/tests/test_serverless/test_modules/test_logger.py b/tests/test_serverless/test_modules/test_logger.py
@@ -135,18 +135,19 @@ def test_log_job_id(self):
             )
 
     def test_log_truncate(self):
-        """Tests that the log method truncates long messages."""
+        """Tests that the log method truncates """
         logger = rp_logger.RunPodLogger()
         job_id = "test_job_id"
-        long_message = "a" * (rp_logger.MAX_MESSAGE_LENGTH + 100)  # Exceeds the max length
+        long_message = "a" * (rp_logger.MAX_MESSAGE_LENGTH + 100)
         expected_start = "a" * (rp_logger.MAX_MESSAGE_LENGTH // 2)
         expected_end = "a" * (rp_logger.MAX_MESSAGE_LENGTH // 2)
-        truncated_message = expected_start + '\n...TRUNCATED...\n' + expected_end
+        truncated_amount = len(long_message) - rp_logger.MAX_MESSAGE_LENGTH
+        truncation_note = f'\n...TRUNCATED {truncated_amount} CHARACTERS...\n'
+        truncated_message = expected_start + truncation_note + expected_end
 
         with patch("builtins.print") as mock_print:
             logger.log(long_message, "INFO", job_id)
 
-            # Construct the expected log output
             expected_log_output = f'INFO   | {job_id} | {truncated_message}'
 
             mock_print.assert_called_once_with(expected_log_output, flush=True)
diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py
@@ -274,8 +274,8 @@ async def test_run_worker_generator_handler_exception(
         except KeyError:
             pass
 
-        with patch("os.environ.get") as mock_get:
-            mock_get.return_value = None
+        with patch("runpod.serverless.worker.os") as mock_os:
+            mock_os.environ.get.return_value = None
             runpod.serverless.start(generator_config)
 
         assert mock_stream_result.call_count == 1
