From 75d0d8ca7111670bd55a3a8a4fe6eced0206e887 Mon Sep 17 00:00:00 2001 From: Sergio Vidiella Pinto <61970661+svidiella@users.noreply.github.com> Date: Wed, 28 Feb 2024 20:24:33 +0100 Subject: [PATCH] tests: add pytest flags (#38) * tests: add pytest flags --------- Co-authored-by: Leonid Kuligin --- libs/vertexai/Makefile | 2 +- libs/vertexai/tests/conftest.py | 64 +++++++++++++++++++ .../tests/integration_tests/test_callbacks.py | 5 ++ .../integration_tests/test_chat_models.py | 16 +++++ .../integration_tests/test_embeddings.py | 6 ++ .../tests/integration_tests/test_gemma.py | 8 +-- .../tests/integration_tests/test_llms.py | 10 +++ .../integration_tests/test_llms_safety.py | 4 +- .../tests/integration_tests/test_tools.py | 5 +- .../integration_tests/test_vectorstores.py | 12 ++-- .../integration_tests/test_vision_models.py | 4 ++ 11 files changed, 121 insertions(+), 15 deletions(-) create mode 100644 libs/vertexai/tests/conftest.py diff --git a/libs/vertexai/Makefile b/libs/vertexai/Makefile index 9ee09b32..6c7e948d 100644 --- a/libs/vertexai/Makefile +++ b/libs/vertexai/Makefile @@ -9,7 +9,7 @@ TEST_FILE ?= tests/unit_tests/ integration_test integration_tests: TEST_FILE = tests/integration_tests/ test tests integration_test integration_tests: - poetry run pytest $(TEST_FILE) + poetry run pytest --release $(TEST_FILE) ###################### diff --git a/libs/vertexai/tests/conftest.py b/libs/vertexai/tests/conftest.py new file mode 100644 index 00000000..c61fab47 --- /dev/null +++ b/libs/vertexai/tests/conftest.py @@ -0,0 +1,64 @@ +""" +Tests configuration to be executed before tests execution. +""" + +from typing import List + +import pytest + +_RELEASE_FLAG = "release" +_GPU_FLAG = "gpu" +_LONG_FLAG = "long" +_EXTENDED_FLAG = "extended" + +_PYTEST_FLAGS = [_RELEASE_FLAG, _GPU_FLAG, _LONG_FLAG, _EXTENDED_FLAG] + + +def pytest_addoption(parser: pytest.Parser) -> None: + """ + Add flags accepted by pytest CLI. + + Args: + parser: The pytest parser object. + + Returns: + + """ + for flag in _PYTEST_FLAGS: + parser.addoption( + f"--{flag}", action="store_true", default=False, help=f"run {flag} tests" + ) + + +def pytest_configure(config: pytest.Config) -> None: + """ + Add pytest custom configuration. + + Args: + config: The pytest config object. + + Returns: + """ + for flag in _PYTEST_FLAGS: + config.addinivalue_line( + "markers", f"{flag}: mark test to run as {flag} only test" + ) + + +def pytest_collection_modifyitems( + config: pytest.Config, items: List[pytest.Item] +) -> None: + """ + Skip tests with a marker from our list that were not explicitly invoked. + + Args: + config: The pytest config object. + items: The list of tests to be executed. + + Returns: + """ + for item in items: + keywords = list(set(item.keywords).intersection(_PYTEST_FLAGS)) + if keywords and not any((config.getoption(f"--{kw}") for kw in keywords)): + skip = pytest.mark.skip(reason=f"need --{keywords[0]} option to run") + item.add_marker(skip) diff --git a/libs/vertexai/tests/integration_tests/test_callbacks.py b/libs/vertexai/tests/integration_tests/test_callbacks.py index f0b9e84c..910cb156 100644 --- a/libs/vertexai/tests/integration_tests/test_callbacks.py +++ b/libs/vertexai/tests/integration_tests/test_callbacks.py @@ -6,6 +6,7 @@ from langchain_google_vertexai.llms import VertexAI +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "text-bison@001", "code-bison@001"], @@ -25,6 +26,7 @@ def test_llm_invoke(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "chat-bison@001", "codechat-bison@001"], @@ -45,6 +47,7 @@ def test_chat_call(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "text-bison@001", "code-bison@001"], @@ -64,6 +67,7 @@ def test_invoke_config(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release def test_llm_stream() -> None: vb = VertexAICallbackHandler() llm = VertexAI(model_name="gemini-pro", temperature=0.0, callbacks=[vb]) @@ -74,6 +78,7 @@ def test_llm_stream() -> None: assert vb.completion_tokens > 0 +@pytest.mark.release def test_chat_stream() -> None: vb = VertexAICallbackHandler() llm = ChatVertexAI(model_name="gemini-pro", temperature=0.0, callbacks=[vb]) diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py index 365c7a21..0357bbb0 100644 --- a/libs/vertexai/tests/integration_tests/test_chat_models.py +++ b/libs/vertexai/tests/integration_tests/test_chat_models.py @@ -18,6 +18,7 @@ model_names_to_test = [None, "codechat-bison", "chat-bison", "gemini-pro"] +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_initialization(model_name: Optional[str]) -> None: """Test chat model initialization.""" @@ -32,6 +33,7 @@ def test_initialization(model_name: Optional[str]) -> None: assert model.model_name == model.client._model_name.split("/")[-1] +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_vertexai_single_call(model_name: Optional[str]) -> None: if model_name: @@ -44,6 +46,7 @@ def test_vertexai_single_call(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release # mark xfail because Vertex API randomly doesn't respect # the n/candidate_count parameter @pytest.mark.xfail @@ -56,6 +59,7 @@ def test_candidates() -> None: assert len(response.generations[0]) == 2 +@pytest.mark.release @pytest.mark.parametrize("model_name", ["chat-bison@001", "gemini-pro"]) async def test_vertexai_agenerate(model_name: str) -> None: model = ChatVertexAI(temperature=0, model_name=model_name) @@ -76,6 +80,7 @@ async def test_vertexai_agenerate(model_name: str) -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.parametrize("model_name", ["chat-bison@001", "gemini-pro"]) def test_vertexai_stream(model_name: str) -> None: model = ChatVertexAI(temperature=0, model_name=model_name) @@ -86,6 +91,7 @@ def test_vertexai_stream(model_name: str) -> None: assert isinstance(chunk, AIMessageChunk) +@pytest.mark.release async def test_vertexai_astream() -> None: model = ChatVertexAI(temperature=0, model_name="gemini-pro") message = HumanMessage(content="Hello") @@ -94,6 +100,7 @@ async def test_vertexai_astream() -> None: assert isinstance(chunk, AIMessageChunk) +@pytest.mark.release def test_vertexai_single_call_with_context() -> None: model = ChatVertexAI() raw_context = ( @@ -110,6 +117,7 @@ def test_vertexai_single_call_with_context() -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_multimodal() -> None: llm = ChatVertexAI(model_name="gemini-pro-vision") gcs_url = ( @@ -129,6 +137,7 @@ def test_multimodal() -> None: assert isinstance(output.content, str) +@pytest.mark.release @pytest.mark.xfail(reason="problem on vertex side") def test_multimodal_history() -> None: llm = ChatVertexAI(model_name="gemini-pro-vision") @@ -158,6 +167,7 @@ def test_multimodal_history() -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_vertexai_single_call_with_examples() -> None: model = ChatVertexAI() raw_context = "My name is Peter. You are my personal assistant." @@ -172,6 +182,7 @@ def test_vertexai_single_call_with_examples() -> None: assert isinstance(response.content, str) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_vertexai_single_call_with_history(model_name: Optional[str]) -> None: if model_name: @@ -188,6 +199,7 @@ def test_vertexai_single_call_with_history(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_vertexai_single_call_fails_no_message() -> None: chat = ChatVertexAI() with pytest.raises(ValueError) as exc_info: @@ -198,6 +210,7 @@ def test_vertexai_single_call_fails_no_message() -> None: ) +@pytest.mark.release @pytest.mark.parametrize("model_name", ["gemini-pro"]) def test_chat_vertexai_gemini_system_message_error(model_name: str) -> None: model = ChatVertexAI(model_name=model_name) @@ -211,6 +224,7 @@ def test_chat_vertexai_gemini_system_message_error(model_name: str) -> None: model([system_message, message1, message2, message3]) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_chat_vertexai_system_message(model_name: Optional[str]) -> None: if model_name: @@ -231,6 +245,7 @@ def test_chat_vertexai_system_message(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_get_num_tokens_from_messages(model_name: str) -> None: if model_name: @@ -243,6 +258,7 @@ def test_get_num_tokens_from_messages(model_name: str) -> None: assert token == 3 +@pytest.mark.release def test_chat_vertexai_gemini_function_calling() -> None: class MyModel(BaseModel): name: str diff --git a/libs/vertexai/tests/integration_tests/test_embeddings.py b/libs/vertexai/tests/integration_tests/test_embeddings.py index 42495629..6772274a 100644 --- a/libs/vertexai/tests/integration_tests/test_embeddings.py +++ b/libs/vertexai/tests/integration_tests/test_embeddings.py @@ -8,11 +8,13 @@ from langchain_google_vertexai.embeddings import VertexAIEmbeddings +@pytest.mark.release def test_initialization() -> None: """Test embedding model initialization.""" VertexAIEmbeddings() +@pytest.mark.release def test_langchain_google_vertexai_embedding_documents() -> None: documents = ["foo bar"] model = VertexAIEmbeddings() @@ -23,6 +25,7 @@ def test_langchain_google_vertexai_embedding_documents() -> None: assert model.model_name == "textembedding-gecko@001" +@pytest.mark.release def test_langchain_google_vertexai_embedding_query() -> None: document = "foo bar" model = VertexAIEmbeddings() @@ -30,6 +33,7 @@ def test_langchain_google_vertexai_embedding_query() -> None: assert len(output) == 768 +@pytest.mark.release def test_langchain_google_vertexai_large_batches() -> None: documents = ["foo bar" for _ in range(0, 251)] model_uscentral1 = VertexAIEmbeddings(location="us-central1") @@ -40,6 +44,7 @@ def test_langchain_google_vertexai_large_batches() -> None: assert model_asianortheast1.instance["batch_size"] < 50 +@pytest.mark.release def test_langchain_google_vertexai_paginated_texts() -> None: documents = [ "foo bar", @@ -58,6 +63,7 @@ def test_langchain_google_vertexai_paginated_texts() -> None: assert model.model_name == model.client._model_id +@pytest.mark.release def test_warning(caplog: pytest.LogCaptureFixture) -> None: _ = VertexAIEmbeddings() assert len(caplog.records) == 1 diff --git a/libs/vertexai/tests/integration_tests/test_gemma.py b/libs/vertexai/tests/integration_tests/test_gemma.py index cb8705c9..d930a9db 100644 --- a/libs/vertexai/tests/integration_tests/test_gemma.py +++ b/libs/vertexai/tests/integration_tests/test_gemma.py @@ -14,7 +14,7 @@ ) -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_gemma_model_garden() -> None: """In order to run this test, you should provide endpoint names. @@ -36,7 +36,7 @@ def test_gemma_model_garden() -> None: assert llm._llm_type == "gemma_vertexai_model_garden" -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_gemma_chat_model_garden() -> None: """In order to run this test, you should provide endpoint names. @@ -67,7 +67,7 @@ def test_gemma_chat_model_garden() -> None: assert len(output.content) > 2 -@pytest.mark.skip("CI testing not set up") +@pytest.mark.gpu def test_gemma_kaggle() -> None: llm = GemmaLocalKaggle(model_name="gemma_2b_en") output = llm.invoke("What is the meaning of life?") @@ -76,7 +76,7 @@ def test_gemma_kaggle() -> None: assert len(output) > 2 -@pytest.mark.skip("CI testing not set up") +@pytest.mark.gpu def test_gemma_chat_kaggle() -> None: llm = GemmaChatLocalKaggle(model_name="gemma_2b_en") text_question1, text_answer1 = "How much is 2+2?", "4" diff --git a/libs/vertexai/tests/integration_tests/test_llms.py b/libs/vertexai/tests/integration_tests/test_llms.py index f9525c8b..cd5e0897 100644 --- a/libs/vertexai/tests/integration_tests/test_llms.py +++ b/libs/vertexai/tests/integration_tests/test_llms.py @@ -13,6 +13,7 @@ model_names_to_test_with_default = [None] + model_names_to_test +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -26,6 +27,7 @@ def test_vertex_initialization(model_name: str) -> None: assert llm.model_name == llm.client._model_name.split("/")[-1] +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -40,6 +42,7 @@ def test_vertex_invoke(model_name: str) -> None: assert isinstance(output, str) +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -58,6 +61,7 @@ def test_vertex_generate(model_name: str) -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates") def test_vertex_generate_multiple_candidates() -> None: llm = VertexAI(temperature=0.3, n=2, model_name="text-bison@001") @@ -67,6 +71,7 @@ def test_vertex_generate_multiple_candidates() -> None: assert len(output.generations[0]) == 2 +@pytest.mark.release @pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates") def test_vertex_generate_code() -> None: llm = VertexAI(temperature=0.3, n=2, model_name="code-bison@001") @@ -79,6 +84,7 @@ def test_vertex_generate_code() -> None: assert int(usage_metadata["candidates_token_count"]) > 1 +@pytest.mark.release async def test_vertex_agenerate() -> None: llm = VertexAI(temperature=0) output = await llm.agenerate(["Please say foo:"]) @@ -88,6 +94,7 @@ async def test_vertex_agenerate() -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -102,6 +109,7 @@ def test_stream(model_name: str) -> None: assert isinstance(token, str) +@pytest.mark.release async def test_vertex_consistency() -> None: llm = VertexAI(temperature=0) output = llm.generate(["Please say foo:"]) @@ -111,12 +119,14 @@ async def test_vertex_consistency() -> None: assert output.generations[0][0].text == async_output.generations[0][0].text +@pytest.mark.release async def test_astream() -> None: llm = VertexAI(temperature=0, model_name="gemini-pro") async for token in llm.astream("I'm Pickle Rick"): assert isinstance(token, str) +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test, diff --git a/libs/vertexai/tests/integration_tests/test_llms_safety.py b/libs/vertexai/tests/integration_tests/test_llms_safety.py index 1b02eb5a..0b3ef041 100644 --- a/libs/vertexai/tests/integration_tests/test_llms_safety.py +++ b/libs/vertexai/tests/integration_tests/test_llms_safety.py @@ -41,7 +41,7 @@ """ -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_gemini_safety_settings_generate() -> None: llm = VertexAI(model_name="gemini-pro", safety_settings=SAFETY_SETTINGS) output = llm.generate(["What do you think about child abuse:"]) @@ -70,7 +70,7 @@ def test_gemini_safety_settings_generate() -> None: assert not generation_info.get("is_blocked") -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended async def test_gemini_safety_settings_agenerate() -> None: llm = VertexAI(model_name="gemini-pro", safety_settings=SAFETY_SETTINGS) output = await llm.agenerate(["What do you think about child abuse:"]) diff --git a/libs/vertexai/tests/integration_tests/test_tools.py b/libs/vertexai/tests/integration_tests/test_tools.py index ef172a1b..9d378d6e 100644 --- a/libs/vertexai/tests/integration_tests/test_tools.py +++ b/libs/vertexai/tests/integration_tests/test_tools.py @@ -44,7 +44,7 @@ def parse(self, text: str) -> Union[AgentAction, AgentFinish]: raise ValueError("Can only parse messages") -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_tools() -> None: from langchain.agents import AgentExecutor from langchain.agents.format_scratchpad import ( @@ -93,7 +93,7 @@ def test_tools() -> None: assert round(float(just_numbers), 2) == 2.16 -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_stream() -> None: from langchain.chains import LLMMathChain @@ -112,6 +112,7 @@ def test_stream() -> None: assert "function_call" in response[0].additional_kwargs +@pytest.mark.release def test_multiple_tools() -> None: from langchain.agents import AgentExecutor from langchain.agents.format_scratchpad import format_to_openai_function_messages diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py index 39face5e..e8fd5b39 100644 --- a/libs/vertexai/tests/integration_tests/test_vectorstores.py +++ b/libs/vertexai/tests/integration_tests/test_vectorstores.py @@ -41,7 +41,7 @@ def sdk_manager() -> VectorSearchSDKManager: return sdk_manager -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_vector_search_sdk_manager(sdk_manager: VectorSearchSDKManager): gcs_client = sdk_manager.get_gcs_client() assert isinstance(gcs_client, storage.Client) @@ -56,7 +56,7 @@ def test_vector_search_sdk_manager(sdk_manager: VectorSearchSDKManager): assert isinstance(endpoint, MatchingEngineIndexEndpoint) -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_gcs_document_storage(sdk_manager: VectorSearchSDKManager): bucket = sdk_manager.get_gcs_bucket(os.environ["GCS_BUCKET_NAME"]) prefix = "integration-test" @@ -80,7 +80,7 @@ def test_gcs_document_storage(sdk_manager: VectorSearchSDKManager): assert original_text == retrieved_text -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_datastore_document_storage(sdk_manager: VectorSearchSDKManager): ds_client = sdk_manager.get_datastore_client(namespace="Foo") @@ -103,7 +103,7 @@ def test_datastore_document_storage(sdk_manager: VectorSearchSDKManager): assert original_text == retrieved_text -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_public_endpoint_vector_searcher(sdk_manager: VectorSearchSDKManager): index = sdk_manager.get_index(os.environ["INDEX_ID"]) endpoint = sdk_manager.get_endpoint(os.environ["ENDPOINT_ID"]) @@ -120,7 +120,7 @@ def test_public_endpoint_vector_searcher(sdk_manager: VectorSearchSDKManager): assert len(matching_neighbors_list) == 2 -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_vector_store(): embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default") @@ -148,7 +148,7 @@ def test_vector_store(): assert isinstance(doc, Document) -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_vector_store_update_index(): embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default") diff --git a/libs/vertexai/tests/integration_tests/test_vision_models.py b/libs/vertexai/tests/integration_tests/test_vision_models.py index a3edaa04..7714373a 100644 --- a/libs/vertexai/tests/integration_tests/test_vision_models.py +++ b/libs/vertexai/tests/integration_tests/test_vision_models.py @@ -10,6 +10,7 @@ ) +@pytest.mark.release def test_vertex_ai_image_captioning_chat(base64_image: str): # This should work model = VertexAIImageCaptioningChat() @@ -43,12 +44,14 @@ def test_vertex_ai_image_captioning_chat(base64_image: str): ) +@pytest.mark.release def test_vertex_ai_image_captioning(base64_image: str): model = VertexAIImageCaptioning() response = model.invoke(base64_image) assert isinstance(response, str) +@pytest.mark.release def test_vertex_ai_visual_qna_chat(base64_image: str): model = VertexAIVisualQnAChat() @@ -112,6 +115,7 @@ def test_vertex_ai_visual_qna_chat(base64_image: str): ) +@pytest.mark.release def test_vertex_ai_image_generation_and_edition(): generator = VertexAIImageGeneratorChat()