import json
from bench.bench_tokens import score_result, _extract_first_json
def test_extraction():
# 1. Clean JSON
assert _extract_first_json('{"a": 1}') == {"a": 1}
# 2. Markdown fenced
assert _extract_first_json('```json\n{"a": 1}\n```') == {"a": 1}
# 3. Chatty prose
assert _extract_first_json('Here is the JSON: {"a": 1} hope it helps!') == {"a": 1}
# 4. Nested structures
assert _extract_first_json('Prefix [1, 2, {"b": 3}] Suffix') == [1, 2, {"b": 3}]
def test_scoring():
# 1. Pure JSON
s1 = score_result('{"a": 1}')
assert s1["parse_ok"] == 1
assert s1["no_prose"] == 1
# 2. Fenced JSON
s2 = score_result('```json\n{"a": 1}\n```')
assert s2["parse_ok"] == 1
assert s2["no_prose"] == 1
# 3. Chatty JSON
s3 = score_result('Here is the JSON: {"a": 1}')
assert s3["parse_ok"] == 1
assert s3["no_prose"] == 0
# 4. Type checking
s4 = score_result('{"a": 1}', expected_type=dict)
assert s4["is_expected_type"] == 1
s5 = score_result('[1, 2]', expected_type=dict)
assert s5["is_expected_type"] == 0
# 5. Structured Content (Preference)
s6 = score_result('Garbage text', structured_content={"a": 1})
assert s6["parse_ok"] == 1
assert s6["no_prose"] == 1 # Structured content is considered clean by definition
if __name__ == "__main__":
print("Running extraction tests...")
test_extraction()
print("Running scoring tests...")
test_scoring()
print("All tests passed!")
