From 69423f02e00d3648b0a4d896d5dea4b91a9e632c Mon Sep 17 00:00:00 2001 From: David Gardner Date: Wed, 13 Dec 2023 15:31:25 -0800 Subject: [PATCH] Mocked sleep times are now fixtures which can be overriden with an environtment variable --- tests/benchmarks/conftest.py | 36 +++++++++++++++---- .../test_bench_agents_simple_pipeline.py | 6 ++-- .../test_bench_vdb_upload_pipeline.py | 6 ++-- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index c65b583765..ff83e4b9a3 100644 --- a/tests/benchmarks/conftest.py +++ b/tests/benchmarks/conftest.py @@ -77,15 +77,38 @@ def pytest_benchmark_update_json(config, benchmarks, output_json): bench['stats']['median_throughput_bytes'] = (byte_count * repeat) / bench['stats']['median'] +@pytest.fixture(name="mock_openai_request_time") +def mock_openai_request_time_fixture(): + return float(os.environ.get("MOCK_OPENAI_REQUEST_TIME", 1.265)) + + +@pytest.fixture(name="mock_nemollm_request_time") +def mock_nemollm_request_time_fixture(): + return float(os.environ.get("MOCK_NEMOLLM_REQUEST_TIME", 0.412)) + + +@pytest.fixture(name="mock_web_scraper_request_time") +def mock_web_scraper_request_time_fixture(): + return float(os.environ.get("MOCK_WEB_SCRAPER_REQUEST_TIME", 0.5)) + + +@pytest.fixture(name="mock_feedparser_request_time") +def mock_feedparser_request_time_fixture(): + return float(os.environ.get("MOCK_FEEDPARSER_REQUEST_TIME", 0.5)) + + +@pytest.fixture(name="mock_serpapi_request_time") +def mock_serpapi_request_time_fixture(): + return float(os.environ.get("MOCK_SERPAPI_REQUEST_TIME", 1.7)) + + @pytest.mark.usefixtures("openai") @pytest.fixture(name="mock_chat_completion") -@pytest.mark.usefixtures() -def mock_chat_completion_fixture(mock_chat_completion: mock.MagicMock): - sleep_time = float(os.environ.get("MOCK_OPENAI_REQUEST_TIME", 1.265)) +def mock_chat_completion_fixture(mock_chat_completion: mock.MagicMock, mock_openai_request_time: float): async def sleep_first(*args, **kwargs): # Sleep time is based on average request time - await asyncio.sleep(sleep_time) + await asyncio.sleep(mock_openai_request_time) return mock.DEFAULT mock_chat_completion.acreate.side_effect = sleep_first @@ -95,13 +118,12 @@ async def sleep_first(*args, **kwargs): @pytest.mark.usefixtures("nemollm") @pytest.fixture(name="mock_nemollm") -def mock_nemollm_fixture(mock_nemollm: mock.MagicMock): - sleep_time = float(os.environ.get("MOCK_NEMOLLM_REQUEST_TIME", 0.412)) +def mock_nemollm_fixture(mock_nemollm: mock.MagicMock, mock_nemollm_request_time: float): # The generate function is a blocking call that returns a future when return_type="async" async def sleep_first(fut: asyncio.Future, value: typing.Any = mock.DEFAULT): # Sleep time is based on average request time - await asyncio.sleep(sleep_time) + await asyncio.sleep(mock_nemollm_request_time) fut.set_result(value) def create_future(*args, **kwargs) -> asyncio.Future: diff --git a/tests/benchmarks/test_bench_agents_simple_pipeline.py b/tests/benchmarks/test_bench_agents_simple_pipeline.py index 45d0774a2b..a59b4e594c 100644 --- a/tests/benchmarks/test_bench_agents_simple_pipeline.py +++ b/tests/benchmarks/test_bench_agents_simple_pipeline.py @@ -101,6 +101,8 @@ def _run_pipeline(config: Config, source_dfs: list[cudf.DataFrame], model_name: @mock.patch("langchain.OpenAI._agenerate", autospec=True) # autospec is needed as langchain will inspect the function def test_agents_simple_pipe(mock_openai_agenerate: mock.AsyncMock, mock_serpapi_aresults: mock.AsyncMock, + mock_openai_request_time: float, + mock_serpapi_request_time: float, benchmark: collections.abc.Callable[[collections.abc.Callable], typing.Any], config: Config): os.environ.update({'OPENAI_API_KEY': 'test_api_key', 'SERPAPI_API_KEY': 'test_api_key'}) @@ -141,13 +143,13 @@ async def _mock_openai_agenerate(self, *args, **kwargs): # pylint: disable=unus # The OpenAI object will raise a ValueError if we attempt to set the attribute directly or use setattr self.__dict__['_unittest_call_count'] = call_count + 1 - await asyncio.sleep(1.265) + await asyncio.sleep(mock_openai_request_time) return response mock_openai_agenerate.side_effect = _mock_openai_agenerate async def _mock_serpapi_aresults(*args, **kwargs): # pylint: disable=unused-argument - await asyncio.sleep(1.7) + await asyncio.sleep(mock_serpapi_request_time) return { 'answer_box': { 'answer': '25 years', 'link': 'http://unit.test', 'people_also_search_for': [] diff --git a/tests/benchmarks/test_bench_vdb_upload_pipeline.py b/tests/benchmarks/test_bench_vdb_upload_pipeline.py index affd8a91e3..9ebbe0de80 100644 --- a/tests/benchmarks/test_bench_vdb_upload_pipeline.py +++ b/tests/benchmarks/test_bench_vdb_upload_pipeline.py @@ -98,6 +98,8 @@ def _run_pipeline(config: Config, @mock.patch('requests.Session') def test_vdb_upload_pipe(mock_requests_session: mock.MagicMock, mock_feedparser_http_get: mock.MagicMock, + mock_web_scraper_request_time: float, + mock_feedparser_request_time: float, benchmark: collections.abc.Callable[[collections.abc.Callable], typing.Any], config: Config, milvus_server_uri: str, @@ -106,8 +108,6 @@ def test_vdb_upload_pipe(mock_requests_session: mock.MagicMock, with open(os.path.join(TEST_DIRS.tests_data_dir, 'service/cisa_web_responses.json'), encoding='utf-8') as fh: web_responses = json.load(fh) - mock_web_scraper_request_time = float(os.environ.get("MOCK_WEB_SCRAPER_REQUEST_TIME", 0.5)) - def mock_get_fn(url: str): mock_response = mock.MagicMock() mock_response.ok = True @@ -119,8 +119,6 @@ def mock_get_fn(url: str): mock_requests_session.return_value = mock_requests_session mock_requests_session.get.side_effect = mock_get_fn - mock_feedparser_request_time = float(os.environ.get("MOCK_FEEDPARSER_REQUEST_TIME", 0.5)) - def mock_feedparser_http_get_fn(*args, **kwargs): # pylint: disable=unused-argument time.sleep(mock_feedparser_request_time) # The RSS Parser expects a bytes string