fix ci

Author: daniel-sanche

tests/system/data/test_system_async.py

@@ -1185,7 +1185,7 @@ async def test_execute_query_simple(self, client, table_id, instance_id):
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
     async def test_execute_against_target(
-        self, client, instance_id, table_id, temp_rows
+        self, client, instance_id, table_id, temp_rows, column_family_config
     ):
         await temp_rows.add_row(b"row_key_1")
         result = await client.execute_query(
@@ -1200,14 +1200,19 @@ async def test_execute_against_target(
         assert family_map[b"q"] == b"test-value"
         assert len(rows[0][TEST_FAMILY_2]) == 0
         md = result.metadata
-        assert len(md) == 3
+        # we expect it to fetch each column family, plus _key
+        # add additional families here if column_family_config changes
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )
 
     @pytest.mark.skipif(
         bool(os.environ.get(BIGTABLE_EMULATOR)),
@@ -1310,7 +1315,7 @@ async def test_execute_query_params(self, client, table_id, instance_id):
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
     async def test_execute_metadata_on_empty_response(
-        self, client, instance_id, table_id, temp_rows
+        self, client, instance_id, table_id, temp_rows, column_family_config
     ):
         await temp_rows.add_row(b"row_key_1")
         result = await client.execute_query(
@@ -1320,11 +1325,16 @@ async def test_execute_metadata_on_empty_response(
 
         assert len(rows) == 0
         md = result.metadata
-        assert len(md) == 3
+        # we expect it to fetch each column family, plus _key
+        # add additional families here if column_family_config change
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )

tests/system/data/test_system_autogen.py

@@ -966,7 +966,9 @@ def test_execute_query_simple(self, client, table_id, instance_id):
     @CrossSync._Sync_Impl.Retry(
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
-    def test_execute_against_target(self, client, instance_id, table_id, temp_rows):
+    def test_execute_against_target(
+        self, client, instance_id, table_id, temp_rows, column_family_config
+    ):
         temp_rows.add_row(b"row_key_1")
         result = client.execute_query("SELECT * FROM `" + table_id + "`", instance_id)
         rows = [r for r in result]
@@ -977,14 +979,17 @@ def test_execute_against_target(self, client, instance_id, table_id, temp_rows):
         assert family_map[b"q"] == b"test-value"
         assert len(rows[0][TEST_FAMILY_2]) == 0
         md = result.metadata
-        assert len(md) == 3
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )
 
     @pytest.mark.skipif(
         bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't support SQL"
@@ -1074,7 +1079,7 @@ def test_execute_query_params(self, client, table_id, instance_id):
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
     def test_execute_metadata_on_empty_response(
-        self, client, instance_id, table_id, temp_rows
+        self, client, instance_id, table_id, temp_rows, column_family_config
     ):
         temp_rows.add_row(b"row_key_1")
         result = client.execute_query(
@@ -1083,11 +1088,14 @@ def test_execute_metadata_on_empty_response(
         rows = [r for r in result]
         assert len(rows) == 0
         md = result.metadata
-        assert len(md) == 3
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )

tests/unit/data/test_mutations.py

@@ -117,6 +117,17 @@ def test_size(self, test_dict):
                 {"delete_from_family": {"family_name": "foo"}},
             ),
             (mutations.DeleteAllFromRow, {"delete_from_row": {}}),
+            (
+                mutations.AddToCell,
+                {
+                    "add_to_cell": {
+                        "family_name": "foo",
+                        "column_qualifier": {"raw_value": b"bar"},
+                        "timestamp": {"raw_timestamp_micros": 12345},
+                        "input": {"int_value": 123},
+                    }
+                },
+            ),
         ],
     )
     def test__from_dict(self, expected_class, input_dict):
@@ -162,6 +173,7 @@ def test__from_dict_wrong_subclass(self):
             mutations.DeleteRangeFromColumn("foo", b"bar"),
             mutations.DeleteAllFromFamily("foo"),
             mutations.DeleteAllFromRow(),
+            mutations.AddToCell("foo", b"bar", 123, 456),
         ]
         for instance in subclasses:
             others = [other for other in subclasses if other != instance]