diff --git a/samples/snippets/conftest.py b/samples/snippets/conftest.py index e8253bc5a7..81595967ec 100644 --- a/samples/snippets/conftest.py +++ b/samples/snippets/conftest.py @@ -12,9 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Iterator +from typing import Generator, Iterator -from google.cloud import bigquery +from google.cloud import bigquery, storage import pytest import test_utils.prefixer @@ -42,11 +42,27 @@ def bigquery_client() -> bigquery.Client: return bigquery_client +@pytest.fixture(scope="session") +def storage_client(project_id: str) -> storage.Client: + return storage.Client(project=project_id) + + @pytest.fixture(scope="session") def project_id(bigquery_client: bigquery.Client) -> str: return bigquery_client.project +@pytest.fixture(scope="session") +def gcs_bucket(storage_client: storage.Client) -> Generator[str, None, None]: + bucket_name = "bigframes_blob_test_with_data_wipeout" + + yield bucket_name + + bucket = storage_client.get_bucket(bucket_name) + for blob in bucket.list_blobs(): + blob.delete() + + @pytest.fixture(autouse=True) def reset_session() -> None: """An autouse fixture ensuring each sample runs in a fresh session. @@ -78,11 +94,6 @@ def dataset_id_eu(bigquery_client: bigquery.Client, project_id: str) -> Iterator bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True) -@pytest.fixture(scope="session") -def gcs_dst_bucket() -> str: - return "gs://bigframes_blob_test" - - @pytest.fixture def random_model_id( bigquery_client: bigquery.Client, project_id: str, dataset_id: str diff --git a/samples/snippets/multimodal_test.py b/samples/snippets/multimodal_test.py index 087299aa0a..1ea6a3f0a6 100644 --- a/samples/snippets/multimodal_test.py +++ b/samples/snippets/multimodal_test.py @@ -13,9 +13,9 @@ # limitations under the License. -def test_multimodal_dataframe(gcs_dst_bucket: str) -> None: +def test_multimodal_dataframe(gcs_bucket: str) -> None: # destination folder must be in a GCS bucket that the BQ connection service account (default or user provided) has write access to. - dst_bucket = gcs_dst_bucket + dst_bucket = f"gs://{gcs_bucket}" # [START bigquery_dataframes_multimodal_dataframe_create] import bigframes diff --git a/samples/snippets/sessions_and_io_test.py b/samples/snippets/sessions_and_io_test.py index 24290c7279..06f0c4ab3c 100644 --- a/samples/snippets/sessions_and_io_test.py +++ b/samples/snippets/sessions_and_io_test.py @@ -13,10 +13,11 @@ # limitations under the License. -def test_sessions_and_io(project_id: str, dataset_id: str) -> None: +def test_sessions_and_io(project_id: str, dataset_id: str, gcs_bucket: str) -> None: YOUR_PROJECT_ID = project_id YOUR_DATASET_ID = dataset_id YOUR_LOCATION = "us" + YOUR_BUCKET = gcs_bucket # [START bigquery_dataframes_create_and_use_session_instance] import bigframes @@ -139,6 +140,15 @@ def test_sessions_and_io(project_id: str, dataset_id: str) -> None: # [END bigquery_dataframes_read_data_from_csv] assert df is not None + # [START bigquery_dataframes_write_data_to_csv] + import bigframes.pandas as bpd + + df = bpd.DataFrame({"my_col": [1, 2, 3]}) + # Write a dataframe to a CSV file in GCS + df.to_csv(f"gs://{YOUR_BUCKET}/myfile*.csv") + # [END bigquery_dataframes_write_data_to_csv] + assert df is not None + # [START bigquery_dataframes_read_data_from_bigquery_table] import bigframes.pandas as bpd