47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from search.evaluate_relevance import evaluate, load_queries
|
|
|
|
|
|
class SearchRelevanceTest(unittest.TestCase):
|
|
def test_loads_queries_and_calculates_document_metrics(self):
|
|
queries = [
|
|
{
|
|
"id": "ru-01",
|
|
"language": "ru",
|
|
"query": "трудовой договор",
|
|
"relevant_document_codes": ["7", "8"],
|
|
},
|
|
{
|
|
"id": "ky-01",
|
|
"language": "ky",
|
|
"query": "эмгек келишими",
|
|
"relevant_document_codes": ["9"],
|
|
},
|
|
]
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
path = Path(temporary) / "queries.json"
|
|
path.write_text(json.dumps(queries, ensure_ascii=False), encoding="utf-8")
|
|
loaded = load_queries(path)
|
|
|
|
responses = [
|
|
{"hits": {"hits": [{"_source": {"document_code": code}} for code in ["7", "10", "8"]]}},
|
|
{"hits": {"hits": [{"_source": {"document_code": code}} for code in ["10", "9"]]}},
|
|
]
|
|
with patch("search.evaluate_relevance.request_json", side_effect=responses) as request:
|
|
result = evaluate(loaded, "http://127.0.0.1:9200", "test", 10)
|
|
|
|
self.assertEqual(result["summary"], {"query_count": 2, "recall_at_10": 1.0, "mrr_at_10": 0.75})
|
|
self.assertEqual(result["queries"][0]["reciprocal_rank_at_10"], 1.0)
|
|
body = json.loads(request.call_args_list[0].args[2])
|
|
self.assertFalse(body["track_total_hits"])
|
|
self.assertEqual(body["collapse"], {"field": "document_code"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|