Print runner IP.

apache · Feb 9, 2025 · 2ac8cca · 2ac8cca
1 parent ca6becb
commit 2ac8cca
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
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
@@ -68,6 +68,7 @@
 ])
 registry.register_coder(MetadataConflictRow, RowCoder)
 
+_LOGGER = logging.getLogger(__name__)
 VECTOR_SIZE = 768
 
 
@@ -141,17 +142,18 @@ def key_on_id(chunk):
   return (int(chunk.id.split('_')[1]), chunk)
 
 
+@unittest.skipUnless(
+    os.environ.get('ALLOYDB_PORT'),
+    "ALLOYDB_PORT environment var is not provided")
 class AlloyDBVectorWriterConfigTest(unittest.TestCase):
   ALLOYDB_TABLE_PREFIX = 'python_rag_alloydb_'
 
   @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.")
     cls.password = os.environ.get('ALLOYDB_PASSWORD')
 
     # Create unique table name suffix
@@ -373,7 +375,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 +490,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
@@ -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 -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 -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
@@ -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