Test Decorators

from elementary_python_sdk.core.tests import boolean_test
import pandas as pd

@boolean_test(
    name="unique_ids",
    severity="ERROR",
    description="All user IDs must be unique",
    column_name="id",
)
def test_unique_ids(df: pd.DataFrame) -> bool:
    ids = df["id"].dropna().tolist()
    return len(ids) == len(set(ids))

from elementary_python_sdk.core.tests import expected_range
import pandas as pd

@expected_range(
    name="average_age",
    min=18,
    max=50,
    severity="ERROR",
    description="Average age should be between 18 and 50",
    column_name="age",
)
def test_average_age(df: pd.DataFrame) -> float:
    return df["age"].mean()

from elementary_python_sdk.core.tests import row_count
import pandas as pd

@row_count(
    name="user_count_range",
    min=1,
    max=1000000,
    severity="WARNING",
    description="Validate user count is within expected range",
)
def get_users_df(df: pd.DataFrame) -> pd.DataFrame:
    """Return the dataframe - decorator calls len() on it."""
    return df

from elementary_python_sdk.core.tests import expected_values
import pandas as pd

@expected_values(
    name="country_count",
    expected=2,
    severity="ERROR",
    description="Should have exactly 2 countries",
    column_name="country",
)
def count_unique_countries(df: pd.DataFrame) -> int:
    return df["country"].nunique()

from elementary_python_sdk.core.cloud.cloud_client import ElementaryCloudClient
from elementary_python_sdk.core.tests import (
    boolean_test,
    elementary_test_context,
    expected_range,
)
from elementary_python_sdk.core.types.asset import TableAsset
import pandas as pd

@boolean_test(
    name="unique_ids",
    description="All user IDs must be unique",
    severity="ERROR",
)
def test_unique_ids(df: pd.DataFrame) -> bool:
    return len(df["id"]) == len(df["id"].unique())

@expected_range(
    name="average_age",
    min=18,
    max=50,
    severity="ERROR",
)
def test_average_age(df: pd.DataFrame) -> float:
    return df["age"].mean()

# Define your asset
asset = TableAsset(
    name="users",
    database_name="prod",
    schema_name="public",
    table_name="users"
)

# Run tests within context
with elementary_test_context(asset=asset) as ctx:
    users_df = pd.DataFrame({"id": [1, 2, 3], "age": [25, 30, 35]})
    
    # Run tests - results are automatically captured
    test_unique_ids(users_df)
    test_average_age(users_df)
    
    # Send results to Elementary Cloud
    PROJECT_ID = "my-python-project"  # Your Python project identifier (used to deduplicate and identify assets)
    API_KEY = "your-api-key"
    URL = "https://app.elementary-data.com/sdk-ingest/{env_id}/batch"
    
    client = ElementaryCloudClient(PROJECT_ID, API_KEY, URL)
    client.send_to_cloud(ctx)

with elementary_test_context(asset=asset) as ctx:
    # Using boolean_test context manager
    with ctx.boolean_test(name="my_test", description="Test description") as my_bool_test:
        my_bool_test.assert_value(False)  # Assert a boolean value
    
    # Using expected_values_test context manager
    with ctx.expected_values_test(
        name="country_count",
        expected=[2, 3],
        allow_none=True,
        metadata={"my_metadata_field": "my_metadata_value"},
    ) as my_expected_values_test:
        my_expected_values_test.assert_value(5)  # This will fail (not in expected list)
        my_expected_values_test.assert_value(3)  # This will pass
    
    # Using expected_range_test context manager
    with ctx.expected_range_test(
        name="age_range",
        min=18,
        max=50,
    ) as my_range_test:
        my_range_test.assert_value(25.5)  # Assert a numeric value
    
    # Using row_count_test context manager
    with ctx.row_count_test(
        name="row_count",
        min=1,
        max=1000,
    ) as my_row_count_test:
        my_row_count_test.assert_value(users_df)  # Assert a Sized object
    
    PROJECT_ID = "your-project-id"
    API_KEY = "your-api-key"
    URL = "https://app.elementary-data.com/sdk-ingest/{env_id}/batch"
    
    client = ElementaryCloudClient(PROJECT_ID, API_KEY, URL)
    client.send_to_cloud(ctx)

from elementary_python_sdk.core.tests import boolean_test

# Existing Great Expectations test
def ge_test(df):
    # Your Great Expectations code here
    return result

# Wrap it with Elementary decorator
@boolean_test(
    name="ge_wrapped_test",
    description="Great Expectations test",
    severity="ERROR",
)
def elementary_wrapped_test(df):
    return ge_test(df)

@expected_range(
    name="average_age",
    min=18,
    max=50,
    description="Average age should be between 18 and 50",
)
def test_average_age(df: pd.DataFrame) -> float:
    # If an exception is raised, it will be captured and reported as a test error
    if df.empty:
        raise ValueError("DataFrame is empty")
    return df["age"].mean()