Fixed flake8 errors with openai module

nv-morpheus · Jan 3, 2024 · 211d66f · 211d66f
1 parent e3af515
commit 211d66f
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 15 deletions.
diff --git 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
@@ -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
@@ -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):