50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import unittest
|
|
from io import BytesIO
|
|
from unittest.mock import patch
|
|
|
|
from flask import Flask
|
|
|
|
from app.routes import domain as domain_routes
|
|
from app.routes.domain import domain_bp
|
|
|
|
|
|
class DomainRouteTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.app = Flask(__name__)
|
|
self.app.register_blueprint(domain_bp)
|
|
self.client = self.app.test_client()
|
|
|
|
def test_clear_domains_also_clears_schema(self):
|
|
app = Flask(__name__)
|
|
|
|
with app.app_context():
|
|
with patch("app.routes.domain.DomainStorage") as domain_storage_cls:
|
|
with patch("app.routes.domain.SchemaStorage") as schema_storage_cls:
|
|
domain_storage_cls.return_value.clear_all.return_value = 2
|
|
|
|
response, status_code = domain_routes.clear_domains()
|
|
|
|
self.assertEqual(status_code, 200)
|
|
domain_storage_cls.return_value.clear_all.assert_called_once_with()
|
|
schema_storage_cls.return_value.delete_file.assert_called_once_with()
|
|
self.assertEqual(
|
|
response.get_json()["data"],
|
|
{"deleted_count": 2, "schema_cleared": True},
|
|
)
|
|
|
|
def test_upload_domains_rejects_empty_file_without_overwriting(self):
|
|
with patch("app.routes.domain.DomainStorage") as domain_storage_cls:
|
|
response = self.client.post(
|
|
"/api/domains/upload",
|
|
data={"file": (BytesIO("domain,note\n".encode("utf-8")), "domains.csv")},
|
|
content_type="multipart/form-data",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, 400)
|
|
self.assertEqual(response.get_json()["message"], "上传失败：未解析到风险领域")
|
|
domain_storage_cls.return_value.save_domains.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|