wehub-resource-sync
122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
# Copyright 2024 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Integration tests for the VertexAISearchClient.
|
|
|
|
This module contains integration tests that interact with the actual
|
|
Vertex AI Search API. These tests require proper configuration of
|
|
environment variables and access to the Vertex AI Search service.
|
|
"""
|
|
|
|
from collections.abc import Generator
|
|
import os
|
|
|
|
import pytest
|
|
from vertex_ai_search_client import VertexAISearchClient, VertexAISearchConfig
|
|
|
|
# Load environment variables
|
|
PROJECT_ID = os.getenv("PROJECT_ID", "your-project")
|
|
LOCATION = os.getenv("LOCATION", "global")
|
|
DATA_STORE_ID = os.getenv("DATA_STORE_ID", "your-data-store")
|
|
ENGINE_DATA_TYPE = os.getenv("ENGINE_DATA_TYPE", "UNSTRUCTURED")
|
|
ENGINE_CHUNK_TYPE = os.getenv("ENGINE_CHUNK_TYPE", "CHUNK")
|
|
SUMMARY_TYPE = os.getenv("SUMMARY_TYPE", "VERTEX_AI_SEARCH")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vertex_ai_search_client() -> Generator[VertexAISearchClient, None, None]:
|
|
"""
|
|
Fixture to create and yield a VertexAISearchClient instance for testing.
|
|
|
|
This fixture creates a VertexAISearchClient instance using the
|
|
environment variables and yields it for use in tests. The client
|
|
is shared across all tests in the module for efficiency.
|
|
|
|
Yields:
|
|
VertexAISearchClient: An instance of the VertexAISearchClient for testing.
|
|
"""
|
|
config = VertexAISearchConfig(
|
|
project_id=PROJECT_ID,
|
|
location=LOCATION,
|
|
data_store_id=DATA_STORE_ID,
|
|
engine_data_type="UNSTRUCTURED",
|
|
engine_chunk_type="DOCUMENT_WITH_EXTRACTIVE_SEGMENTS",
|
|
summary_type="VERTEX_AI_SEARCH",
|
|
)
|
|
client = VertexAISearchClient(config)
|
|
yield client
|
|
|
|
|
|
def test_search_integration(client: VertexAISearchClient) -> None:
|
|
"""
|
|
Test the search functionality of VertexAISearchClient with the actual API.
|
|
|
|
This test performs a search using the VertexAISearchClient and verifies
|
|
that the results have the expected structure and content types.
|
|
|
|
Args:
|
|
vertex_ai_search_client (VertexAISearchClient): The client instance to test.
|
|
"""
|
|
# Perform a search
|
|
query = "test query"
|
|
results = client.search(query)
|
|
|
|
# Check the structure of the results
|
|
assert "simplified_results" in results
|
|
assert isinstance(results["simplified_results"], list)
|
|
|
|
if results["simplified_results"]:
|
|
first_result = results["simplified_results"][0]
|
|
assert "metadata" in first_result
|
|
assert "page_content" in first_result
|
|
|
|
# Check for other expected fields
|
|
assert "total_size" in results
|
|
assert isinstance(results["total_size"], int)
|
|
|
|
if "summary" in results:
|
|
assert "summary_text" in results["summary"]
|
|
|
|
|
|
def test_unstructured_summary() -> None:
|
|
"""
|
|
Test VertexAISearchClient with unstructured data and summary generation.
|
|
|
|
This test creates a new VertexAISearchClient instance with specific
|
|
settings for unstructured data and summary generation, then performs
|
|
a search to verify the results.
|
|
"""
|
|
config = VertexAISearchConfig(
|
|
project_id=PROJECT_ID,
|
|
location=LOCATION,
|
|
data_store_id=DATA_STORE_ID,
|
|
engine_data_type="UNSTRUCTURED",
|
|
engine_chunk_type="DOCUMENT_WITH_EXTRACTIVE_SEGMENTS",
|
|
summary_type="VERTEX_AI_SEARCH",
|
|
)
|
|
client = VertexAISearchClient(config)
|
|
results = client.search("What is the name of the company?")
|
|
# Check the structure of the results
|
|
assert "simplified_results" in results
|
|
assert isinstance(results["simplified_results"], list)
|
|
|
|
if results["simplified_results"]:
|
|
first_result = results["simplified_results"][0]
|
|
assert "metadata" in first_result
|
|
assert "page_content" in first_result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|