Spaces:
Running
Running
| """Tests for gazet.lm — prompt templates, postprocessing, and GGUF helpers.""" | |
| from gazet.lm import ( | |
| _PLACES_SYSTEM_PROMPT, | |
| _SYSTEM_PROMPT_TEMPLATE, | |
| _USER_PROMPT_TEMPLATE, | |
| _postprocess_sql, | |
| ) | |
| from gazet.schemas import PlacesResult | |
| class TestPostprocessSql: | |
| def test_plain_sql(self): | |
| assert _postprocess_sql("SELECT * FROM foo") == "SELECT * FROM foo" | |
| def test_sql_fence_with_lang(self): | |
| raw = "```sql\nSELECT id FROM bar\n```" | |
| assert _postprocess_sql(raw) == "SELECT id FROM bar" | |
| def test_sql_fence_no_lang(self): | |
| raw = "```\nSELECT 1\n```" | |
| assert _postprocess_sql(raw) == "SELECT 1" | |
| def test_fence_with_trailing_text(self): | |
| raw = "```sql\nSELECT * FROM foo\n```\nDone!" | |
| result = _postprocess_sql(raw) | |
| assert "SELECT" in result | |
| assert "FROM" in result | |
| def test_nested_backticks_preserved(self): | |
| raw = "```sql\nSELECT `col`\n```" | |
| assert _postprocess_sql(raw) == "SELECT `col`" | |
| def test_whitespace_only(self): | |
| assert _postprocess_sql(" ") == "" | |
| def test_empty_string(self): | |
| assert _postprocess_sql("") == "" | |
| def test_multiline_query(self): | |
| raw = """```sql | |
| SELECT | |
| id, | |
| names."primary" AS name | |
| FROM read_parquet('divisions_area') | |
| WHERE country = 'IN' | |
| ```""" | |
| result = _postprocess_sql(raw) | |
| assert "SELECT" in result | |
| assert "WHERE" in result | |
| assert "```" not in result | |
| def test_no_fence_just_whitespace_surround(self): | |
| raw = " SELECT * FROM table " | |
| assert _postprocess_sql(raw) == "SELECT * FROM table" | |
| class TestPromptTemplates: | |
| def test_system_prompt_has_schema_placeholder(self): | |
| assert "{schema}" in _SYSTEM_PROMPT_TEMPLATE | |
| def test_system_prompt_mentions_st_asgeojson(self): | |
| assert "ST_AsGeoJSON" in _SYSTEM_PROMPT_TEMPLATE | |
| def test_user_prompt_has_candidates_placeholder(self): | |
| assert "{candidates_csv}" in _USER_PROMPT_TEMPLATE | |
| def test_user_prompt_has_question_placeholder(self): | |
| assert "{question}" in _USER_PROMPT_TEMPLATE | |
| def test_system_prompt_formatting(self): | |
| formatted = _SYSTEM_PROMPT_TEMPLATE.format(schema="dummy_schema") | |
| assert "dummy_schema" in formatted | |
| assert "{schema}" not in formatted | |
| def test_user_prompt_formatting(self): | |
| formatted = _USER_PROMPT_TEMPLATE.format( | |
| candidates_csv="source,id,name", | |
| question="get boundary", | |
| ) | |
| assert "source,id,name" in formatted | |
| assert "get boundary" in formatted | |
| def test_places_system_prompt_has_examples(self): | |
| assert "EXAMPLES:" in _PLACES_SYSTEM_PROMPT | |
| assert "Puri, Odisha" in _PLACES_SYSTEM_PROMPT | |
| assert "Amazon basin" in _PLACES_SYSTEM_PROMPT | |
| class TestGeneratePlacesFallback: | |
| """Test that generate_places handles parse errors gracefully. | |
| We can't test the real llama-server path without a running server, | |
| but we can verify the PlacesResult model behavior used in fallback. | |
| """ | |
| def test_places_result_from_raw_query(self): | |
| # Simulates the fallback path in generate_places | |
| query = "get me the boundary of Paris" | |
| result = PlacesResult(places=[{"place": query}]) | |
| assert len(result.places) == 1 | |
| assert result.places[0].place == query | |
| def test_places_result_validation(self): | |
| data = {"places": [{"place": "Chad"}, {"place": "Lake Chad"}]} | |
| r = PlacesResult.model_validate(data) | |
| assert len(r.places) == 2 | |