From d7ce449134d096663bf49ad97450d9975904c156 Mon Sep 17 00:00:00 2001
From: Chelsea Lin <chelsealin@google.com>
Date: Tue, 19 Mar 2024 21:28:09 +0000
Subject: [PATCH] fix: renable to_csv and to_json related tests

---
 tests/system/small/test_dataframe_io.py | 12 ++++++++----
 tests/system/small/test_series.py       |  6 ++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/tests/system/small/test_dataframe_io.py b/tests/system/small/test_dataframe_io.py
index adc729565e..afb97c3595 100644
--- a/tests/system/small/test_dataframe_io.py
+++ b/tests/system/small/test_dataframe_io.py
@@ -115,7 +115,6 @@ def test_to_pandas_batches_w_correct_dtypes(scalars_df_default_index):
         pd.testing.assert_series_equal(actual, expected)
 
 
-@pytest.mark.skip(reason="Disable to unblock kokoro tests")
 @pytest.mark.parametrize(
     ("index"),
     [True, False],
@@ -154,6 +153,7 @@ def test_to_csv_index(
         dtype=dtype,
         date_format={"timestamp_col": "YYYY-MM-DD HH:MM:SS Z"},
         index_col=index_col,
+        storage_options=dict(expand=True),
     )
     convert_pandas_dtypes(gcs_df, bytes_col=True)
     gcs_df.index.name = scalars_df.index.name
@@ -164,7 +164,6 @@ def test_to_csv_index(
     pd.testing.assert_frame_equal(gcs_df, scalars_pandas_df)
 
 
-@pytest.mark.skip(reason="Disable to unblock kokoro tests")
 def test_to_csv_tabs(
     scalars_dfs: Tuple[bigframes.dataframe.DataFrame, pd.DataFrame],
     gcs_folder: str,
@@ -194,6 +193,7 @@ def test_to_csv_tabs(
         dtype=dtype,
         date_format={"timestamp_col": "YYYY-MM-DD HH:MM:SS Z"},
         index_col=index_col,
+        storage_options=dict(expand=True),
     )
     convert_pandas_dtypes(gcs_df, bytes_col=True)
     gcs_df.index.name = scalars_df.index.name
@@ -415,7 +415,6 @@ def test_to_json_index_invalid_lines(
         scalars_df.to_json(path, index=index)
 
 
-@pytest.mark.skip(reason="Disable to unblock kokoro tests")
 @pytest.mark.parametrize(
     ("index"),
     [True, False],
@@ -435,7 +434,12 @@ def test_to_json_index_records_orient(
     """ Test the `to_json` API with `orient` is `records` and `lines` is True"""
     scalars_df.to_json(path, index=index, orient="records", lines=True)
 
-    gcs_df = pd.read_json(path, lines=True, convert_dates=["datetime_col"])
+    gcs_df = pd.read_json(
+        path,
+        lines=True,
+        convert_dates=["datetime_col"],
+        storage_options=dict(expand=True),
+    )
     convert_pandas_dtypes(gcs_df, bytes_col=True)
     if index and scalars_df.index.name is not None:
         gcs_df = gcs_df.set_index(scalars_df.index.name)
diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py
index f63ea977ff..3c593a75bf 100644
--- a/tests/system/small/test_series.py
+++ b/tests/system/small/test_series.py
@@ -2390,11 +2390,10 @@ def test_to_frame(scalars_dfs):
     assert_pandas_df_equal(bf_result, pd_result)
 
 
-@pytest.mark.skip(reason="Disable to unblock kokoro tests")
 def test_to_json(gcs_folder, scalars_df_index, scalars_pandas_df_index):
     path = gcs_folder + "test_series_to_json*.jsonl"
     scalars_df_index["int64_col"].to_json(path, lines=True, orient="records")
-    gcs_df = pd.read_json(path, lines=True)
+    gcs_df = pd.read_json(path, lines=True, storage_options=dict(expand=True))
 
     pd.testing.assert_series_equal(
         gcs_df["int64_col"].astype(pd.Int64Dtype()),
@@ -2404,11 +2403,10 @@ def test_to_json(gcs_folder, scalars_df_index, scalars_pandas_df_index):
     )
 
 
-@pytest.mark.skip(reason="Disable to unblock kokoro tests")
 def test_to_csv(gcs_folder, scalars_df_index, scalars_pandas_df_index):
     path = gcs_folder + "test_series_to_csv*.csv"
     scalars_df_index["int64_col"].to_csv(path)
-    gcs_df = pd.read_csv(path)
+    gcs_df = pd.read_csv(path, storage_options=dict(expand=True))
 
     pd.testing.assert_series_equal(
         gcs_df["int64_col"].astype(pd.Int64Dtype()),