add more tests, remove warning

blossomtune_gradio/util.py

@@ -22,7 +22,7 @@ def validate_email(email_address: str) -> bool:
     # DNS MX record validation
     try:
         domain = email_address.rsplit("@", 1)[-1]
-        dns.resolver.query(domain, "MX")
+        dns.resolver.resolve(domain, "MX")
         return True
     except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN, IndexError):
         return False

tests/test_federation.py

@@ -0,0 +1,253 @@
+import sqlite3
+import pytest
+from unittest.mock import MagicMock
+
+from blossomtune_gradio import federation as fed
+from blossomtune_gradio import database
+
+
+@pytest.fixture
+def in_memory_db(mocker):
+    """
+    Fixture to set up and tear down an in-memory SQLite database for tests.
+    It ensures that the same connection object is used for both schema
+    initialization and the test execution.
+    """
+    con = sqlite3.connect(":memory:")
+    mocker.patch("sqlite3.connect", return_value=con)
+    database.init()
+    yield con
+    con.close()
+
+
+@pytest.fixture
+def mock_settings(mocker):
+    """Fixture to mock the settings module."""
+    # The lambda returns a formatted string to simulate Jinja2's behavior
+    mock_get = MagicMock(
+        side_effect=lambda key, **kwargs: f"mock_{key}".format(**kwargs)
+    )
+    mocker.patch("blossomtune_gradio.federation.settings.get_text", mock_get)
+    return mock_get
+
+
+@pytest.fixture
+def mock_mail(mocker):
+    """Fixture to mock the mail module."""
+    return mocker.patch("blossomtune_gradio.mail.send_activation_email")
+
+
+def test_generate_participant_id():
+    """Test the generation of a participant ID."""
+    pid = fed.generate_participant_id()
+    assert isinstance(pid, str)
+    assert len(pid) == 6
+    assert pid.isalnum() and pid.isupper()
+
+
+def test_generate_activation_code():
+    """Test the generation of an activation code."""
+    code = fed.generate_activation_code()
+    assert isinstance(code, str)
+    assert len(code) == 8
+    assert code.isalnum() and code.isupper()
+
+
+class TestCheckParticipantStatus:
+    """Test suite for the check_participant_status function."""
+
+    def test_new_user_registration_success(
+        self, in_memory_db, mock_settings, mock_mail
+    ):
+        """Verify successful registration for a new user."""
+        mock_mail.return_value = (True, "")
+        success, message = fed.check_participant_status(
+            "new_user", "new@example.com", ""
+        )
+        assert success is True
+        assert message == "mock_registration_submitted_md"
+
+        # Verify the user was added to the database
+        cursor = in_memory_db.cursor()
+        cursor.execute("SELECT hf_handle FROM requests WHERE hf_handle = 'new_user'")
+        assert cursor.fetchone() is not None
+
+    def test_new_user_invalid_email(self, in_memory_db, mock_settings):
+        """Verify registration fails with an invalid email."""
+        success, message = fed.check_participant_status("user", "invalid-email", "")
+        assert success is False
+        assert message == "mock_invalid_email_md"
+
+    def test_new_user_federation_full(self, in_memory_db, mock_settings, mocker):
+        """Verify registration fails when the federation is full."""
+        mocker.patch("blossomtune_gradio.federation.cfg.MAX_NUM_NODES", 0)
+        success, message = fed.check_participant_status(
+            "another_user", "another@example.com", ""
+        )
+        assert success is False
+        assert message == "mock_federation_full_md"
+
+    def test_user_activation_success(self, in_memory_db, mock_settings):
+        """Verify a user can successfully activate their account."""
+        # Setup: Add a pending, non-activated user
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
+            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
+        )
+        in_memory_db.commit()
+
+        success, message = fed.check_participant_status(
+            "test_user", "test@example.com", "ABCDEF"
+        )
+        assert success is True
+        assert message == "mock_activation_successful_md"
+        # Verify the user is now activated
+        cursor.execute(
+            "SELECT is_activated FROM requests WHERE hf_handle = 'test_user'"
+        )
+        assert cursor.fetchone()[0] == 1
+
+    def test_user_activation_invalid_code(self, in_memory_db, mock_settings):
+        """Verify activation fails with an invalid code."""
+        # Setup: Add a pending, non-activated user
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
+            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
+        )
+        in_memory_db.commit()
+
+        success, message = fed.check_participant_status(
+            "test_user", "test@example.com", "WRONGCODE"
+        )
+        assert success is False
+        assert message == "mock_activation_invalid_md"
+
+    def test_status_check_approved(self, in_memory_db, mock_settings):
+        """Verify the status check for an approved user."""
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated, partition_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+            (
+                "PID456",
+                "approved",
+                "now",
+                "approved_user",
+                "approved@example.com",
+                "GHIJKL",
+                1,
+                5,
+            ),
+        )
+        in_memory_db.commit()
+        success, message = fed.check_participant_status(
+            "approved_user", "approved@example.com", "GHIJKL"
+        )
+        assert success is True
+        assert "mock_status_approved_md" in message
+
+
+class TestManageRequest:
+    """Test suite for the manage_request function."""
+
+    def test_approve_success(self, in_memory_db):
+        """Verify successful approval of a participant."""
+        # Setup: Add an activated, pending user
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
+            (
+                "PENDING1",
+                "pending",
+                "now",
+                "pending_user",
+                "pending@example.com",
+                "CODE",
+                1,
+            ),
+        )
+        in_memory_db.commit()
+
+        success, message = fed.manage_request("PENDING1", "10", "approve")
+        assert success is True
+        assert "is allowed to join" in message
+
+        # Verify status in DB
+        cursor.execute(
+            "SELECT status, partition_id FROM requests WHERE participant_id = 'PENDING1'"
+        )
+        status, partition_id = cursor.fetchone()
+        assert status == "approved"
+        assert partition_id == 10
+
+    def test_approve_not_activated(self, in_memory_db, mock_settings):
+        """Verify approval fails if the user is not activated."""
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
+            (
+                "PENDING2",
+                "pending",
+                "now",
+                "pending_user2",
+                "pending2@example.com",
+                "CODE",
+                0,
+            ),
+        )
+        in_memory_db.commit()
+        success, message = fed.manage_request("PENDING2", "11", "approve")
+        assert success is False
+        assert message == "mock_participant_not_activated_warning_md"
+
+    def test_deny_success(self, in_memory_db):
+        """Verify successful denial of a participant."""
+        cursor = in_memory_db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
+            (
+                "PENDING3",
+                "pending",
+                "now",
+                "pending_user3",
+                "pending3@example.com",
+                "CODE",
+                1,
+            ),
+        )
+        in_memory_db.commit()
+        success, message = fed.manage_request("PENDING3", "", "deny")
+        assert success is True
+        assert "is not allowed to join" in message
+
+        # Verify status in DB
+        cursor.execute("SELECT status FROM requests WHERE participant_id = 'PENDING3'")
+        assert cursor.fetchone()[0] == "denied"
+
+
+def test_get_next_partition_id(in_memory_db):
+    """Verify the logic for finding the next available partition ID."""
+    cursor = in_memory_db.cursor()
+    # No approved users yet
+    assert fed.get_next_partion_id() == 0
+
+    # Add some approved users with assigned partitions, including the required timestamp
+    cursor.execute(
+        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
+        ("P1", "approved", "now", 0),
+    )
+    cursor.execute(
+        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
+        ("P2", "approved", "now", 1),
+    )
+    in_memory_db.commit()
+    assert fed.get_next_partion_id() == 2
+
+    # Add another user, skipping a partition ID
+    cursor.execute(
+        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
+        ("P3", "approved", "now", 3),
+    )
+    in_memory_db.commit()
+    assert fed.get_next_partion_id() == 2

tests/test_util.py

@@ -9,7 +9,7 @@ def test_validate_email_valid(monkeypatch):
     """Tests a syntactically valid email with an existing MX record."""
     # Mock the dns.resolver.query to return a successful result.
     mock_query = MagicMock()
-    monkeypatch.setattr(dns.resolver, "query", mock_query)
+    monkeypatch.setattr(dns.resolver, "resolve", mock_query)
 
     email = "test@google.com"
     assert validate_email(email) is True
@@ -28,7 +28,7 @@ def test_validate_email_no_mx_record(monkeypatch):
     """Tests a domain that exists but has no MX record."""
     # Mock the dns.resolver.query to raise a NoAnswer exception.
     mock_query = MagicMock(side_effect=dns.resolver.NoAnswer)
-    monkeypatch.setattr(dns.resolver, "query", mock_query)
+    monkeypatch.setattr(dns.resolver, "resolve", mock_query)
 
     email = "user@example.com"
     assert validate_email(email) is False
@@ -39,7 +39,7 @@ def test_validate_email_non_existent_domain(monkeypatch):
     """Tests a domain that does not exist."""
     # Mock the dns.resolver.query to raise an NXDOMAIN exception.
     mock_query = MagicMock(side_effect=dns.resolver.NXDOMAIN)
-    monkeypatch.setattr(dns.resolver, "query", mock_query)
+    monkeypatch.setattr(dns.resolver, "resolve", mock_query)
 
     email = "user@not-a-real-domain-123.com"
     assert validate_email(email) is False