166 lines
7.4 KiB
Python
166 lines
7.4 KiB
Python
import base64
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from http.server import ThreadingHTTPServer
|
|
|
|
import requests
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from app import downloader
|
|
Downloader, make_handler, fingerprints = downloader.Downloader, downloader.make_handler, downloader.fingerprints
|
|
from article_writer import save_article
|
|
|
|
|
|
def article(i):
|
|
return dict(title=f'新闻{i}', content=f'<p>第{i}篇新闻内容</p>', url=f'https://example.com/{i}', src='测试来源', time='2026-10-31')
|
|
|
|
|
|
class DownloadTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp.name)
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def run_job(self, manager, target):
|
|
manager.start('test-key', '头条', target)
|
|
manager.thread.join(10)
|
|
self.assertFalse(manager.thread.is_alive())
|
|
return manager.snapshot()
|
|
|
|
def test_format_and_local_image(self):
|
|
data = article(1)
|
|
image = 'data:image/png;base64,' + base64.b64encode(b'\x89PNG\r\n\x1a\nexample').decode()
|
|
data['content'] += f'<img src="{image}"><img src="{image}">'
|
|
result = save_article(data, self.root)
|
|
body = Path(result.markdown).read_text()
|
|
self.assertTrue(body.startswith('# 新闻1\n\n第1篇新闻内容'))
|
|
self.assertNotIn('发布时间：2026-10-31', body)
|
|
self.assertNotIn('\n---\n', body)
|
|
self.assertTrue(body.endswith('来源：测试来源，如有侵权，请联系删除。\n'))
|
|
self.assertEqual(len(result.images), 1)
|
|
self.assertEqual(body.count('image001.png'), 2)
|
|
self.assertNotIn('data:image', body)
|
|
second = save_article(data, self.root)
|
|
self.assertEqual(Path(second.markdown).name, '新闻1 (2).md')
|
|
|
|
def test_bad_image_cleans_partial(self):
|
|
data = article(1)
|
|
data['content'] += '<img src="data:image/png;base64,aW52YWxpZA==">'
|
|
with self.assertRaises(Exception):
|
|
save_article(data, self.root)
|
|
self.assertEqual(list(self.root.iterdir()), [])
|
|
|
|
def test_duplicate_then_fill_and_restart(self):
|
|
items = [article(1), article(1), article(2), article(3)]
|
|
calls = []
|
|
def request(key, endpoint, params):
|
|
calls.append(params.copy())
|
|
return {'code': 1, 'data': {'list': items[params['start']:params['start']+params['num']]}}
|
|
manager = Downloader(self.root, request)
|
|
state = self.run_job(manager, 3)
|
|
self.assertEqual((state['status'], state['saved'], state['skipped']), ('completed', 3, 1))
|
|
self.assertEqual([c['start'] for c in calls], [0, 3])
|
|
restarted = Downloader(self.root, request)
|
|
state = self.run_job(restarted, 3)
|
|
self.assertEqual((state['status'], state['saved'], state['skipped']), ('partial', 0, 4))
|
|
self.assertEqual(len(list(self.root.glob('*/*.md'))), 3)
|
|
|
|
def test_offset_limit_and_partial_progress(self):
|
|
offsets = []
|
|
def request(key, endpoint, params):
|
|
offsets.append(params['start'])
|
|
return {'code': 1, 'data': {'list': [article(i) for i in range(params['start'], params['start']+params['num'])]}}
|
|
state = self.run_job(Downloader(self.root, request), 500)
|
|
self.assertEqual(state['status'], 'partial')
|
|
self.assertEqual(state['saved'], 440)
|
|
self.assertEqual(max(offsets), 400)
|
|
|
|
def test_error_and_key_redaction(self):
|
|
def request(*args):
|
|
return {'code': 10001, 'msg': 'invalid test-key'}
|
|
manager = Downloader(self.root, request)
|
|
state = self.run_job(manager, 10)
|
|
self.assertEqual(state['status'], 'error')
|
|
self.assertNotIn('test-key', json.dumps(state))
|
|
|
|
def test_stop_and_concurrent_start(self):
|
|
entered, release = threading.Event(), threading.Event()
|
|
def request(*args):
|
|
entered.set()
|
|
release.wait(2)
|
|
return {'code': 1, 'data': {'list': [article(1)]}}
|
|
manager = Downloader(self.root, request)
|
|
manager.start('key', '头条', 10)
|
|
self.assertTrue(entered.wait(2))
|
|
with self.assertRaises(ValueError):
|
|
manager.start('key', '头条', 10)
|
|
manager.stop.set()
|
|
release.set()
|
|
manager.thread.join(3)
|
|
self.assertEqual(manager.snapshot()['status'], 'cancelled')
|
|
self.assertEqual(manager.snapshot()['saved'], 0)
|
|
|
|
def test_failed_save_can_retry_next_job(self):
|
|
attempts = []
|
|
def writer(data, root):
|
|
attempts.append(1)
|
|
if len(attempts) == 1:
|
|
raise ValueError('模拟图片失败')
|
|
return save_article(data, root)
|
|
def request(key, endpoint, params):
|
|
return {'code': 1, 'data': {'list': [article(1)] if params['start'] == 0 else []}}
|
|
manager = Downloader(self.root, request, writer)
|
|
self.assertEqual(self.run_job(manager, 1)['failed'], 1)
|
|
self.assertEqual(self.run_job(manager, 1)['saved'], 1)
|
|
|
|
def test_fingerprint_content_and_tracking(self):
|
|
a, b = article(1), article(2)
|
|
b['content'] = a['content'].replace('<p>', '<div>').replace('</p>', '</div>')
|
|
self.assertTrue(set(fingerprints(a)) & set(fingerprints(b)))
|
|
b['content'] = '其他内容'
|
|
b['url'] = a['url'] + '?utm_source=abc#fragment'
|
|
self.assertTrue(set(fingerprints(a)) & set(fingerprints(b)))
|
|
|
|
def test_log_cap(self):
|
|
manager = Downloader(self.root)
|
|
for i in range(400):
|
|
manager.log(str(i))
|
|
self.assertEqual(len(manager.snapshot()['logs']), 300)
|
|
self.assertEqual(manager.snapshot()['logs'][0]['message'], '100')
|
|
|
|
def test_http_routes_and_access(self):
|
|
def request(key, endpoint, params=None):
|
|
return {'code': 1, 'data': {'list': ['头条', '科技']}}
|
|
manager = Downloader(self.root, request)
|
|
server = ThreadingHTTPServer(('127.0.0.1', 0), make_handler(manager, 'token'))
|
|
worker = threading.Thread(target=server.serve_forever, daemon=True)
|
|
worker.start()
|
|
url = f'http://127.0.0.1:{server.server_port}'
|
|
try:
|
|
self.assertIn('新闻下载器', requests.get(url).text)
|
|
self.assertEqual(requests.get(url+'/api/state').json()['status'], 'idle')
|
|
self.assertEqual(requests.post(url+'/api/channels', json={'key': 'test'}).status_code, 403)
|
|
self.assertEqual(requests.get(url, headers={'Host':'evil.example'}).status_code, 403)
|
|
headers = {'X-Session-Token': 'token'}
|
|
self.assertEqual(requests.post(url+'/api/channels', headers=headers, json={'key':'test'}).json()['channels'], ['头条','科技'])
|
|
invalid = requests.post(url+'/api/start', headers=headers, json={'key':'test', 'target':0,'channel':'头条'})
|
|
self.assertEqual(invalid.status_code, 400)
|
|
invalid = requests.post(url+'/api/debug', headers=headers, json={'key':'test','endpoint':'index','num':41,'start':0})
|
|
self.assertEqual(invalid.status_code, 400)
|
|
debug = requests.post(url+'/api/debug', headers=headers, json={'key':'test','endpoint':'channel'})
|
|
self.assertEqual(debug.status_code, 200)
|
|
self.assertNotIn('key=test', debug.text)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
worker.join()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|