|
6 | 6 | import os |
7 | 7 | import sys |
8 | 8 | from unittest import mock |
9 | | -from unittest.mock import patch, mock_open, Mock, MagicMock |
10 | | - |
11 | | -from unittest import IsolatedAsyncioTestCase |
12 | | -import nest_asyncio |
| 9 | +from unittest import TestCase |
| 10 | +from unittest.mock import patch, mock_open, Mock, MagicMock, AsyncMock |
13 | 11 |
|
14 | 12 | import runpod |
15 | 13 | from runpod.serverless.modules.rp_logger import RunPodLogger |
16 | 14 | from runpod.serverless.modules.rp_scale import _handle_uncaught_exception |
17 | 15 | from runpod.serverless import _signal_handler |
18 | 16 |
|
19 | | -nest_asyncio.apply() |
20 | | - |
21 | 17 |
|
22 | | -class TestWorker(IsolatedAsyncioTestCase): |
| 18 | +class TestWorker(TestCase): |
23 | 19 | """Tests for Runpod serverless worker.""" |
24 | 20 |
|
25 | | - async def asyncSetUp(self): |
| 21 | + def setUp(self): |
26 | 22 | self.mock_handler = mock.Mock(return_value="test") |
27 | 23 | self.mock_config = { |
28 | 24 | "handler": self.mock_handler, |
@@ -105,10 +101,10 @@ def test_signal_handler(self, mock_exit, mock_logger): |
105 | 101 | assert mock_logger.info.called |
106 | 102 |
|
107 | 103 |
|
108 | | -class TestWorkerTestInput(IsolatedAsyncioTestCase): |
| 104 | +class TestWorkerTestInput(TestCase): |
109 | 105 | """Tests for runpod | serverless| worker""" |
110 | 106 |
|
111 | | - async def asyncSetUp(self): |
| 107 | + def setUp(self): |
112 | 108 | self.mock_handler = Mock() |
113 | 109 | self.mock_handler.return_value = {} |
114 | 110 |
|
@@ -176,28 +172,38 @@ def test_generator_handler_exception(): |
176 | 172 | assert True, "Exception was caught as expected" |
177 | 173 |
|
178 | 174 |
|
179 | | -class TestRunWorker(IsolatedAsyncioTestCase): |
| 175 | +class TestRunWorker(TestCase): |
180 | 176 | """Tests for runpod | serverless| worker""" |
181 | 177 |
|
182 | | - async def asyncSetUp(self): |
| 178 | + def setUp(self): |
183 | 179 | os.environ["RUNPOD_WEBHOOK_GET_JOB"] = "https://test.com" |
184 | 180 |
|
| 181 | + # run_worker() runs real fitness checks (GPU/memory probes) that exit the |
| 182 | + # process when unmet; they have dedicated coverage in |
| 183 | + # test_modules/test_fitness/. Stub them so these tests are deterministic |
| 184 | + # regardless of host state. |
| 185 | + fitness_patcher = patch( |
| 186 | + "runpod.serverless.worker.run_fitness_checks", new=AsyncMock() |
| 187 | + ) |
| 188 | + fitness_patcher.start() |
| 189 | + self.addCleanup(fitness_patcher.stop) |
| 190 | + |
185 | 191 | # Set up the config |
186 | 192 | self.config = { |
187 | 193 | "handler": MagicMock(), |
188 | 194 | "refresh_worker": True, |
189 | 195 | "rp_args": {"rp_debugger": True, "rp_log_level": "DEBUG"}, |
190 | 196 | } |
191 | 197 |
|
192 | | - async def asyncTearDown(self): |
| 198 | + def tearDown(self): |
193 | 199 | sys.excepthook = sys.__excepthook__ |
194 | 200 |
|
195 | 201 | @patch("runpod.serverless.modules.rp_scale.AsyncClientSession") |
196 | 202 | @patch("runpod.serverless.modules.rp_scale.get_job") |
197 | 203 | @patch("runpod.serverless.modules.rp_job.run_job") |
198 | 204 | @patch("runpod.serverless.modules.rp_job.stream_result") |
199 | 205 | @patch("runpod.serverless.modules.rp_job.send_result") |
200 | | - async def test_run_worker( |
| 206 | + def test_run_worker( |
201 | 207 | self, |
202 | 208 | mock_send_result, |
203 | 209 | mock_stream_result, |
@@ -228,7 +234,7 @@ async def test_run_worker( |
228 | 234 | @patch("runpod.serverless.modules.rp_job.run_job") |
229 | 235 | @patch("runpod.serverless.modules.rp_job.stream_result") |
230 | 236 | @patch("runpod.serverless.modules.rp_job.send_result") |
231 | | - async def test_run_worker_generator_handler( |
| 237 | + def test_run_worker_generator_handler( |
232 | 238 | self, mock_send_result, mock_stream_result, mock_run_job, mock_get_job |
233 | 239 | ): |
234 | 240 | """ |
@@ -258,7 +264,7 @@ async def test_run_worker_generator_handler( |
258 | 264 | @patch("runpod.serverless.modules.rp_job.run_job") |
259 | 265 | @patch("runpod.serverless.modules.rp_job.stream_result") |
260 | 266 | @patch("runpod.serverless.modules.rp_job.send_result") |
261 | | - async def test_run_worker_generator_handler_exception( |
| 267 | + def test_run_worker_generator_handler_exception( |
262 | 268 | self, mock_send_result, mock_stream_result, mock_run_job, mock_get_job |
263 | 269 | ): |
264 | 270 | """ |
@@ -303,7 +309,7 @@ async def test_run_worker_generator_handler_exception( |
303 | 309 | @patch("runpod.serverless.modules.rp_job.run_job") |
304 | 310 | @patch("runpod.serverless.modules.rp_job.stream_result") |
305 | 311 | @patch("runpod.serverless.modules.rp_job.send_result") |
306 | | - async def test_run_worker_generator_aggregate_handler( |
| 312 | + def test_run_worker_generator_aggregate_handler( |
307 | 313 | self, mock_send_result, mock_stream_result, mock_run_job, mock_get_job |
308 | 314 | ): |
309 | 315 | """ |
@@ -343,7 +349,7 @@ async def test_run_worker_generator_aggregate_handler( |
343 | 349 | @patch("runpod.serverless.modules.rp_job.run_job") |
344 | 350 | @patch("runpod.serverless.modules.rp_job.stream_result") |
345 | 351 | @patch("runpod.serverless.modules.rp_job.send_result") |
346 | | - async def test_run_worker_concurrency( |
| 352 | + def test_run_worker_concurrency( |
347 | 353 | self, |
348 | 354 | mock_send_result, |
349 | 355 | mock_stream_result, |
@@ -420,7 +426,7 @@ def concurrency_modifier(current_concurrency): |
420 | 426 | @patch("runpod.serverless.modules.rp_job.run_job") |
421 | 427 | @patch("runpod.serverless.modules.rp_job.stream_result") |
422 | 428 | @patch("runpod.serverless.modules.rp_job.send_result") |
423 | | - async def test_run_worker_multi_processing( |
| 429 | + def test_run_worker_multi_processing( |
424 | 430 | self, |
425 | 431 | mock_send_result, |
426 | 432 | mock_stream_result, |
@@ -480,7 +486,7 @@ async def test_run_worker_multi_processing( |
480 | 486 |
|
481 | 487 | @patch("runpod.serverless.modules.rp_scale.get_job") |
482 | 488 | @patch("runpod.serverless.modules.rp_job.run_job") |
483 | | - async def test_run_worker_multi_processing_scaling_up( |
| 489 | + def test_run_worker_multi_processing_scaling_up( |
484 | 490 | self, mock_run_job, mock_get_job |
485 | 491 | ): |
486 | 492 | """ |
|
0 commit comments