"""Admin-configurable UPI QR code (upload / replace / remove / fallback). Reuses the existing public-static image upload + validation helpers, so these tests exercise the real flow: storage under a UUID name, replacement cleanup, removal + fallback to the generated QR, validation, admin preview and the payment page display. Existing Premium/Payment behaviour is preserved (backward compat). """ import io import os import shutil import tempfile import unittest from pathlib import Path os.environ["TEST_DATABASE_URL"] = "sqlite:///:memory:" from werkzeug.security import generate_password_hash from app import create_app from app.extensions import db from app.models import User, UserRole class PremiumQrTests(unittest.TestCase): def setUp(self): self.app = create_app("testing") self.ctx = self.app.app_context() self.ctx.push() db.create_all() self._dir = Path(tempfile.mkdtemp()) self.app.config["PREMIUM_QR_UPLOAD_FOLDER"] = self._dir u = User(email="admin@c.edu", password_hash=generate_password_hash("Passw0rd!"), is_verified=True, role=UserRole.SUPER_ADMIN) db.session.add(u) db.session.commit() self.client = self.app.test_client() self.client.post("/login", data={"email": "admin@c.edu", "password": "Passw0rd!"}) def tearDown(self): shutil.rmtree(self._dir, ignore_errors=True) db.session.remove() db.drop_all() self.ctx.pop() # --- helpers --------------------------------------------------------------- def _real_png(self, name="qr.png"): """A genuine, decodable PNG (passes the strengthened content check).""" from PIL import Image buf = io.BytesIO() Image.new("RGB", (16, 16), "red").save(buf, "PNG") buf.seek(0) return (buf, name) def _raw(self, name, data): """Raw bytes with an arbitrary name (for type/size/corruption cases).""" return (io.BytesIO(data), name) def _upload(self, file=None, follow=True): return self.client.post( "/app/premium/admin/settings/qr", data={"action": "upload", "qr_image": file or self._real_png()}, content_type="multipart/form-data", follow_redirects=follow) def _remove(self): return self.client.post( "/app/premium/admin/settings/qr", data={"action": "remove"}, follow_redirects=True) def _settings(self): from app.premium import config return config.get_payment_settings() def _stored_files(self): return [p.name for p in self._dir.iterdir() if p.is_file()] def _become_student(self): """Switch the single test client's identity to a student. Uses one client (log out admin, log in student) rather than a second logged-in client — two concurrently-authenticated clients bleed identity in this harness (admin identity wins), which would mask authorization. """ if not db.session.scalar(db.select(User).filter_by(email="s@c.edu")): db.session.add(User(email="s@c.edu", password_hash=generate_password_hash("Passw0rd!"), is_verified=True, role=UserRole.STUDENT)) db.session.commit() self.client.post("/logout") self.client.post("/login", data={"email": "s@c.edu", "password": "Passw0rd!"}) # --- upload / preview ------------------------------------------------------ def test_upload_sets_qr_and_stores_file(self): self._upload() settings = self._settings() self.assertIsNotNone(settings.qr_image_path) self.assertTrue(settings.qr_image_path.startswith("/static/uploads/premium/")) self.assertEqual(len(self._stored_files()), 1) self.assertEqual(os.path.basename(settings.qr_image_path), self._stored_files()[0]) def test_uuid_filename(self): self._upload(self._real_png(name="my personal qr!.png")) stored = self._stored_files()[0] self.assertNotIn("personal", stored) # client name never used self.assertTrue(stored.endswith(".png")) self.assertEqual(len(stored.split(".")[0]), 32) # uuid4 hex def test_admin_preview_shows_uploaded_qr(self): self._upload() html = self.client.get("/app/premium/admin/settings").get_data(as_text=True) self.assertIn(self._settings().qr_image_path, html) self.assertIn("Remove QR Code", html) # remove control appears self.assertIn("Uploaded QR", html) def test_admin_placeholder_when_none(self): html = self.client.get("/app/premium/admin/settings").get_data(as_text=True) self.assertIn("No QR code uploaded.", html) self.assertIn("UPI QR Code", html) self.assertIn("Upload QR Code", html) self.assertNotIn("Remove QR Code", html) # nothing to remove # --- replace (cleanup) ----------------------------------------------------- def test_replace_deletes_old_file(self): self._upload() old = self._stored_files()[0] self._upload(self._real_png(name="new.png")) files = self._stored_files() self.assertEqual(len(files), 1) # old removed, only new remains self.assertNotIn(old, files) self.assertEqual(os.path.basename(self._settings().qr_image_path), files[0]) # --- remove / fallback ----------------------------------------------------- def test_remove_deletes_file_and_reverts_to_generated(self): from app.premium import config self._upload() self.assertEqual(len(self._stored_files()), 1) self._remove() self.assertIsNone(self._settings().qr_image_path) self.assertEqual(self._stored_files(), []) # asset deleted # Falls back to the generated QR (default seeded UPI id) — never empty. src = config.payment_qr_src() self.assertTrue(src and src.startswith("data:image/png;base64,")) def test_payment_qr_src_prefers_uploaded(self): from app.premium import config self.assertTrue(config.payment_qr_src().startswith("data:")) # generated first self._upload() self.assertEqual(config.payment_qr_src(), self._settings().qr_image_path) # --- validation ------------------------------------------------------------ def test_invalid_type_rejected(self): resp = self._upload(self._raw("notes.txt", b"just text")) self.assertIsNone(self._settings().qr_image_path) self.assertEqual(self._stored_files(), []) self.assertIn("Unsupported image type", resp.get_data(as_text=True)) def test_oversized_rejected(self): self.app.config["MAX_AD_IMAGE_SIZE"] = 1024 # 1 KB cap for the test resp = self._upload(self._raw("big.png", b"\x00" * 4096)) # ~4 KB self.assertIsNone(self._settings().qr_image_path) self.assertEqual(self._stored_files(), []) self.assertIn("too large", resp.get_data(as_text=True)) def test_corrupted_image_rejected(self): # Right extension + under size, but not a decodable image → content check. resp = self._upload(self._raw("fake.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 64)) self.assertIsNone(self._settings().qr_image_path) self.assertEqual(self._stored_files(), []) self.assertIn("not a valid image", resp.get_data(as_text=True)) def test_missing_file_rejected(self): resp = self.client.post("/app/premium/admin/settings/qr", data={"action": "upload"}, content_type="multipart/form-data", follow_redirects=True) self.assertIsNone(self._settings().qr_image_path) self.assertIn("Choose a QR image", resp.get_data(as_text=True)) # --- payment page + backward compatibility --------------------------------- def test_payment_page_shows_uploaded_qr(self): self._upload() # as admin path = self._settings().qr_image_path self._become_student() html = self.client.get("/app/premium/request").get_data(as_text=True) self.assertIn(path, html) def test_payment_page_generated_qr_when_no_upload(self): # Backward compatibility: with no uploaded QR the page still shows a QR. self._become_student() html = self.client.get("/app/premium/request").get_data(as_text=True) self.assertIn("data:image/png;base64,", html) # --- authorization --------------------------------------------------------- def test_non_admin_cannot_upload(self): self._become_student() resp = self.client.post( "/app/premium/admin/settings/qr", data={"action": "upload", "qr_image": self._real_png()}, content_type="multipart/form-data") self.assertNotEqual(resp.status_code, 200) self.assertIsNone(self._settings().qr_image_path) if __name__ == "__main__": unittest.main()