Files
pipeline-lifetime/tests/test_auth_registration.py
T
jiang c7ee2adb82 feat(auth): add password reset flow
Add admin-generated reset links, reset UI, timezone-aware expiry display, and registration captcha coverage.
2026-07-06 17:55:10 +08:00

524 lines
22 KiB
Python

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("第 <span class=\"font-bold text-textMain\">2</span> / 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()