fix unit tests

Author: tswast

bigframes/dataframe.py

@@ -4630,24 +4630,24 @@ def to_string(
     ) -> str | None:
         return self.to_pandas(allow_large_results=allow_large_results).to_string(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,  # type: ignore
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            line_width,
-            min_rows,
-            max_colwidth,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,  # type: ignore
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            line_width=line_width,
+            min_rows=min_rows,
+            max_colwidth=max_colwidth,
+            encoding=encoding,
         )
 
     def to_html(
@@ -4680,28 +4680,28 @@ def to_html(
     ) -> str:
         return self.to_pandas(allow_large_results=allow_large_results).to_html(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,  # type: ignore
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            bold_rows,
-            classes,
-            escape,
-            notebook,
-            border,
-            table_id,
-            render_links,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,  # type: ignore
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            bold_rows=bold_rows,
+            classes=classes,
+            escape=escape,
+            notebook=notebook,
+            border=border,
+            table_id=table_id,
+            render_links=render_links,
+            encoding=encoding,
         )
 
     def to_markdown(
@@ -4713,7 +4713,7 @@ def to_markdown(
         allow_large_results: Optional[bool] = None,
         **kwargs,
     ) -> str | None:
-        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode, index, **kwargs)  # type: ignore
+        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode=mode, index=index, **kwargs)  # type: ignore
 
     def to_pickle(self, path, *, allow_large_results=None, **kwargs) -> None:
         return self.to_pandas(allow_large_results=allow_large_results).to_pickle(

bigframes/testing/mocks.py

@@ -143,6 +143,7 @@ def query_and_wait_mock(query, *args, job_config=None, **kwargs):
 
     bqclient.query.side_effect = query_mock
     bqclient.query_and_wait.side_effect = query_and_wait_mock
+    bqclient._query_and_wait_bigframes.side_effect = query_and_wait_mock
 
     clients_provider = mock.create_autospec(bigframes.session.clients.ClientsProvider)
     type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient)

tests/unit/session/test_session.py

@@ -247,7 +247,7 @@ def test_read_gbq_cached_table():
         table,
     )
 
-    session.bqclient.query_and_wait = mock.MagicMock(
+    session.bqclient._query_and_wait_bigframes = mock.MagicMock(
         return_value=({"total_count": 3, "distinct_count": 2},)
     )
     session.bqclient.get_table.return_value = table
@@ -278,7 +278,7 @@ def test_read_gbq_cached_table_doesnt_warn_for_anonymous_tables_and_doesnt_inclu
         table,
     )
 
-    session.bqclient.query_and_wait = mock.MagicMock(
+    session.bqclient._query_and_wait_bigframes = mock.MagicMock(
         return_value=({"total_count": 3, "distinct_count": 2},)
     )
     session.bqclient.get_table.return_value = table
@@ -306,7 +306,9 @@ def test_default_index_warning_raised_by_read_gbq(table):
     bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
     bqclient.project = "test-project"
     bqclient.get_table.return_value = table
-    bqclient.query_and_wait.return_value = ({"total_count": 3, "distinct_count": 2},)
+    bqclient._query_and_wait_bigframes.return_value = (
+        {"total_count": 3, "distinct_count": 2},
+    )
     session = mocks.create_bigquery_session(
         bqclient=bqclient,
         # DefaultIndexWarning is only relevant for strict mode.
@@ -333,7 +335,9 @@ def test_default_index_warning_not_raised_by_read_gbq_index_col_sequential_int64
     bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
     bqclient.project = "test-project"
     bqclient.get_table.return_value = table
-    bqclient.query_and_wait.return_value = ({"total_count": 4, "distinct_count": 3},)
+    bqclient._query_and_wait_bigframes.return_value = (
+        {"total_count": 4, "distinct_count": 3},
+    )
     session = mocks.create_bigquery_session(
         bqclient=bqclient,
         # DefaultIndexWarning is only relevant for strict mode.
@@ -382,7 +386,7 @@ def test_default_index_warning_not_raised_by_read_gbq_index_col_columns(
     bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
     bqclient.project = "test-project"
     bqclient.get_table.return_value = table
-    bqclient.query_and_wait.return_value = (
+    bqclient._query_and_wait_bigframes.return_value = (
         {"total_count": total_count, "distinct_count": distinct_count},
     )
     session = mocks.create_bigquery_session(
@@ -492,6 +496,7 @@ def query_mock(query, *args, **kwargs):
         return session_query_mock(query, *args, **kwargs)
 
     session.bqclient.query_and_wait = query_mock
+    session.bqclient._query_and_wait_bigframes = query_mock
 
     def get_table_mock(table_ref):
         table = google.cloud.bigquery.Table(