YG-Rules/tests/test_domain_routes.py

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()