""" Tests for URL classifier. """ from local_deep_research.content_fetcher.url_classifier import ( URLClassifier, URLType, ) class TestURLClassifier: """Test URL classification.""" def test_classify_arxiv_abs(self): """Test arXiv abstract URL classification.""" url = "https://arxiv.org/abs/2301.12345" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_arxiv_pdf(self): """Test arXiv PDF URL classification.""" url = "https://arxiv.org/pdf/2301.12345.pdf" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_arxiv_html(self): """Test arXiv HTML URL classification.""" url = "https://arxiv.org/html/2301.12345" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_arxiv_old_format(self): """Test old arXiv format URL classification.""" url = "https://arxiv.org/abs/cond-mat/0501234" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_ar5iv(self): """Test ar5iv (HTML arXiv) URL classification.""" url = "https://ar5iv.org/abs/2301.12345" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_pubmed(self): """Test PubMed URL classification.""" url = "https://pubmed.ncbi.nlm.nih.gov/12345678" assert URLClassifier.classify(url) == URLType.PUBMED def test_classify_pubmed_old_format(self): """Test old PubMed URL format.""" url = "https://www.ncbi.nlm.nih.gov/pubmed/12345678" assert URLClassifier.classify(url) == URLType.PUBMED def test_classify_pmc(self): """Test PubMed Central URL classification.""" url = "https://www.ncbi.nlm.nih.gov/pmc/articles/PMC1234567" assert URLClassifier.classify(url) == URLType.PMC def test_classify_europe_pmc(self): """Test Europe PMC URL classification.""" url = "https://europepmc.org/article/PMC/1234567" assert URLClassifier.classify(url) == URLType.PMC def test_classify_semantic_scholar(self): """Test Semantic Scholar URL classification.""" # Semantic Scholar uses 40-char hex paper IDs - this is a fake test ID, not a secret url = "https://www.semanticscholar.org/paper/Attention-Is-All-You-Need-Vaswani-Shazeer/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" # DevSkim: ignore DS173237 assert URLClassifier.classify(url) == URLType.SEMANTIC_SCHOLAR def test_classify_semantic_scholar_api(self): """Test Semantic Scholar API URL classification.""" url = "https://api.semanticscholar.org/v1/paper/abcd1234" assert URLClassifier.classify(url) == URLType.SEMANTIC_SCHOLAR def test_classify_biorxiv(self): """Test bioRxiv URL classification.""" url = "https://www.biorxiv.org/content/10.1101/2021.01.01.12345v1" assert URLClassifier.classify(url) == URLType.BIORXIV def test_classify_medrxiv(self): """Test medRxiv URL classification.""" url = "https://www.medrxiv.org/content/10.1101/2021.01.01.12345v1" assert URLClassifier.classify(url) == URLType.MEDRXIV def test_classify_doi(self): """Test DOI URL classification.""" url = "https://doi.org/10.1234/example.2021.01" assert URLClassifier.classify(url) == URLType.DOI def test_classify_dx_doi(self): """Test dx.doi.org URL classification.""" url = "https://dx.doi.org/10.1234/example" assert URLClassifier.classify(url) == URLType.DOI def test_classify_direct_pdf(self): """Test direct PDF URL classification.""" url = "https://example.com/papers/document.pdf" assert URLClassifier.classify(url) == URLType.PDF def test_classify_html_generic(self): """Test generic HTML URL classification.""" url = "https://www.example.com/article/123" assert URLClassifier.classify(url) == URLType.HTML def test_classify_wikipedia(self): """Test Wikipedia URL (should be HTML).""" url = "https://en.wikipedia.org/wiki/Machine_learning" assert URLClassifier.classify(url) == URLType.HTML def test_classify_github(self): """Test GitHub URL (should be HTML).""" url = "https://github.com/anthropics/claude-code" assert URLClassifier.classify(url) == URLType.HTML def test_classify_news_site(self): """Test news site URL (should be HTML).""" url = "https://news.ycombinator.com/item?id=12345" assert URLClassifier.classify(url) == URLType.HTML class TestURLClassifierExtractID: """Test ID extraction from URLs.""" def test_extract_arxiv_id_new_format(self): """Test arXiv ID extraction (new format).""" url = "https://arxiv.org/abs/2301.12345" assert URLClassifier.extract_id(url) == "2301.12345" def test_extract_arxiv_id_with_version(self): """Test arXiv ID extraction with version.""" url = "https://arxiv.org/abs/2301.12345v2" assert URLClassifier.extract_id(url) == "2301.12345v2" def test_extract_arxiv_id_old_format(self): """Test arXiv ID extraction (old format).""" url = "https://arxiv.org/abs/cond-mat/0501234" assert URLClassifier.extract_id(url) == "cond-mat/0501234" def test_extract_pubmed_id(self): """Test PubMed ID extraction.""" url = "https://pubmed.ncbi.nlm.nih.gov/12345678/" assert URLClassifier.extract_id(url) == "12345678" def test_extract_pmc_id(self): """Test PMC ID extraction.""" url = "https://www.ncbi.nlm.nih.gov/pmc/articles/PMC1234567/" assert URLClassifier.extract_id(url) == "PMC1234567" def test_extract_semantic_scholar_id(self): """Test Semantic Scholar ID extraction.""" # Semantic Scholar uses 40-char hex paper IDs - this is a fake test ID, not a secret url = "https://www.semanticscholar.org/paper/Title/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" # DevSkim: ignore DS173237 extracted = URLClassifier.extract_id(url) assert extracted == "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" def test_extract_doi(self): """Test DOI extraction.""" url = "https://doi.org/10.1234/example.2021" assert URLClassifier.extract_id(url) == "10.1234/example.2021" def test_extract_id_unknown_url(self): """Test ID extraction from unknown URL type.""" url = "https://example.com/article/123" assert URLClassifier.extract_id(url) is None class TestURLClassifierSourceName: """Test human-readable source names.""" def test_source_name_arxiv(self): """Test arXiv source name.""" assert URLClassifier.get_source_name(URLType.ARXIV) == "arXiv" def test_source_name_pubmed(self): """Test PubMed source name.""" assert URLClassifier.get_source_name(URLType.PUBMED) == "PubMed" def test_source_name_pmc(self): """Test PMC source name.""" assert URLClassifier.get_source_name(URLType.PMC) == "PubMed Central" def test_source_name_semantic_scholar(self): """Test Semantic Scholar source name.""" assert ( URLClassifier.get_source_name(URLType.SEMANTIC_SCHOLAR) == "Semantic Scholar" ) def test_source_name_html(self): """Test HTML source name.""" assert URLClassifier.get_source_name(URLType.HTML) == "Web Page" def test_source_name_pdf(self): """Test PDF source name.""" assert URLClassifier.get_source_name(URLType.PDF) == "PDF" class TestURLClassifierEdgeCases: """Test edge cases and malformed URLs.""" def test_classify_empty_url(self): """Test empty URL defaults to HTML.""" assert URLClassifier.classify("") == URLType.HTML def test_classify_invalid_url(self): """Test invalid URL defaults to HTML.""" assert URLClassifier.classify("not-a-valid-url") == URLType.HTML def test_classify_case_insensitive(self): """Test URL classification is case-insensitive.""" url = "https://ARXIV.ORG/abs/2301.12345" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_with_query_params(self): """Test URL with query parameters.""" url = "https://arxiv.org/abs/2301.12345?ref=homepage" assert URLClassifier.classify(url) == URLType.ARXIV def test_classify_with_fragment(self): """Test URL with fragment.""" url = "https://pubmed.ncbi.nlm.nih.gov/12345678#abstract" assert URLClassifier.classify(url) == URLType.PUBMED class TestURLClassifierSecurity: """Test security-related URL classification.""" def test_javascript_url_rejected(self): """Test that javascript: URLs are classified as INVALID.""" url = "javascript:alert('XSS')" assert URLClassifier.classify(url) == URLType.INVALID def test_data_url_rejected(self): """Test that data: URLs are classified as INVALID.""" url = "data:text/html," assert URLClassifier.classify(url) == URLType.INVALID def test_file_url_rejected(self): """Test that file: URLs are classified as INVALID.""" url = "file:///etc/passwd" assert URLClassifier.classify(url) == URLType.INVALID def test_vbscript_url_rejected(self): """Test that vbscript: URLs are classified as INVALID.""" url = "vbscript:msgbox('XSS')" assert URLClassifier.classify(url) == URLType.INVALID def test_about_url_rejected(self): """Test that about: URLs are classified as INVALID.""" url = "about:blank" assert URLClassifier.classify(url) == URLType.INVALID def test_ftp_url_rejected(self): """Test that ftp: URLs are classified as INVALID.""" url = "ftp://ftp.example.com/file.txt" assert URLClassifier.classify(url) == URLType.INVALID def test_http_url_accepted(self): """Test that http: URLs are accepted.""" url = "http://example.com/page" assert URLClassifier.classify(url) == URLType.HTML def test_https_url_accepted(self): """Test that https: URLs are accepted.""" url = "https://example.com/page" assert URLClassifier.classify(url) == URLType.HTML def test_invalid_source_name(self): """Test INVALID URL source name.""" assert URLClassifier.get_source_name(URLType.INVALID) == "Invalid URL" def test_javascript_url_case_insensitive(self): """Test that JavaScript: URLs (mixed case) are rejected.""" url = "JavaScript:alert('XSS')" assert URLClassifier.classify(url) == URLType.INVALID def test_whitespace_trimmed(self): """Test that URLs with whitespace are trimmed.""" url = " https://example.com/page " assert URLClassifier.classify(url) == URLType.HTML