forked from Gitlink/gitlink-cli
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Assert-based unit tests for spark.py pure parsers. Run: python test_spark.py"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
from spark import parse_arxiv_atom, parse_github_search, extract_method_keywords
|
|
|
|
SAMPLE_ARXIV = """<?xml version="1.0" encoding="UTF-8"?>
|
|
<feed xmlns="http://www.w3.org/2005/Atom">
|
|
<entry>
|
|
<id>http://arxiv.org/abs/2403.12345v1</id>
|
|
<title>Graph Attention Networks with Sparse Transformers</title>
|
|
<summary>We propose a new graph attention mechanism using sparse attention.</summary>
|
|
<published>2024-03-15T00:00:00Z</published>
|
|
</entry>
|
|
<entry>
|
|
<id>http://arxiv.org/abs/2404.99999v2</id>
|
|
<title>Federated Learning on Heterogeneous Graphs</title>
|
|
<summary>A federated approach for heterogeneous graph neural networks.</summary>
|
|
<published>2024-04-20T00:00:00Z</published>
|
|
</entry>
|
|
</feed>"""
|
|
|
|
def test_parse_arxiv_atom():
|
|
papers = parse_arxiv_atom(SAMPLE_ARXIV)
|
|
assert len(papers) == 2, f"expected 2 papers, got {len(papers)}"
|
|
assert papers[0]["arxiv_id"] == "2403.12345v1", papers[0]["arxiv_id"]
|
|
assert "Graph Attention" in papers[0]["title"]
|
|
assert papers[0]["published"] == "2024-03-15"
|
|
assert "sparse" in papers[0]["abstract"].lower()
|
|
print("test_parse_arxiv_atom OK")
|
|
|
|
def test_parse_github_search():
|
|
import json as _j
|
|
sample = _j.dumps({"total_count": 1543, "items": [{"full_name": "a/b", "stargazers_count": 3534}]})
|
|
res = parse_github_search(sample)
|
|
assert res["total_count"] == 1543
|
|
assert res["top"][0]["full_name"] == "a/b"
|
|
assert res["top"][0]["stars"] == 3534
|
|
print("test_parse_github_search OK")
|
|
|
|
def test_extract_method_keywords():
|
|
kws = extract_method_keywords("Graph Attention Networks", "We propose a sparse attention mechanism for graphs.", max_k=5)
|
|
assert "graph" in kws and "attention" in kws
|
|
assert "propose" not in kws # 'propose' is in the stop set, filtered out
|
|
print("test_extract_method_keywords OK")
|
|
|
|
if __name__ == "__main__":
|
|
test_parse_arxiv_atom()
|
|
test_parse_github_search()
|
|
test_extract_method_keywords()
|
|
print("ALL TESTS PASSED")
|