unstructured/test_unstructured_ingest/python/test-gcs-output.py

#!/usr/bin/env python
import click
from google.cloud import storage
from google.oauth2 import service_account


@click.group(name="gcs-ingest")
def cli():
    pass


@cli.command()
@click.option("--service-account-file", type=click.Path(), required=True)
@click.option("--bucket", type=str, required=True)
@click.option("--blob-path", type=str, required=True)
def down(service_account_file: str, bucket: str, blob_path: str):
    credentials = service_account.Credentials.from_service_account_file(
        filename=service_account_file
    )

    storage_client = storage.Client(credentials=credentials)
    for blob in storage_client.list_blobs(bucket_or_name=bucket, prefix=blob_path):
        print(f"deleting {blob.name}")
        blob.delete()


@cli.command()
@click.option("--service-account-file", type=click.Path(), required=True)
@click.option("--bucket", type=str, required=True)
@click.option("--blob-path", type=str, required=True)
@click.option("--expected-files", type=int, required=True)
def check(service_account_file: str, bucket: str, blob_path: str, expected_files: int):
    credentials = service_account.Credentials.from_service_account_file(
        filename=service_account_file
    )

    storage_client = storage.Client(credentials=credentials)
    blob_json_list = [
        f.name
        for f in storage_client.list_blobs(bucket_or_name=bucket, prefix=blob_path)
        if f.name.endswith("json")
    ]
    found = len(blob_json_list)
    print(
        f"Checking that the number of files found ({found}) "
        f"matches what's expected: {expected_files}"
    )
    assert (
        found == expected_files
    ), f"number of files found ({found}) doesn't match what's expected: {expected_files}"
    print("successfully checked the number of files!")


if __name__ == "__main__":
    cli()