Spaces:
Sleeping
Sleeping
File size: 2,280 Bytes
431a8be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import pandas as pd
import pytest
import requests
from datapilot.analyst import (
ai_context,
build_data_dictionary,
gemini_dataset_summary,
inspect_dataset,
)
def test_immediate_profile_and_target_ranking():
frame = pd.DataFrame(
{
"customer_id": [1, 2, 3, 4],
"age": [24, None, 39, 41],
"churn": ["no", "no", "yes", "yes"],
}
)
profile = inspect_dataset(frame)
assert profile["brief"].rows == 4
assert profile["brief"].missing_cells == 1
assert profile["targets"][0]["column"] == "churn"
def test_dictionary_and_ai_context_redact_pii():
frame = pd.DataFrame(
{"email": ["a@x.com", "b@x.com"], "revenue": [99.0, 101.0], "secret": ["x", "y"]}
)
profile = inspect_dataset(frame)
dictionary = build_data_dictionary(frame).set_index("column")
context = ai_context(frame, profile, ["secret"])
assert "Potential PII" in dictionary.loc["email", "issues"]
assert context["sample"] == [{"revenue": 99.0}, {"revenue": 101.0}]
assert set(context["excluded_columns"]) == {"email", "secret"}
def test_gemini_requires_a_key():
frame = pd.DataFrame({"x": [1, 2], "target": [0, 1]})
with pytest.raises(ValueError, match="API key"):
gemini_dataset_summary(frame, inspect_dataset(frame), "", "gemini", [])
def test_gemini_rest_success(monkeypatch):
frame = pd.DataFrame({"x": [1, 2], "target": [0, 1]})
class Response:
status_code = 200
ok = True
@staticmethod
def json():
return {"candidates": [{"content": {"parts": [{"text": "## Finding\nGrounded"}]}}]}
monkeypatch.setattr(requests, "post", lambda *args, **kwargs: Response())
result = gemini_dataset_summary(
frame, inspect_dataset(frame), "test-key", "gemini-2.5-flash", []
)
assert "Grounded" in result
def test_gemini_rest_quota_error(monkeypatch):
frame = pd.DataFrame({"x": [1, 2], "target": [0, 1]})
class Response:
status_code = 429
ok = False
monkeypatch.setattr(requests, "post", lambda *args, **kwargs: Response())
with pytest.raises(ValueError, match="quota"):
gemini_dataset_summary(frame, inspect_dataset(frame), "test-key", "gemini-2.5-flash", [])
|