diff --git a/tests/conftest.py b/tests/conftest.py index 30296e7a6c..de4626c5e6 100755 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1091,7 +1091,7 @@ def mock_openai_fixture(chat_completion): @pytest.mark.usefixtures("openai") @pytest.fixture(name="mock_async_openai") -def mock_async_openai_fixture(chat_completion): +def mock_async_openai_fixture(chat_completion:): with mock.patch("openai.AsyncOpenAI") as mock_async_openai: mock_async_openai_instance = mock.MagicMock() mock_create = mock.AsyncMock() diff --git a/tests/llm/services/test_llm_service_pipe.py b/tests/llm/services/test_llm_service_pipe.py index ad2cb896c2..b0d1be90f0 100644 --- a/tests/llm/services/test_llm_service_pipe.py +++ b/tests/llm/services/test_llm_service_pipe.py @@ -90,11 +90,12 @@ def test_completion_pipe_nemo( @pytest.mark.use_python -def test_completion_pipe_openai(config: Config, - chat_completion, - mock_async_openai: mock.MagicMock, - country_prompts: list[str], - capital_responses: list[str]): +def test_completion_pipe_openai( + config: Config, + chat_completion, + mock_async_openai: mock.MagicMock, + country_prompts: list[str], + capital_responses: list[str]): chat_completions = [] for response in capital_responses: diff --git a/tests/llm/test_rag_standalone_pipe.py b/tests/llm/test_rag_standalone_pipe.py index cc55b5fcba..3c366a1dd8 100644 --- a/tests/llm/test_rag_standalone_pipe.py +++ b/tests/llm/test_rag_standalone_pipe.py @@ -20,7 +20,6 @@ from unittest import mock import pytest -from openai.types.chat.chat_completion import ChatCompletion import cudf @@ -170,14 +169,15 @@ def test_rag_standalone_pipe_nemo( @pytest.mark.use_cudf @pytest.mark.parametrize("repeat_count", [5]) @pytest.mark.import_mod(os.path.join(TEST_DIRS.examples_dir, 'llm/common/utils.py')) -def test_rag_standalone_pipe_openai(config: Config, - chat_completion: ChatCompletion, - mock_openai: mock.MagicMock, - mock_async_openai: mock.MagicMock, - dataset: DatasetManager, - milvus_server_uri: str, - repeat_count: int, - import_mod: types.ModuleType): +def test_rag_standalone_pipe_openai( + config: Config, + chat_completion, + mock_openai: mock.MagicMock, + mock_async_openai: mock.MagicMock, + dataset: DatasetManager, + milvus_server_uri: str, + repeat_count: int, + import_mod: types.ModuleType): chat_completions = [] for _ in range(repeat_count):