use @patch and reformat

Author: joeyzhao2018

datadog_lambda/tracing.py

@@ -344,12 +344,14 @@ def extract_context_custom_extractor(extractor, event, lambda_context):
 
         return None, None, None
 
+
 def is_authorizer_response(response) -> bool:
     try:
         return response["principalId"] and response["policyDocument"]
     except:
         return False
 
+
 def get_injected_authorizer_data(event, is_http_api) -> dict:
     try:
         authorizer_headers = event.get("requestContext", {}).get("authorizer")

tests/test_wrapper.py

@@ -545,10 +545,9 @@ def lambda_handler(event, context):
         self.assertEquals(inject_data[TraceHeader.SAMPLING_PRIORITY], "1")
         self.assertEquals(result["context"]["scope"], "still here")
 
-    def test_different_return_type_no_error(self):
-        TEST_RESULTS = ["a str to return", 42, {'value': 42}, ['A', 42]]
-        import traceback as _importedTraceback
-        _importedTraceback.print_exc = MagicMock()
+    @patch("traceback.print_exc")
+    def test_different_return_type_no_error(self, MockPrintExc):
+        TEST_RESULTS = ["a str to return", 42, {"value": 42}, ["A", 42]]
         mock_context = get_mock_context()
         for test_result in TEST_RESULTS:
 
@@ -558,5 +557,4 @@ def return_a_str(event, context):
 
             result = return_a_str({}, mock_context)
             self.assertEquals(result, test_result)
-            _importedTraceback.print_exc.assert_not_called()
-
+            self.assertFalse(MockPrintExc.called)