test: Add unit tests for SequentialUIDGenerator (#1813)

This commit introduces unit tests for the `get_uid_stream` method in the `SequentialUIDGenerator` class.

Fixes internal issue 416487613

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -167,6 +167,7 @@ def compile_readtable(self, node: nodes.ReadTableNode, *args):
             table.table_id,
             col_names=[col.source_id for col in node.scan_list.items],
             alias_names=[col.id.sql for col in node.scan_list.items],
+            uid_gen=self.uid_gen,
         )
 
     @_compile_node.register

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -114,6 +114,7 @@ def from_table(
         table_id: str,
         col_names: typing.Sequence[str],
         alias_names: typing.Sequence[str],
+        uid_gen: guid.SequentialUIDGenerator,
     ) -> SQLGlotIR:
         selections = [
             sge.Alias(
@@ -128,7 +129,7 @@ def from_table(
             catalog=sg.to_identifier(project_id, quoted=cls.quoted),
         )
         select_expr = sge.Select().select(*selections).from_(table_expr)
-        return cls(expr=select_expr)
+        return cls(expr=select_expr, uid_gen=uid_gen)
 
     @classmethod
     def from_query_string(
@@ -164,10 +165,10 @@ def select(
         squashed_selections = _squash_selections(self.expr.expressions, selections)
         if squashed_selections != []:
             new_expr = self.expr.select(*squashed_selections, append=False)
-            return SQLGlotIR(expr=new_expr)
+            return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
         else:
             new_expr = self._encapsulate_as_cte().select(*selections, append=False)
-            return SQLGlotIR(expr=new_expr)
+            return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def project(
         self,
@@ -181,7 +182,7 @@ def project(
             for id, expr in projected_cols
         ]
         new_expr = self._encapsulate_as_cte().select(*projected_cols_expr, append=False)
-        return SQLGlotIR(expr=new_expr)
+        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def insert(
         self,

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql

@@ -1,4 +1,4 @@
-WITH `bfcte_2` AS (
+WITH `bfcte_0` AS (
   SELECT
     `rowindex` AS `bfcol_0`,
     `int64_col` AS `bfcol_1`,
@@ -13,4 +13,4 @@ SELECT
   `bfcol_2` AS `string_col`,
   `bfcol_3` AS `float64_col`,
   `bfcol_4` AS `bool_col`
-FROM `bfcte_2`
+FROM `bfcte_0`

tests/unit/core/test_guid.py

@@ -0,0 +1,41 @@
+import types
+import unittest
+
+from bigframes.core.guid import SequentialUIDGenerator
+
+
+class TestSequentialUIDGenerator(unittest.TestCase):
+    def test_get_uid_stream_returns_generator(self):
+        generator = SequentialUIDGenerator()
+        stream = generator.get_uid_stream("prefix")
+        self.assertIsInstance(stream, types.GeneratorType)
+
+    def test_generator_yields_correct_uids(self):
+        generator = SequentialUIDGenerator()
+        stream = generator.get_uid_stream("prefix")
+        self.assertEqual(next(stream), "prefix0")
+        self.assertEqual(next(stream), "prefix1")
+        self.assertEqual(next(stream), "prefix2")
+
+    def test_generator_yields_different_uids_for_different_prefixes(self):
+        generator = SequentialUIDGenerator()
+        stream_a = generator.get_uid_stream("prefixA")
+        stream_b = generator.get_uid_stream("prefixB")
+        self.assertEqual(next(stream_a), "prefixA0")
+        self.assertEqual(next(stream_b), "prefixB0")
+        self.assertEqual(next(stream_a), "prefixA1")
+        self.assertEqual(next(stream_b), "prefixB1")
+
+    def test_multiple_calls_continue_generation(self):
+        generator = SequentialUIDGenerator()
+        stream1 = generator.get_uid_stream("prefix")
+        self.assertEqual(next(stream1), "prefix0")
+        self.assertEqual(next(stream1), "prefix1")
+
+        stream2 = generator.get_uid_stream("prefix")
+        self.assertEqual(next(stream2), "prefix2")
+        self.assertEqual(next(stream2), "prefix3")
+
+
+if __name__ == "__main__":
+    unittest.main()