bigquery-etl/conftest.py

"""PyTest configuration."""

from google.cloud import bigquery
from google.cloud import storage
import os
import pytest
import random
import string


TEST_BUCKET = "bigquery-etl-integration-test-bucket"


pytest_plugins = [
    "bigquery_etl.pytest_plugin.sql",
    "bigquery_etl.pytest_plugin.udf",
    "bigquery_etl.pytest_plugin.script_lint.black",
    "bigquery_etl.pytest_plugin.script_lint.docstyle",
    "bigquery_etl.pytest_plugin.script_lint.flake8",
    "bigquery_etl.pytest_plugin.script_lint.mypy",
]


def pytest_collection_modifyitems(config, items):
    keywordexpr = config.option.keyword
    markexpr = config.option.markexpr
    if keywordexpr or markexpr:
        return

    skip_integration = pytest.mark.skip(reason="integration marker not selected")

    for item in items:
        if "integration" in item.keywords:
            item.add_marker(skip_integration)


@pytest.fixture
def project_id():
    """Provide a BigQuery project ID."""
    # GOOGLE_PROJECT_ID needs to be set for integration tests to run
    project_id = os.environ["GOOGLE_PROJECT_ID"]

    return project_id


@pytest.fixture
def bigquery_client():
    """Provide a BigQuery client."""
    project_id = os.environ["GOOGLE_PROJECT_ID"]
    return bigquery.Client(project_id)


@pytest.fixture
def temporary_dataset():
    """Fixture for creating a random temporary BigQuery dataset."""
    # generate a random test dataset to avoid conflicts when running tests in parallel
    test_dataset = "test_" + "".join(
        random.choice(string.ascii_lowercase) for i in range(12)
    )

    project_id = os.environ["GOOGLE_PROJECT_ID"]
    client = bigquery.Client(project_id)
    client.create_dataset(test_dataset)

    yield test_dataset

    # cleanup and remove temporary dataset
    client.delete_dataset(test_dataset, delete_contents=True, not_found_ok=True)


@pytest.fixture
def test_bucket():
    """Provide a test bucket instance."""
    storage_client = storage.Client()
    bucket = storage_client.bucket(TEST_BUCKET)

    yield bucket

    # cleanup test bucket
    bucket.delete_blobs(bucket.list_blobs())


@pytest.fixture
def storage_client():
    """Provide a client instance for cloud storage."""
    yield storage.Client()