data-query-agent/tests/test_models.py

from data_query_agent.models import QueryAnalysis, QueryDecision, SkillParameters


def test_structured_output_schemas_require_every_property() -> None:
    for model in (SkillParameters, QueryDecision, QueryAnalysis):
        schema = model.model_json_schema()
        assert set(schema["required"]) == set(schema["properties"])
        assert schema["additionalProperties"] is False


def test_query_decision_supports_direct_natural_language_reply() -> None:
    decision = QueryDecision(
        status="ready",
        reply="产品效率报表中的 ROV 是回流 UV 与 DAU 的比值。",
        title="数据助手回复",
        selected_skill="odps-product-efficiency-report",
        execution_mode="direct_reply",
        sql=None,
        parameters=SkillParameters(
            user_id=None,
            date=None,
            apptype=None,
            realtime=None,
            app_type=None,
            date_from=None,
            date_to=None,
            data_mode=None,
            bucket_position_from_end=None,
            experiment_buckets=None,
            control_buckets=None,
            version=None,
            first_layer_rule=None,
            exclude_qywx=None,
        ),
        assumptions=[],
    )

    assert decision.execution_mode == "direct_reply"
    assert decision.sql is None