Update test_worker.py

justinmerrell · justinmerrell · commit 2c0e04bf407f · 2024-02-12T16:48:55.000-05:00
diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py
@@ -268,7 +268,11 @@ async def test_run_worker_generator_handler_exception(
         generator_config = {
             "handler": generator_handler_exception, "refresh_worker": True}
 
-        # os.environ.pop("RUNPOD_ENDPOINT_ID")
+        try:
+            os.environ.pop("RUNPOD_ENDPOINT_ID")
+        except KeyError:
+            pass
+
         runpod.serverless.start(generator_config)
 
         assert mock_stream_result.call_count == 1
