conftest.py

import os
import uuid

import pytest
import snowflake.connector
from dotenv import load_dotenv, dotenv_values

from snowcap.enums import ResourceType

# Load environment variables from tests/.env if it exists
load_dotenv("tests/.env")

TEST_ROLE = os.environ.get("TEST_SNOWFLAKE_ROLE")
TEST_WAREHOUSE = os.environ.get("TEST_SNOWFLAKE_WAREHOUSE")


def connection_params():
    params = {
        "account": os.environ["TEST_SNOWFLAKE_ACCOUNT"],
        "user": os.environ["TEST_SNOWFLAKE_USER"],
        "role": TEST_ROLE,
    }

    # Key-pair auth (preferred) or password auth
    private_key_path = os.environ.get("TEST_SNOWFLAKE_PRIVATE_KEY_PATH")
    if private_key_path:
        params["private_key_file"] = private_key_path
        passphrase = os.environ.get("TEST_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")
        if passphrase:
            params["private_key_file_pwd"] = passphrase
    else:
        params["password"] = os.environ["TEST_SNOWFLAKE_PASSWORD"]

    # Optional warehouse
    warehouse = os.environ.get("TEST_SNOWFLAKE_WAREHOUSE")
    if warehouse:
        params["warehouse"] = warehouse

    return params


def pytest_addoption(parser):
    parser.addoption(
        "--snowflake",
        action="store_true",
        default=False,
        help="Runs tests that require a Snowflake connection",
    )


def pytest_runtest_setup(item):
    if "requires_snowflake" in item.keywords and not item.config.getoption("--snowflake"):
        pytest.skip("need --snowflake option to run this test")


def pytest_collection_modifyitems(items):
    for item in items:
        if not item.get_closest_marker("enterprise"):
            item.add_marker("standard")


@pytest.fixture(scope="session")
def blueprint_vars():
    if os.path.exists("env/.vars.test_account"):
        vars = dotenv_values("env/.vars.test_account")
        vars.pop("static_user_rsa_public_key", None)
        vars.pop("static_user_mfa_password", None)
        return vars
    else:
        return {key[4:].lower(): value for key, value in os.environ.items() if key.startswith("VAR_")}


@pytest.fixture(scope="session")
def suffix():
    return str(uuid.uuid4())[:8].upper()


@pytest.fixture(scope="session")
def test_db(suffix):
    return f"TEST_DB_RUN_{suffix}"


@pytest.fixture(scope="session")
def marked_for_cleanup() -> list:
    """List to keep track of resources created during tests."""
    return []


def _verify_static_resources(cursor):
    """Verify that required static test resources exist in Snowflake."""
    required_resources = [
        ("SHOW ROLES LIKE 'STATIC_ROLE'", "STATIC_ROLE"),
        ("SHOW DATABASES LIKE 'STATIC_DATABASE'", "STATIC_DATABASE"),
        ("SHOW WAREHOUSES LIKE 'STATIC_WAREHOUSE'", "STATIC_WAREHOUSE"),
    ]
    missing = []
    for query, name in required_resources:
        result = cursor.execute(query).fetchall()
        if not result:
            missing.append(name)

    if missing:
        pytest.exit(
            f"\n\nMissing required static test resources: {', '.join(missing)}\n\n"
            f"Please run 'make setup-test-resources' to create them.\n"
            f"See tests/fixtures/static_resources/README.md for details.\n",
            returncode=1,
        )


@pytest.fixture(scope="session")
def cursor(suffix, test_db, marked_for_cleanup):
    session = snowflake.connector.connect(**connection_params())
    with session.cursor(snowflake.connector.DictCursor) as cur:
        cur.execute(f"ALTER SESSION set query_tag='snowcap_package:test::{suffix}'")
        _verify_static_resources(cur)
        cur.execute(f"CREATE DATABASE {test_db}")
        try:
            yield cur
            cur.execute(f"USE ROLE {TEST_ROLE}")
            cur.execute(f"USE DATABASE {test_db}")

            # Prioritize NetworkPolicy resources for cleanup
            network_policies = [res for res in marked_for_cleanup if res.resource_type == ResourceType.NETWORK_POLICY]
            other_resources = [res for res in marked_for_cleanup if res.resource_type != ResourceType.NETWORK_POLICY]

            for res in network_policies + other_resources:
                try:
                    cur.execute(res.drop_sql(if_exists=True))
                except snowflake.connector.errors.ProgrammingError as err:
                    if err.errno in (2003, 393950):
                        pass
                    else:
                        raise

        finally:
            cur.execute(f"DROP DATABASE {test_db}")


@pytest.fixture(scope="session")
def dummy_cursor(request):
    if not request.config.getoption("--snowflake"):
        yield None
    else:
        yield request.getfixturevalue("cursor")


@pytest.fixture(autouse=True)
def reset_cursor_context(dummy_cursor, test_db):
    """
    This fixture resets the cursor's context to the initial test database before each test.
    It uses `autouse=True` to automatically apply it to each test without needing to explicitly include it.
    """
    cursor = dummy_cursor
    if cursor:
        cursor.execute(f"USE ROLE {TEST_ROLE}")
        if TEST_WAREHOUSE:
            cursor.execute(f"USE WAREHOUSE {TEST_WAREHOUSE}")
        cursor.execute(f"USE DATABASE {test_db}")
    yield