from __future__ import annotations import re import unittest from datetime import datetime, timedelta from tempfile import TemporaryDirectory from urllib.parse import urlparse from app import create_app from app.config import Config from app.extensions import db from app.models import AppSetting, PasswordResetToken, UploadRecord, User class RegistrationRoutesTest(unittest.TestCase): def create_test_app(self, temp_dir: str, *, allow_registration: bool): class TestConfig(Config): TESTING = True SECRET_KEY = "test-secret" SECRET_KEY_GENERATED = False SQLALCHEMY_DATABASE_URI = f"sqlite:///{temp_dir}/test.db" ALLOW_REGISTRATION = allow_registration ADMIN_PASSWORD = None return create_app(TestConfig, load_model_on_start=False) def csrf_token_from(self, html: bytes) -> str: match = re.search(rb'name="csrf_token" value="([^"]+)"', html) self.assertIsNotNone(match) return match.group(1).decode() def create_user(self, app, username: str, password: str, *, is_admin: bool = False) -> int: with app.app_context(): user = User(username=username, is_admin=is_admin) user.set_password(password) db.session.add(user) db.session.commit() return user.id def add_upload_records(self, app, username: str, count: int) -> None: with app.app_context(): user = User.query.filter_by(username=username).one() base_time = datetime(2026, 1, 1, 12, 0, 0) for index in range(count): db.session.add( UploadRecord( user_id=user.id, original_filename=f"{username}-file-{index:02d}.xlsx", saved_path=f"/tmp/{username}-original-{index:02d}.xlsx", prediction_path=f"/tmp/{username}-prediction-{index:02d}.xlsx", image_path=f"/tmp/{username}-image-{index:02d}.png", upload_time=base_time + timedelta(minutes=index), ) ) db.session.commit() def login(self, client, username: str, password: str) -> None: response = client.get("/login") token = self.csrf_token_from(response.data) with client.session_transaction() as session: captcha = session["captcha"] login_response = client.post( "/login", data={ "csrf_token": token, "username": username, "password": password, "captcha": captcha, }, ) self.assertEqual(login_response.status_code, 302) def login_attempt(self, client, username: str, password: str): response = client.get("/login") token = self.csrf_token_from(response.data) with client.session_transaction() as session: captcha = session["captcha"] return client.post( "/login", data={ "csrf_token": token, "username": username, "password": password, "captcha": captcha, }, ) def create_reset_link(self, app, client, user_id: int) -> str: token = self.csrf_token_from(client.get("/admin").data) response = client.post( f"/admin/users/{user_id}/password-reset-link", data={"csrf_token": token}, ) self.assertEqual(response.status_code, 200) return urlparse(response.get_json()["reset_url"]).path def test_login_page_always_shows_register_entry_when_registration_is_closed(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) response = app.test_client().get("/login") self.assertEqual(response.status_code, 200) self.assertIn('href="/register"', response.get_data(as_text=True)) def test_register_page_is_visible_but_disabled_when_registration_is_closed(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) response = app.test_client().get("/register") html = response.get_data(as_text=True) self.assertEqual(response.status_code, 200) self.assertIn("当前未开放自助注册", html) self.assertIn("disabled", html) def test_register_page_shows_captcha(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=True) client = app.test_client() response = client.get("/register") html = response.get_data(as_text=True) with client.session_transaction() as session: captcha = session["captcha"] self.assertEqual(response.status_code, 200) self.assertIn('name="captcha"', html) self.assertIn(captcha, html) def test_register_post_does_not_create_user_when_registration_is_closed(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) client = app.test_client() token = self.csrf_token_from(client.get("/register").data) with client.session_transaction() as session: captcha = session["captcha"] response = client.post( "/register", data={ "csrf_token": token, "username": "new-user", "password": "secret123", "captcha": captcha, }, ) self.assertEqual(response.status_code, 403) self.assertIn("当前未开放自助注册", response.get_data(as_text=True)) with app.app_context(): self.assertIsNone(User.query.filter_by(username="new-user").first()) def test_register_post_checks_captcha_before_registration_setting(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) client = app.test_client() token = self.csrf_token_from(client.get("/register").data) response = client.post( "/register", data={ "csrf_token": token, "username": "new-user", "password": "secret123", "captcha": "WRONG", }, ) self.assertEqual(response.status_code, 400) self.assertIn("验证码错误", response.get_data(as_text=True)) with app.app_context(): self.assertIsNone(User.query.filter_by(username="new-user").first()) def test_register_post_rejects_wrong_captcha(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=True) client = app.test_client() token = self.csrf_token_from(client.get("/register").data) response = client.post( "/register", data={ "csrf_token": token, "username": "new-user", "password": "secret123", "captcha": "WRONG", }, ) self.assertEqual(response.status_code, 400) self.assertIn("验证码错误", response.get_data(as_text=True)) with app.app_context(): self.assertIsNone(User.query.filter_by(username="new-user").first()) def test_register_post_creates_user_when_registration_is_open(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=True) client = app.test_client() token = self.csrf_token_from(client.get("/register").data) with client.session_transaction() as session: captcha = session["captcha"] response = client.post( "/register", data={ "csrf_token": token, "username": "new-user", "password": "secret123", "captcha": captcha, }, ) self.assertEqual(response.status_code, 200) self.assertIn("注册成功，请登录", response.get_data(as_text=True)) with app.app_context(): user = User.query.filter_by(username="new-user").first() self.assertIsNotNone(user) self.assertFalse(user.is_admin) def test_admin_can_enable_registration_from_admin_panel(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) client = app.test_client() self.login(client, "admin", "secret123") token = self.csrf_token_from(client.get("/admin").data) response = client.post( "/admin/registration", data={"csrf_token": token, "allow_registration": "on"}, follow_redirects=True, ) self.assertEqual(response.status_code, 200) self.assertIn("已开放用户自助注册", response.get_data(as_text=True)) with app.app_context(): self.assertTrue(AppSetting.get_bool("allow_registration")) register_page = client.get("/register").get_data(as_text=True) self.assertNotIn("当前未开放自助注册", register_page) def test_admin_can_disable_registration_from_admin_panel(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=True) self.create_user(app, "admin", "secret123", is_admin=True) client = app.test_client() self.login(client, "admin", "secret123") token = self.csrf_token_from(client.get("/admin").data) response = client.post( "/admin/registration", data={"csrf_token": token}, follow_redirects=True, ) self.assertEqual(response.status_code, 200) self.assertIn("已关闭用户自助注册", response.get_data(as_text=True)) with app.app_context(): self.assertFalse(AppSetting.get_bool("allow_registration", True)) def test_admin_can_create_password_reset_link_for_regular_user(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") client = app.test_client() self.login(client, "admin", "secret123") token = self.csrf_token_from(client.get("/admin").data) response = client.post( f"/admin/users/{user_id}/password-reset-link", data={"csrf_token": token}, ) data = response.get_json() self.assertEqual(response.status_code, 200) self.assertIn("/password-reset/", data["reset_url"]) self.assertEqual(data["username"], "alice") with app.app_context(): self.assertEqual(PasswordResetToken.query.count(), 1) def test_admin_password_reset_section_lists_registered_users_without_uploads(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=True) self.create_user(app, "admin", "secret123", is_admin=True) client = app.test_client() register_page = client.get("/register") token = self.csrf_token_from(register_page.data) with client.session_transaction() as session: captcha = session["captcha"] register_response = client.post( "/register", data={ "csrf_token": token, "username": "registered-user", "password": "secret123", "captcha": captcha, }, ) self.assertEqual(register_response.status_code, 200) self.login(client, "admin", "secret123") admin_page = client.get("/admin").get_data(as_text=True) self.assertIn("用户密码重置", admin_page) self.assertIn("registered-user", admin_page) self.assertIn("生成重置链接", admin_page) def test_password_reset_link_requires_admin(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) user_id = self.create_user(app, "alice", "oldpass") client = app.test_client() token = self.csrf_token_from(client.get("/login").data) anonymous_response = client.post( f"/admin/users/{user_id}/password-reset-link", data={"csrf_token": token}, ) self.assertEqual(anonymous_response.status_code, 302) self.login(client, "alice", "oldpass") token = self.csrf_token_from(client.get("/home").data) user_response = client.post( f"/admin/users/{user_id}/password-reset-link", data={"csrf_token": token}, ) self.assertEqual(user_response.status_code, 403) def test_admin_cannot_create_password_reset_link_for_admin_user(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) admin_id = self.create_user(app, "admin", "secret123", is_admin=True) client = app.test_client() self.login(client, "admin", "secret123") token = self.csrf_token_from(client.get("/admin").data) response = client.post( f"/admin/users/{admin_id}/password-reset-link", data={"csrf_token": token}, ) self.assertEqual(response.status_code, 403) self.assertIn("管理员账号", response.get_json()["error"]) def test_password_reset_changes_password_and_consumes_link(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") admin_client = app.test_client() self.login(admin_client, "admin", "secret123") reset_path = self.create_reset_link(app, admin_client, user_id) client = app.test_client() reset_page = client.get(reset_path) token = self.csrf_token_from(reset_page.data) response = client.post( reset_path, data={ "csrf_token": token, "password": "newpass123", "password_confirm": "newpass123", }, ) self.assertEqual(response.status_code, 200) self.assertIn("密码已重置", response.get_data(as_text=True)) self.assertEqual(self.login_attempt(client, "alice", "oldpass").status_code, 400) self.assertEqual(self.login_attempt(client, "alice", "newpass123").status_code, 302) self.assertEqual(client.get(reset_path).status_code, 400) with app.app_context(): reset_token = PasswordResetToken.query.one() self.assertIsNotNone(reset_token.used_at) def test_password_reset_page_displays_expiry_in_configured_timezone(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") admin_client = app.test_client() self.login(admin_client, "admin", "secret123") reset_path = self.create_reset_link(app, admin_client, user_id) with app.app_context(): reset_token = PasswordResetToken.query.one() reset_token.expires_at = datetime(2027, 1, 1, 0, 0, 0) db.session.commit() response = app.test_client().get(reset_path) self.assertEqual(response.status_code, 200) self.assertIn("2027-01-01 08:00:00", response.get_data(as_text=True)) def test_expired_password_reset_link_cannot_change_password(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") admin_client = app.test_client() self.login(admin_client, "admin", "secret123") reset_path = self.create_reset_link(app, admin_client, user_id) with app.app_context(): reset_token = PasswordResetToken.query.one() reset_token.expires_at = datetime.utcnow() - timedelta(minutes=1) db.session.commit() client = app.test_client() token = self.csrf_token_from(client.get("/login").data) response = client.post( reset_path, data={ "csrf_token": token, "password": "newpass123", "password_confirm": "newpass123", }, ) self.assertEqual(response.status_code, 400) self.assertEqual(self.login_attempt(client, "alice", "oldpass").status_code, 302) def test_new_password_reset_link_invalidates_previous_link(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") client = app.test_client() self.login(client, "admin", "secret123") first_path = self.create_reset_link(app, client, user_id) second_path = self.create_reset_link(app, client, user_id) self.assertEqual(client.get(first_path).status_code, 400) self.assertEqual(client.get(second_path).status_code, 200) def test_password_reset_validation_does_not_consume_link(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) user_id = self.create_user(app, "alice", "oldpass") admin_client = app.test_client() self.login(admin_client, "admin", "secret123") reset_path = self.create_reset_link(app, admin_client, user_id) client = app.test_client() reset_page = client.get(reset_path) token = self.csrf_token_from(reset_page.data) short_response = client.post( reset_path, data={ "csrf_token": token, "password": "short", "password_confirm": "short", }, ) self.assertEqual(short_response.status_code, 400) reset_page = client.get(reset_path) token = self.csrf_token_from(reset_page.data) mismatch_response = client.post( reset_path, data={ "csrf_token": token, "password": "newpass123", "password_confirm": "different", }, ) self.assertEqual(mismatch_response.status_code, 400) self.assertEqual(client.get(reset_path).status_code, 200) with app.app_context(): reset_token = PasswordResetToken.query.one() self.assertIsNone(reset_token.used_at) def test_history_page_paginates_upload_records(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "alice", "secret123") self.add_upload_records(app, "alice", 12) client = app.test_client() self.login(client, "alice", "secret123") first_page = client.get("/history").get_data(as_text=True) second_page = client.get("/history?page=2").get_data(as_text=True) self.assertIn("共 12 条", first_page) self.assertIn("alice-file-11.xlsx", first_page) self.assertIn("alice-file-02.xlsx", first_page) self.assertNotIn("alice-file-01.xlsx", first_page) self.assertIn("第 2 / 2 页", second_page) self.assertIn("alice-file-01.xlsx", second_page) self.assertIn("alice-file-00.xlsx", second_page) def test_admin_page_paginates_upload_records(self) -> None: with TemporaryDirectory() as temp_dir: app = self.create_test_app(temp_dir, allow_registration=False) self.create_user(app, "admin", "secret123", is_admin=True) self.create_user(app, "alice", "secret123") self.add_upload_records(app, "alice", 12) client = app.test_client() self.login(client, "admin", "secret123") first_page = client.get("/admin").get_data(as_text=True) second_page = client.get("/admin?page=2").get_data(as_text=True) self.assertIn("共 12 条", first_page) self.assertIn("alice-file-11.xlsx", first_page) self.assertIn("alice-file-02.xlsx", first_page) self.assertNotIn("alice-file-01.xlsx", first_page) self.assertIn("alice-file-01.xlsx", second_page) self.assertIn("alice-file-00.xlsx", second_page) if __name__ == "__main__": unittest.main()