Mocked sleep times are now fixtures which can be overriden with an en…

…virontment variable

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:

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': []

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