split-trainer/server/tests/test_store.py

"""Тесты хранилища SQLite: атомарность одноразовых токенов, версии состояния, миграция, чистка просроченного."""
import os
import sqlite3
import tempfile
import threading
import unittest

from splitapi.store import Store

NOW = 1_000_000


class StoreTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.path = os.path.join(self.tmp.name, "t.sqlite3")
        self.store = Store(self.path)
        self.addCleanup(self.tmp.cleanup)

    def make_user(self, email="u@example.com", pw_hash="hash"):
        uid = self.store.create_user(email, pw_hash, NOW)
        self.assertIsNotNone(uid)
        return uid


class UserTests(StoreTestCase):
    def test_create_and_lookup(self):
        uid = self.make_user()
        user = self.store.user_by_email("u@example.com")
        self.assertEqual((user["id"], user["pw_hash"]), (uid, "hash"))
        self.assertIsNone(self.store.user_by_email("none@example.com"))

    def test_duplicate_email_returns_none(self):
        self.make_user()
        self.assertIsNone(self.store.create_user("u@example.com", "other", NOW))

    def test_set_password(self):
        uid = self.make_user()
        self.assertTrue(self.store.set_password(uid, "new"))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "new")

    def test_set_password_can_be_conditional_on_the_previous_hash(self):
        """Перехеширование при входе не должно затирать пароль, который сменили, пока считался новый хеш."""
        uid = self.make_user(pw_hash="old")
        self.assertTrue(self.store.set_password(uid, "upgraded", expect_pw_hash="old"))
        self.assertFalse(self.store.set_password(uid, "stale-upgrade", expect_pw_hash="old"))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "upgraded")


class RegistrationTests(StoreTestCase):
    def test_registration_stores_no_password(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        conn = sqlite3.connect(self.path)
        try:
            columns = [row[1] for row in conn.execute("PRAGMA table_info(registrations)")]
        finally:
            conn.close()
        self.assertNotIn("pw_hash", columns)

    def test_peek_does_not_consume(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        self.assertEqual(self.store.peek_registration("th1", NOW), "u@example.com")
        self.assertEqual(self.store.peek_registration("th1", NOW), "u@example.com")
        self.assertIsNone(self.store.peek_registration("nope", NOW))

    def test_expired_registration_is_invisible(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        self.assertIsNone(self.store.peek_registration("th1", NOW + 101))
        self.assertFalse(self.store.activate_registration("th1", "pw", NOW + 101))

    def test_activation_creates_user_once(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        self.assertTrue(self.store.activate_registration("th1", "pw-hash", NOW))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "pw-hash")
        self.assertFalse(self.store.activate_registration("th1", "other", NOW))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "pw-hash")

    def test_activation_fails_cleanly_if_the_address_was_taken(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        self.make_user("u@example.com", "existing")
        self.assertFalse(self.store.activate_registration("th1", "intruder", NOW))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "existing")
        self.assertIsNone(self.store.peek_registration("th1", NOW))          # ссылка погашена

    def test_new_registration_replaces_previous_for_same_email(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        self.store.save_registration("u@example.com", "th2", NOW + 100, NOW)
        self.assertIsNone(self.store.peek_registration("th1", NOW))
        self.assertEqual(self.store.peek_registration("th2", NOW), "u@example.com")

    def test_concurrent_activations_create_exactly_one_user(self):
        self.store.save_registration("u@example.com", "th1", NOW + 100, NOW)
        results, gate = [], threading.Barrier(8)

        def worker(i):
            gate.wait()
            results.append(self.store.activate_registration("th1", f"hash{i}", NOW))

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.assertEqual(results.count(True), 1)
        self.assertEqual(self.store.count("users"), 1)


class ResetTokenTests(StoreTestCase):
    def test_peek_does_not_consume(self):
        uid = self.make_user()
        self.store.save_reset_token(uid, "r1", NOW + 100)
        self.assertEqual(self.store.peek_reset_token("r1", NOW)["email"], "u@example.com")
        self.assertEqual(self.store.peek_reset_token("r1", NOW)["id"], uid)
        self.assertIsNone(self.store.peek_reset_token("nope", NOW))

    def test_reset_password_is_atomic_and_single_use(self):
        uid = self.make_user()
        self.store.create_session(uid, "s1", NOW + 1000, NOW)
        self.store.save_reset_token(uid, "r1", NOW + 100)
        result = self.store.reset_password("r1", "new-hash", NOW)
        self.assertEqual((result["id"], result["email"]), (uid, "u@example.com"))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "new-hash")
        self.assertIsNone(self.store.session_user("s1", NOW))                  # сессии завершены в той же транзакции
        self.assertIsNone(self.store.reset_password("r1", "again", NOW))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "new-hash")

    def test_expired_token_changes_nothing(self):
        uid = self.make_user()
        self.store.create_session(uid, "s1", NOW + 1000, NOW)
        self.store.save_reset_token(uid, "r1", NOW + 100)
        self.assertIsNone(self.store.peek_reset_token("r1", NOW + 101))
        self.assertIsNone(self.store.reset_password("r1", "new-hash", NOW + 101))
        self.assertEqual(self.store.user_by_email("u@example.com")["pw_hash"], "hash")
        self.assertIsNotNone(self.store.session_user("s1", NOW))

    def test_new_token_invalidates_previous(self):
        uid = self.make_user()
        self.store.save_reset_token(uid, "r1", NOW + 100)
        self.store.save_reset_token(uid, "r2", NOW + 100)
        self.assertIsNone(self.store.reset_password("r1", "x", NOW))
        self.assertIsNotNone(self.store.reset_password("r2", "x", NOW))


class SessionTests(StoreTestCase):
    def test_session_lookup_and_expiry(self):
        uid = self.make_user()
        self.assertTrue(self.store.create_session(uid, "s1", NOW + 100, NOW))
        self.assertEqual(self.store.session_user("s1", NOW)["email"], "u@example.com")
        self.assertIsNone(self.store.session_user("s1", NOW + 101))
        self.assertIsNone(self.store.session_user("nope", NOW))

    def test_session_is_refused_if_password_changed_meanwhile(self):
        """Гонка вход/сброс: пароль сменили между проверкой старого пароля и созданием сессии."""
        uid = self.make_user(pw_hash="old")
        self.store.set_password(uid, "new")
        self.assertFalse(self.store.create_session(uid, "s1", NOW + 100, NOW, expect_pw_hash="old"))
        self.assertIsNone(self.store.session_user("s1", NOW))
        self.assertTrue(self.store.create_session(uid, "s2", NOW + 100, NOW, expect_pw_hash="new"))

    def test_delete_session_and_all_sessions(self):
        uid = self.make_user()
        for name in ("s1", "s2"):
            self.store.create_session(uid, name, NOW + 100, NOW)
        self.store.delete_session("s1")
        self.assertIsNone(self.store.session_user("s1", NOW))
        self.assertIsNotNone(self.store.session_user("s2", NOW))
        self.store.delete_user_sessions(uid)
        self.assertIsNone(self.store.session_user("s2", NOW))

    def test_only_latest_sessions_are_kept(self):
        uid = self.make_user()
        for i in range(15):
            self.store.create_session(uid, f"s{i}", NOW + 1000, NOW + i, keep=10)
        alive = [i for i in range(15) if self.store.session_user(f"s{i}", NOW + 20)]
        self.assertEqual(alive, list(range(5, 15)))


class StateTests(StoreTestCase):
    def test_first_write_needs_version_zero(self):
        uid = self.make_user()
        self.assertIsNone(self.store.get_state(uid))
        self.assertEqual(self.store.put_state(uid, 0, '{"a":1}', NOW), 1)
        self.assertEqual(self.store.get_state(uid), (1, '{"a":1}'))

    def test_stale_version_is_a_conflict(self):
        uid = self.make_user()
        self.store.put_state(uid, 0, "{}", NOW)
        self.assertIsNone(self.store.put_state(uid, 0, '{"x":1}', NOW))
        self.assertEqual(self.store.put_state(uid, 1, '{"x":1}', NOW), 2)
        self.assertEqual(self.store.get_state(uid), (2, '{"x":1}'))

    def test_concurrent_writes_with_one_version_let_exactly_one_through(self):
        uid = self.make_user()
        results, gate = [], threading.Barrier(8)

        def worker(i):
            gate.wait()
            results.append(self.store.put_state(uid, 0, f'{{"w":{i}}}', NOW))

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.assertEqual(len([r for r in results if r is not None]), 1)

    def test_state_is_per_user(self):
        a, b = self.make_user("a@example.com"), self.make_user("b@example.com")
        self.store.put_state(a, 0, '{"who":"a"}', NOW)
        self.assertIsNone(self.store.get_state(b))

    def test_state_is_deleted_with_user(self):
        uid = self.make_user()
        self.store.put_state(uid, 0, "{}", NOW)
        self.store.create_session(uid, "s1", NOW + 100, NOW)
        self.store.delete_user(uid)
        self.assertIsNone(self.store.get_state(uid))
        self.assertIsNone(self.store.user_by_email("u@example.com"))
        self.assertIsNone(self.store.session_user("s1", NOW))


class PurgeTests(StoreTestCase):
    def test_purge_removes_only_expired(self):
        uid = self.make_user()
        self.store.save_registration("old@example.com", "reg-old", NOW - 1, NOW - 100)
        self.store.save_registration("new@example.com", "reg-new", NOW + 100, NOW)
        self.store.save_reset_token(uid, "rt-old", NOW - 1)
        self.store.create_session(uid, "s-old", NOW - 1, NOW - 100)
        self.store.create_session(uid, "s-new", NOW + 100, NOW)
        self.store.purge_expired(NOW)
        self.assertEqual(self.store.count("registrations"), 1)
        self.assertEqual(self.store.count("reset_tokens"), 0)
        self.assertEqual(self.store.count("sessions"), 1)


class SchemaTests(StoreTestCase):
    def test_expiry_columns_are_indexed(self):
        conn = sqlite3.connect(self.path)
        try:
            indexed = {row[0] for row in conn.execute("SELECT tbl_name FROM sqlite_master WHERE type='index' AND name LIKE '%expires%'")}
        finally:
            conn.close()
        self.assertEqual(indexed, {"registrations", "reset_tokens", "sessions"})

    def test_old_schema_with_password_column_is_migrated(self):
        old = os.path.join(self.tmp.name, "old.sqlite3")
        conn = sqlite3.connect(old)
        conn.executescript(
            "CREATE TABLE registrations (email TEXT PRIMARY KEY, pw_hash TEXT NOT NULL, token_hash TEXT NOT NULL UNIQUE, "
            "expires_at INTEGER NOT NULL, created_at INTEGER NOT NULL);"
            "INSERT INTO registrations VALUES ('x@example.com', 'h', 't', 99999999999, 1);"
        )
        conn.commit()
        conn.close()
        store = Store(old)                                      # открытие мигрирует схему
        store.save_registration("y@example.com", "th9", NOW + 100, NOW)
        self.assertEqual(store.peek_registration("th9", NOW), "y@example.com")
        self.assertIsNone(store.peek_registration("t", NOW))    # старые ожидающие регистрации с паролем отброшены


if __name__ == "__main__":
    unittest.main()