diff --git a/sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py b/sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py index a038f3760f24..cf9fbe0a828f 100644 --- a/sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py +++ b/sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py @@ -31,15 +31,12 @@ from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports -try: - from google.api_core.exceptions import BadRequest - from apache_beam.transforms.enrichment import Enrichment - from apache_beam.ml.rag.enrichment.bigquery_vector_search import \ - BigQueryVectorSearchEnrichmentHandler - from apache_beam.ml.rag.enrichment.bigquery_vector_search import \ - BigQueryVectorSearchParameters -except ImportError: - raise unittest.SkipTest('BigQuery dependencies not installed') +from google.api_core.exceptions import BadRequest +from apache_beam.transforms.enrichment import Enrichment +from apache_beam.ml.rag.enrichment.bigquery_vector_search import \ + BigQueryVectorSearchEnrichmentHandler +from apache_beam.ml.rag.enrichment.bigquery_vector_search import \ + BigQueryVectorSearchParameters _LOGGER = logging.getLogger(__name__) diff --git a/sdks/python/apache_beam/ml/rag/ingestion/alloydb_it_test.py b/sdks/python/apache_beam/ml/rag/ingestion/alloydb_it_test.py index 38266e2c7ab1..596bca6aee8b 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/alloydb_it_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/alloydb_it_test.py @@ -68,6 +68,7 @@ ]) registry.register_coder(MetadataConflictRow, RowCoder) +_LOGGER = logging.getLogger(__name__) VECTOR_SIZE = 768 @@ -146,12 +147,12 @@ class AlloyDBVectorWriterConfigTest(unittest.TestCase): @classmethod def setUpClass(cls): - cls.host = os.environ.get('ALLOYDB_HOST', '10.0.0.2') + cls.host = os.environ.get('ALLOYDB_HOST', '10.119.0.22') cls.port = os.environ.get('ALLOYDB_PORT', '5432') cls.database = os.environ.get('ALLOYDB_DATABASE', 'postgres') cls.username = os.environ.get('ALLOYDB_USERNAME', 'postgres') if not os.environ.get('ALLOYDB_PASSWORD'): - raise ValueError("ALLOYDB_PASSWORD is not set.") + raise ValueError('ALLOYDB_PASSWORD env not set') cls.password = os.environ.get('ALLOYDB_PASSWORD') # Create unique table name suffix @@ -373,7 +374,7 @@ def custom_row_to_chunk(row): # Extract timestamp from custom_id timestamp = row.custom_id.split('timestamp_')[1] # Extract index from timestamp - i = int(timestamp.split('T wrong')[1][:2]) + i = int(timestamp.split('T')[1][:2]) # Parse embedding vector embedding_list = [ @@ -488,7 +489,7 @@ def custom_row_to_chunk(row): # Verify count count_result = rows | "Count All" >> beam.combiners.Count.Globally() - assert_that(count_result, equal_to([num_records]), label='count_check') + assert_that(count_result, equal_to([10]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') diff --git a/sdks/python/scripts/run_pytest.sh b/sdks/python/scripts/run_pytest.sh index ad35b48972b6..84e752838105 100755 --- a/sdks/python/scripts/run_pytest.sh +++ b/sdks/python/scripts/run_pytest.sh @@ -41,10 +41,10 @@ echo "pytest_args: $pytest_args" echo "posargs: $posargs" # Run with pytest-xdist and without. -pytest -o junit_suite_name=${envname} \ +pytest -v -rs -o junit_suite_name=${envname} \ --junitxml=pytest_${envname}.xml -m 'not no_xdist' -n 6 --import-mode=importlib ${pytest_args} --pyargs ${posargs} status1=$? -pytest -o junit_suite_name=${envname}_no_xdist \ +pytest -v -rs -o junit_suite_name=${envname}_no_xdist \ --junitxml=pytest_${envname}_no_xdist.xml -m 'no_xdist' --import-mode=importlib ${pytest_args} --pyargs ${posargs} status2=$? diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 5c917582f34b..c88b0415df2d 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -117,6 +117,7 @@ commands = # many packages do not support py3.12 # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests. setenv = +passenv = ALLOYDB_PASSWORD extras = test,gcp,dataframe,p312_ml_test commands = # Log tensorflow version for debugging