fix: UA test

Author: justinmerrell

tests/test_serverless/test_modules/test_logger.py

@@ -120,7 +120,7 @@ def test_log_job_id(self):
             logger.log("test_message", "INFO", job_id)
 
             mock_print.assert_called_once_with(
-                "'INFO   | test_job_id | test_message', flush=True",
+                "'INFO   | test_job_id | test_message'",
                 flush=True
             )
 

tests/test_user_agent.py

@@ -38,7 +38,10 @@ def test_user_agent_with_integration(
             self, mock_python_version, mock_machine, mock_release, mock_system):
         """Test the User-Agent string with an integration method specified."""
         expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot"  # pylint: disable=line-too-long
+
+        os.environ['RUNPOD_UA_INTEGRATION'] = 'SkyPilot'
         self.assertEqual(construct_user_agent(), expected_ua)
+        os.environ.pop('RUNPOD_UA_INTEGRATION')
 
         assert mock_python_version.called
         assert mock_machine.called