chore: update tests to be compatible with Python 3.14, set 3.14 testing workflow as required

PiperOrigin-RevId: 825273830

Author: speedstorm1

testing/constraints-3.14.txt

@@ -13,4 +13,4 @@ ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
 google-adk==0.0.2
 google-genai>=1.10.0
 google-vizier==0.1.21
-pyarrow>=18.0.0
+pyarrow>=22.0.0

tests/unit/aiplatform/test_prediction.py

@@ -3287,7 +3287,10 @@ def test_health(self, model_server_env_mock, importlib_import_module_mock_twice)
 
         assert response.status_code == 200
 
-    def test_predict(self, model_server_env_mock, importlib_import_module_mock_twice):
+    @pytest.mark.asyncio
+    async def test_predict(
+        self, model_server_env_mock, importlib_import_module_mock_twice
+    ):
         model_server = CprModelServer()
         client = TestClient(model_server.app)
 

tests/unit/vertex_rag/test_rag_data.py

@@ -419,9 +419,10 @@ def create_transformation_config(
 def rag_corpus_eq(returned_corpus, expected_corpus):
     assert returned_corpus.name == expected_corpus.name
     assert returned_corpus.display_name == expected_corpus.display_name
-    assert returned_corpus.backend_config.__eq__(expected_corpus.backend_config)
-    assert returned_corpus.vertex_ai_search_config.__eq__(
-        expected_corpus.vertex_ai_search_config
+    assert returned_corpus.backend_config == expected_corpus.backend_config
+    assert (
+        returned_corpus.vertex_ai_search_config
+        == expected_corpus.vertex_ai_search_config
     )
 
 
@@ -464,8 +465,8 @@ def import_files_request_eq(returned_request, expected_request):
 
 def rag_engine_config_eq(returned_config, expected_config):
     assert returned_config.name == expected_config.name
-    assert returned_config.rag_managed_db_config.__eq__(
-        expected_config.rag_managed_db_config
+    assert (
+        returned_config.rag_managed_db_config == expected_config.rag_managed_db_config
     )
 
 

tests/unit/vertex_rag/test_rag_data_preview.py

@@ -674,12 +674,13 @@ def create_transformation_config(
 def rag_corpus_eq(returned_corpus, expected_corpus):
     assert returned_corpus.name == expected_corpus.name
     assert returned_corpus.display_name == expected_corpus.display_name
-    assert returned_corpus.vector_db.__eq__(expected_corpus.vector_db)
-    assert returned_corpus.backend_config.__eq__(expected_corpus.backend_config)
-    assert returned_corpus.vertex_ai_search_config.__eq__(
-        expected_corpus.vertex_ai_search_config
+    assert returned_corpus.vector_db == expected_corpus.vector_db
+    assert returned_corpus.backend_config == expected_corpus.backend_config
+    assert (
+        returned_corpus.vertex_ai_search_config
+        == expected_corpus.vertex_ai_search_config
     )
-    assert returned_corpus.corpus_type_config.__eq__(expected_corpus.corpus_type_config)
+    assert returned_corpus.corpus_type_config == expected_corpus.corpus_type_config
 
 
 def rag_file_eq(returned_file, expected_file):
@@ -729,8 +730,8 @@ def import_files_request_eq(returned_request, expected_request):
 
 def rag_engine_config_eq(returned_config, expected_config):
     assert returned_config.name == expected_config.name
-    assert returned_config.rag_managed_db_config.__eq__(
-        expected_config.rag_managed_db_config
+    assert (
+        returned_config.rag_managed_db_config == expected_config.rag_managed_db_config
     )
 
 

vertexai/preview/rag/utils/_gapic_utils.py

@@ -185,6 +185,24 @@ def _check_rag_embedding_model_config(
         return gapic_vector_db.rag_embedding_model_config.ByteSize() > 0
 
 
+def _check_document_corpus(
+    gapic_corpus_type_config: GapicRagCorpus.CorpusTypeConfig,
+) -> bool:
+    try:
+        return gapic_corpus_type_config.__contains__("document_corpus")
+    except AttributeError:
+        return gapic_corpus_type_config.document_corpus.ByteSize() > 0
+
+
+def _check_memory_corpus(
+    gapic_corpus_type_config: GapicRagCorpus.CorpusTypeConfig,
+) -> bool:
+    try:
+        return gapic_corpus_type_config.__contains__("memory_corpus")
+    except AttributeError:
+        return gapic_corpus_type_config.memory_corpus.ByteSize() > 0
+
+
 def _convert_gapic_to_rag_managed_db(
     gapic_rag_managed_db: GapicRagVectorDbConfig.RagManagedDb,
 ) -> RagManagedDb:
@@ -252,10 +270,10 @@ def convert_gapic_to_vector_db(
 
 
 def convert_gapic_to_vertex_ai_search_config(
-    gapic_vertex_ai_search_config: VertexAiSearchConfig,
-) -> VertexAiSearchConfig:
+    gapic_vertex_ai_search_config: GapicVertexAiSearchConfig,
+) -> Optional[VertexAiSearchConfig]:
     """Convert Gapic VertexAiSearchConfig to VertexAiSearchConfig."""
-    if gapic_vertex_ai_search_config.serving_config:
+    if gapic_vertex_ai_search_config.ByteSize() > 0:
         return VertexAiSearchConfig(
             serving_config=gapic_vertex_ai_search_config.serving_config,
         )
@@ -293,6 +311,8 @@ def convert_gapic_to_backend_config(
     gapic_vector_db: GapicRagVectorDbConfig,
 ) -> RagVectorDbConfig:
     """Convert Gapic RagVectorDbConfig to VertexVectorSearch, Pinecone, or RagManagedDb."""
+    if not gapic_vector_db or not gapic_vector_db.ByteSize():
+        return None
     vector_config = RagVectorDbConfig()
     if _check_pinecone(gapic_vector_db):
         vector_config.vector_db = Pinecone(
@@ -314,16 +334,21 @@ def convert_gapic_to_backend_config(
                 gapic_vector_db.rag_embedding_model_config
             )
         )
+    if (
+        vector_config.vector_db is None
+        and vector_config.rag_embedding_model_config is None
+    ):
+        return None
     return vector_config
 
 
 def convert_gapic_to_rag_corpus_type_config(
     gapic_rag_corpus_type_config: GapicRagCorpus.CorpusTypeConfig,
 ) -> RagCorpusTypeConfig:
     """Convert GapicRagCorpus.CorpusTypeConfig to RagCorpusTypeConfig."""
-    if gapic_rag_corpus_type_config.document_corpus:
+    if _check_document_corpus(gapic_rag_corpus_type_config):
         return RagCorpusTypeConfig(corpus_type_config=DocumentCorpus())
-    elif gapic_rag_corpus_type_config.memory_corpus:
+    elif _check_memory_corpus(gapic_rag_corpus_type_config):
         return RagCorpusTypeConfig(
             corpus_type_config=MemoryCorpus(
                 llm_parser=LlmParserConfig(

vertexai/rag/utils/_gapic_utils.py

@@ -175,6 +175,8 @@ def convert_gapic_to_backend_config(
     gapic_vector_db: GapicRagVectorDbConfig,
 ) -> RagVectorDbConfig:
     """Convert Gapic RagVectorDbConfig to VertexVectorSearch, Pinecone, or RagManagedDb."""
+    if not gapic_vector_db:
+        return None
     vector_config = RagVectorDbConfig()
     if _check_pinecone(gapic_vector_db):
         vector_config.vector_db = Pinecone(
@@ -198,10 +200,10 @@ def convert_gapic_to_backend_config(
 
 
 def convert_gapic_to_vertex_ai_search_config(
-    gapic_vertex_ai_search_config: VertexAiSearchConfig,
-) -> VertexAiSearchConfig:
+    gapic_vertex_ai_search_config: GapicVertexAiSearchConfig,
+) -> Optional[VertexAiSearchConfig]:
     """Convert Gapic VertexAiSearchConfig to VertexAiSearchConfig."""
-    if gapic_vertex_ai_search_config.serving_config:
+    if gapic_vertex_ai_search_config.ByteSize() > 0:
         return VertexAiSearchConfig(
             serving_config=gapic_vertex_ai_search_config.serving_config,
         )