From 9d7b727fabdcf384feba9659581a7cca1065519b Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:15:14 -0400 Subject: [PATCH 01/15] fix: cli generator and test output --- .github/workflows/CD-publish_to_pypi.yml | 1 + runpod/serverless/__init__.py | 4 +++ runpod/serverless/modules/rp_fastapi.py | 14 ++++++-- runpod/serverless/modules/rp_handler.py | 8 +++++ runpod/serverless/modules/rp_local.py | 8 +++++ runpod/serverless/worker.py | 6 ++-- .../test_modules/test_handler.py | 33 +++++++++++++++++++ 7 files changed, 69 insertions(+), 5 deletions(-) create mode 100644 runpod/serverless/modules/rp_handler.py create mode 100644 tests/test_serverless/test_modules/test_handler.py diff --git a/.github/workflows/CD-publish_to_pypi.yml b/.github/workflows/CD-publish_to_pypi.yml index d00b533d..356f449b 100644 --- a/.github/workflows/CD-publish_to_pypi.yml +++ b/.github/workflows/CD-publish_to_pypi.yml @@ -69,6 +69,7 @@ jobs: "runpod-workers/worker-controlnet", "runpod-workers/worker-blip", "runpod-workers/worker-deforum", + runpod-workers/mock-worker, ] runs-on: ubuntu-latest diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index dfd86a7b..abf60633 100644 --- a/runpod/serverless/__init__.py +++ b/runpod/serverless/__init__.py @@ -66,6 +66,10 @@ def _set_config_args(config) -> dict: if config["rp_args"]["test_input"]: config["rp_args"]["test_input"] = json.loads(config["rp_args"]["test_input"]) + # Parse the test output from JSON + if config["rp_args"].get("test_output", None): + config["rp_args"]["test_output"] = json.loads(config["rp_args"]["test_output"]) + # Set the log level if config["rp_args"]["rp_log_level"]: log.set_level(config["rp_args"]["rp_log_level"]) diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 1ba83f44..dfd102ad 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -9,7 +9,8 @@ from fastapi.encoders import jsonable_encoder from pydantic import BaseModel -from .rp_job import run_job +from .rp_handler import is_generator +from .rp_job import run_job, run_job_generator from .worker_state import Jobs from .rp_ping import Heartbeat from ...version import __version__ as runpod_version @@ -125,7 +126,16 @@ async def _debug_run(self, job: TestJob): # Set the current job ID. job_list.add_job(job.id) - job_results = await run_job(self.config["handler"], job.__dict__) + if is_generator(self.config["handler"]): + generator_output = run_job_generator(self.config["handler"], job) + job_results = {"output": []} + async for stream_output in generator_output: + if "error" in stream_output: + job_results = stream_output + break + job_results["output"].append(stream_output["output"]) + else: + job_results = await run_job(self.config["handler"], job.__dict__) job_results["id"] = job.id diff --git a/runpod/serverless/modules/rp_handler.py b/runpod/serverless/modules/rp_handler.py new file mode 100644 index 00000000..40f62128 --- /dev/null +++ b/runpod/serverless/modules/rp_handler.py @@ -0,0 +1,8 @@ +"""Retrieve handler info. """ + +import inspect +from typing import Callable + +def is_generator(handler: Callable) -> bool: + """Check if handler is a generator function. """ + return inspect.isgeneratorfunction(handler) or inspect.isasyncgenfunction(handler) diff --git a/runpod/serverless/modules/rp_local.py b/runpod/serverless/modules/rp_local.py index 5716531d..7db3c69c 100644 --- a/runpod/serverless/modules/rp_local.py +++ b/runpod/serverless/modules/rp_local.py @@ -47,5 +47,13 @@ async def run_local(config: Dict[str, Any]) -> None: log.info(f"Job {local_job['id']} completed successfully.") log.info(f"Job result: {job_result}") + # Compare to sample output, if provided + if config['rp_args'].get('test_output', None): + log.info("test_output set, comparing output to test_output.") + if job_result != config['rp_args']['test_output']: + log.error("Job output does not match test_output.") + sys.exit(1) + log.info("Job output matches test_output.") + log.info("Local testing complete, exiting.") sys.exit(0) diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 82d6883f..f24518d1 100644 --- a/runpod/serverless/worker.py +++ b/runpod/serverless/worker.py @@ -12,6 +12,7 @@ from runpod.serverless.modules.rp_logger import RunPodLogger from runpod.serverless.modules.rp_scale import JobScaler from .modules import rp_local +from .modules.rp_handler import is_generator from .modules.rp_ping import Heartbeat from .modules.rp_job import run_job, run_job_generator from .modules.rp_http import send_result, stream_result @@ -46,11 +47,10 @@ def _is_local(config) -> bool: async def _process_job(job, session, job_scaler, config): - if inspect.isgeneratorfunction(config["handler"]) \ - or inspect.isasyncgenfunction(config["handler"]): + if is_generator(config["handler"]): generator_output = run_job_generator(config["handler"], job) - log.debug("Handler is a generator, streaming results.") + job_result = {'output': []} async for stream_output in generator_output: if 'error' in stream_output: diff --git a/tests/test_serverless/test_modules/test_handler.py b/tests/test_serverless/test_modules/test_handler.py new file mode 100644 index 00000000..d055a337 --- /dev/null +++ b/tests/test_serverless/test_modules/test_handler.py @@ -0,0 +1,33 @@ +""" Unit tests for the handler module. +""" +import unittest + +from runpod.serverless.modules.rp_handler import is_generator + + +class TestIsGenerator(unittest.TestCase): + """Tests for the is_generator function.""" + + def test_regular_function(self): + """Test that a regular function is not a generator.""" + def regular_func(): + return "I'm a regular function!" + self.assertFalse(is_generator(regular_func)) + + def test_generator_function(self): + """Test that a generator function is a generator.""" + def generator_func(): + yield "I'm a generator function!" + self.assertTrue(is_generator(generator_func)) + + def test_async_function(self): + """Test that an async function is not a generator.""" + async def async_func(): + return "I'm an async function!" + self.assertFalse(is_generator(async_func)) + + def test_async_generator_function(self): + """Test that an async generator function is a generator.""" + async def async_gen_func(): + yield "I'm an async generator function!" + self.assertTrue(is_generator(async_gen_func)) From fc3ff409b02f88ea9e76b10c9fed8ce299b5459d Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:22:27 -0400 Subject: [PATCH 02/15] fix: added tests --- CHANGELOG.md | 1 + runpod/serverless/worker.py | 1 - tests/test_serverless/test_worker.py | 3 +-- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9dea7cf7..0120ce90 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Added - BETA: CLI DevEx functionality to create development projects. +- `test_output` can be passed in as an arg to compare the results of `test_input` ## Release 1.3.0 (10/12/23) diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index f24518d1..f8694863 100644 --- a/runpod/serverless/worker.py +++ b/runpod/serverless/worker.py @@ -4,7 +4,6 @@ """ import os import asyncio -import inspect from typing import Dict, Any import aiohttp diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index a94a67e5..93c4ebda 100644 --- a/tests/test_serverless/test_worker.py +++ b/tests/test_serverless/test_worker.py @@ -44,7 +44,7 @@ def test_is_local(self): with patch("runpod.serverless.worker.os") as mock_os: mock_os.environ.get.return_value = None assert runpod.serverless.worker._is_local({"rp_args": {}}) is True # pylint: disable=protected-access - assert runpod.serverless.worker._is_local({"rp_args": {"test_input": "something"}}) is True # pylint: disable=protected-access, line-too-long + assert runpod.serverless.worker._is_local({"rp_args": {"test_input": "something", "test_output": "something"}}) is True # pylint: disable=protected-access, line-too-long mock_os.environ.get.return_value = "something" assert runpod.serverless.worker._is_local(self.mock_config) is False # pylint: disable=protected-access @@ -75,7 +75,6 @@ def test_local_api(self): ''' Test local FastAPI setup. ''' - known_args = argparse.Namespace() known_args.rp_log_level = None known_args.rp_debugger = None From 17d3805544e1cf1ed1f818558f6e208abe483858 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:29:22 -0400 Subject: [PATCH 03/15] fix: tests --- CHANGELOG.md | 1 + runpod/serverless/modules/rp_fastapi.py | 3 --- tests/test_serverless/test_modules/test_fastapi.py | 7 +++++++ tests/test_serverless/test_worker.py | 3 ++- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0120ce90..3d08acf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - BETA: CLI DevEx functionality to create development projects. - `test_output` can be passed in as an arg to compare the results of `test_input` +- Generator/Streaming handlers supported with local testing ## Release 1.3.0 (10/12/23) diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index dfd102ad..8194ba78 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -130,9 +130,6 @@ async def _debug_run(self, job: TestJob): generator_output = run_job_generator(self.config["handler"], job) job_results = {"output": []} async for stream_output in generator_output: - if "error" in stream_output: - job_results = stream_output - break job_results["output"].append(stream_output["output"]) else: job_results = await run_job(self.config["handler"], job.__dict__) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index e85bf99b..05c08538 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -87,4 +87,11 @@ def test_run(self): self.assertTrue(mock_ping.called) + # Test with generator handler + generator_handler = Mock() + generator_handler.return_value = [{"result": "success"}] + generator_run_return = asyncio.run(worker_api._run(job_object)) # pylint: disable=protected-access + + assert generator_run_return == {"output": {"result": "success"}} + loop.close() diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index 93c4ebda..888022bf 100644 --- a/tests/test_serverless/test_worker.py +++ b/tests/test_serverless/test_worker.py @@ -44,7 +44,7 @@ def test_is_local(self): with patch("runpod.serverless.worker.os") as mock_os: mock_os.environ.get.return_value = None assert runpod.serverless.worker._is_local({"rp_args": {}}) is True # pylint: disable=protected-access - assert runpod.serverless.worker._is_local({"rp_args": {"test_input": "something", "test_output": "something"}}) is True # pylint: disable=protected-access, line-too-long + assert runpod.serverless.worker._is_local({"rp_args": {"test_input": "something"}}) is True # pylint: disable=protected-access, line-too-long mock_os.environ.get.return_value = "something" assert runpod.serverless.worker._is_local(self.mock_config) is False # pylint: disable=protected-access @@ -125,6 +125,7 @@ def test_worker_bad_local(self): known_args.rp_api_concurrency = 1 known_args.rp_api_host = "localhost" known_args.test_input = '{"test": "test"}' + known_args.test_output = '{"test": "test"}' with patch("argparse.ArgumentParser.parse_known_args") as mock_parse_known_args, \ self.assertRaises(SystemExit): From 3fc506fc9caa395063595a213ad0385ce9a5f6f6 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:34:48 -0400 Subject: [PATCH 04/15] fix: tests --- tests/test_serverless/test_modules/test_fastapi.py | 5 +++-- tests/test_serverless/test_modules/test_local.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index 05c08538..7580645a 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -88,8 +88,9 @@ def test_run(self): self.assertTrue(mock_ping.called) # Test with generator handler - generator_handler = Mock() - generator_handler.return_value = [{"result": "success"}] + def generator_handler(): + yield {"result": "success"} + generator_run_return = asyncio.run(worker_api._run(job_object)) # pylint: disable=protected-access assert generator_run_return == {"output": {"result": "success"}} diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index f242b4b5..e650b4ed 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -21,6 +21,10 @@ async def test_run_local_with_test_input(self, mock_file, mock_run): "test_input": { "input": "test", "id": "test_id" + }, + "test_output": { + "output": "test", + "id": "test_id" } } } From 522492225b4b72d69be5e56c5a0f1407d64a3418 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:51:13 -0400 Subject: [PATCH 05/15] Update test_fastapi.py --- tests/test_serverless/test_modules/test_fastapi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index 7580645a..90903ad9 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -91,8 +91,8 @@ def test_run(self): def generator_handler(): yield {"result": "success"} - generator_run_return = asyncio.run(worker_api._run(job_object)) # pylint: disable=protected-access - + generator_worker_api = rp_fastapi.WorkerAPI(handler=generator_handler) + generator_run_return = asyncio.run(generator_worker_api._run(job_object)) # pylint: disable=protected-access assert generator_run_return == {"output": {"result": "success"}} loop.close() From 77fabae3d3f5ec2ceb01e97d0c22e4f75a9d5394 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:55:44 -0400 Subject: [PATCH 06/15] Update test_fastapi.py --- tests/test_serverless/test_modules/test_fastapi.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index 90903ad9..0a5c7bc6 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -88,7 +88,8 @@ def test_run(self): self.assertTrue(mock_ping.called) # Test with generator handler - def generator_handler(): + def generator_handler(job): + del job yield {"result": "success"} generator_worker_api = rp_fastapi.WorkerAPI(handler=generator_handler) From 67c11291a030f67ee21e9c72d3fad65305b61e1b Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 18:57:37 -0400 Subject: [PATCH 07/15] Update test_fastapi.py --- tests/test_serverless/test_modules/test_fastapi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index 0a5c7bc6..dbabf4d9 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -93,7 +93,7 @@ def generator_handler(job): yield {"result": "success"} generator_worker_api = rp_fastapi.WorkerAPI(handler=generator_handler) - generator_run_return = asyncio.run(generator_worker_api._run(job_object)) # pylint: disable=protected-access + generator_run_return = asyncio.run(generator_worker_api._debug_run(job_object)) # pylint: disable=protected-access assert generator_run_return == {"output": {"result": "success"}} loop.close() From ef96d2f99f71c206ddb5ed0f9037733e45e89c73 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:03:33 -0400 Subject: [PATCH 08/15] Update rp_fastapi.py --- runpod/serverless/modules/rp_fastapi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 8194ba78..aa4b1a8f 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -127,7 +127,7 @@ async def _debug_run(self, job: TestJob): job_list.add_job(job.id) if is_generator(self.config["handler"]): - generator_output = run_job_generator(self.config["handler"], job) + generator_output = run_job_generator(self.config["handler"], job.__dict__) job_results = {"output": []} async for stream_output in generator_output: job_results["output"].append(stream_output["output"]) From d4b7959c42427f7b845ec4637c91941291ee581b Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:08:48 -0400 Subject: [PATCH 09/15] Update test_fastapi.py --- tests/test_serverless/test_modules/test_fastapi.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_serverless/test_modules/test_fastapi.py b/tests/test_serverless/test_modules/test_fastapi.py index dbabf4d9..d8dedc80 100644 --- a/tests/test_serverless/test_modules/test_fastapi.py +++ b/tests/test_serverless/test_modules/test_fastapi.py @@ -94,6 +94,9 @@ def generator_handler(job): generator_worker_api = rp_fastapi.WorkerAPI(handler=generator_handler) generator_run_return = asyncio.run(generator_worker_api._debug_run(job_object)) # pylint: disable=protected-access - assert generator_run_return == {"output": {"result": "success"}} + assert generator_run_return == { + "id": "test_job_id", + "output": [{"result": "success"}] + } loop.close() From 6eeecd330c8967b4c72d75431a27ab5b2bb649bb Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:13:30 -0400 Subject: [PATCH 10/15] Update test_local.py --- tests/test_serverless/test_modules/test_local.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index e650b4ed..b92f8727 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -23,7 +23,7 @@ async def test_run_local_with_test_input(self, mock_file, mock_run): "id": "test_id" }, "test_output": { - "output": "test", + "input": "test", "id": "test_id" } } From d077063cc2d92747676afe0eff5f9846e1df918c Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:20:15 -0400 Subject: [PATCH 11/15] Update test_local.py --- tests/test_serverless/test_modules/test_local.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index b92f8727..761f9ec4 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -22,10 +22,7 @@ async def test_run_local_with_test_input(self, mock_file, mock_run): "input": "test", "id": "test_id" }, - "test_output": { - "input": "test", - "id": "test_id" - } + "test_output": {} } } with self.assertRaises(SystemExit) as sys_exit: From 32b79cae0c1052af386885b077501741bf5449ab Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:28:51 -0400 Subject: [PATCH 12/15] Update test_local.py --- tests/test_serverless/test_modules/test_local.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index 761f9ec4..87cadebb 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -9,7 +9,7 @@ class TestRunLocal(IsolatedAsyncioTestCase): ''' Tests for run_local function ''' - @patch("runpod.serverless.modules.rp_local.run_job", return_value={}) + @patch("runpod.serverless.modules.rp_local.run_job", return_value={"result": "success"}) @patch("builtins.open", new_callable=mock_open, read_data='{"input": "test"}') async def test_run_local_with_test_input(self, mock_file, mock_run): ''' @@ -22,7 +22,9 @@ async def test_run_local_with_test_input(self, mock_file, mock_run): "input": "test", "id": "test_id" }, - "test_output": {} + "test_output": { + "result": "success" + } } } with self.assertRaises(SystemExit) as sys_exit: From 6d3bc56e05096b3e9acfbcc6f4f05d871abb91bb Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:33:37 -0400 Subject: [PATCH 13/15] Update test_local.py --- tests/test_serverless/test_modules/test_local.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index 87cadebb..cba0ffcf 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -29,7 +29,12 @@ async def test_run_local_with_test_input(self, mock_file, mock_run): } with self.assertRaises(SystemExit) as sys_exit: await rp_local.run_local(config) - self.assertEqual(sys_exit.exception.code, 0) + self.assertEqual(sys_exit.exception.code, 0) + + config["rp_args"]["test_output"] = {"result": "fail"} + with self.assertRaises(SystemExit) as sys_exit: + await rp_local.run_local(config) + self.assertEqual(sys_exit.exception.code, 1) assert mock_file.called is False assert mock_run.called From 709cafa0ee83d0eecf06554ed2ea34baf08bb1d5 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:54:32 -0400 Subject: [PATCH 14/15] Update rp_ping.py --- runpod/serverless/modules/rp_ping.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/runpod/serverless/modules/rp_ping.py b/runpod/serverless/modules/rp_ping.py index 5d1a6e96..9e26f6d8 100644 --- a/runpod/serverless/modules/rp_ping.py +++ b/runpod/serverless/modules/rp_ping.py @@ -52,6 +52,10 @@ def start_ping(self, test=False): ''' Sends heartbeat pings to the Runpod server. ''' + if os.environ.get('RUNPOD_AI_API_KEY') is None: + log.debug("Not deployed on RunPod serverless, pings will not be sent.") + return + if os.environ.get('RUNPOD_POD_ID') is None: log.info("Not running on RunPod, pings will not be sent.") return From d3012ae74382ad164fa7ee742b7d8bc6c78286b4 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Fri, 27 Oct 2023 19:55:30 -0400 Subject: [PATCH 15/15] Update test_ping.py --- tests/test_serverless/test_modules/test_ping.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_serverless/test_modules/test_ping.py b/tests/test_serverless/test_modules/test_ping.py index 3e57560c..751d45ea 100644 --- a/tests/test_serverless/test_modules/test_ping.py +++ b/tests/test_serverless/test_modules/test_ping.py @@ -46,6 +46,13 @@ def test_start_ping(self, mock_get_return): ''' Tests that the start_ping function works correctly ''' + # No RUNPOD_AI_API_KEY case + with patch("threading.Thread.start") as mock_thread_start: + rp_ping.Heartbeat().start_ping(test=True) + assert mock_thread_start.call_count == 0 + + os.environ["RUNPOD_AI_API_KEY"] = "test_key" + # No RUNPOD_POD_ID case with patch("threading.Thread.start") as mock_thread_start: rp_ping.Heartbeat().start_ping(test=True)