Print runner IP.

Author: claudevdm

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')

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