208 lines
8.9 KiB
Python
208 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import unittest
|
|
from datetime import datetime, timedelta
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from app import create_app
|
|
from app.config import Config
|
|
from app.extensions import db
|
|
from app.models import AppSetting, 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) -> None:
|
|
with app.app_context():
|
|
user = User(username=username, is_admin=is_admin)
|
|
user.set_password(password)
|
|
db.session.add(user)
|
|
db.session.commit()
|
|
|
|
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 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_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)
|
|
|
|
response = client.post(
|
|
"/register",
|
|
data={"csrf_token": token, "username": "new-user", "password": "secret123"},
|
|
)
|
|
|
|
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_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)
|
|
|
|
response = client.post(
|
|
"/register",
|
|
data={"csrf_token": token, "username": "new-user", "password": "secret123"},
|
|
)
|
|
|
|
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_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()
|