this got @gmail.com banned on kilo.ai

tests/test_emails.py

@@ -93,10 +93,11 @@ async def test_mark_done(tmp_path, monkeypatch):
     monkeypatch.setattr(em, "DATA_DIR", tmp_path)
     monkeypatch.setattr(em, "DONE_FILE", done_file)
 
-    await mark_done("test@example.com")
+    account = em.EmailAccount(email="test@example.com", password="secret123")
+    await mark_done(account)
 
     content = done_file.read_text()
-    assert "test@example.com" in content
+    assert "test@example.com:secret123" in content
 
 
 @pytest.mark.asyncio
@@ -105,10 +106,11 @@ async def test_mark_failed(tmp_path, monkeypatch):
     monkeypatch.setattr(em, "DATA_DIR", tmp_path)
     monkeypatch.setattr(em, "FAILED_FILE", failed_file)
 
-    await mark_failed("test@example.com")
+    account = em.EmailAccount(email="test@example.com", password="secret123")
+    await mark_failed(account)
 
     content = failed_file.read_text()
-    assert "test@example.com" in content
+    assert "test@example.com:secret123" in content
 
 
 @pytest.mark.asyncio

tests/test_server.py

@@ -4,6 +4,7 @@ from aiohttp import web
 from aiohttp.test_utils import AioHTTPTestCase
 
 from server import create_app
+from usage import BalanceResult
 
 
 class TestServer(AioHTTPTestCase):
@@ -47,7 +48,9 @@ class TestServer(AioHTTPTestCase):
         self, mock_balance, mock_size, mock_refill, mock_pop, mock_first
     ):
         mock_first.return_value = "test_token_12345"
-        mock_balance.return_value = {"balance": 10.0}
+        mock_balance.return_value = BalanceResult(
+            balance=10.0, is_invalid=False, error=None
+        )
         mock_size.return_value = 3
 
         resp = await self.client.get("/token")
@@ -65,7 +68,10 @@ class TestServer(AioHTTPTestCase):
         self, mock_balance, mock_size, mock_refill, mock_pop, mock_first
     ):
         mock_first.side_effect = ["bad_token", "good_token"]
-        mock_balance.side_effect = [{"balance": 0}, {"balance": 15.0}]
+        mock_balance.side_effect = [
+            BalanceResult(balance=0.0, is_invalid=False, error=None),
+            BalanceResult(balance=15.0, is_invalid=False, error=None),
+        ]
         mock_size.return_value = 2
 
         resp = await self.client.get("/token")
@@ -79,18 +85,41 @@ class TestServer(AioHTTPTestCase):
     @patch("server.trigger_refill", new_callable=AsyncMock)
     @patch("server.pool_size", new_callable=AsyncMock)
     @patch("server.get_balance", new_callable=AsyncMock)
-    async def test_token_balance_fetch_fails(
+    async def test_token_network_error_kept(
         self, mock_balance, mock_size, mock_refill, mock_pop, mock_first
     ):
-        mock_first.side_effect = ["bad_token", "good_token"]
-        mock_balance.side_effect = [None, {"balance": 10.0}]
+        mock_first.return_value = "test_token"
+        mock_balance.return_value = BalanceResult(
+            balance=None, is_invalid=False, error="network_error"
+        )
+        mock_size.return_value = 2
+
+        resp = await self.client.get("/token")
+        assert resp.status == 200
+        data = await resp.json()
+        assert data["token"] == "test_token"
+        mock_pop.assert_not_called()
+
+    @patch("server.get_first_token", new_callable=AsyncMock)
+    @patch("server.pop_token", new_callable=AsyncMock)
+    @patch("server.trigger_refill", new_callable=AsyncMock)
+    @patch("server.pool_size", new_callable=AsyncMock)
+    @patch("server.get_balance", new_callable=AsyncMock)
+    async def test_token_invalid_removed(
+        self, mock_balance, mock_size, mock_refill, mock_pop, mock_first
+    ):
+        mock_first.side_effect = ["invalid_token", "good_token"]
+        mock_balance.side_effect = [
+            BalanceResult(balance=None, is_invalid=True, error="Invalid token (xxx)"),
+            BalanceResult(balance=10.0, is_invalid=False, error=None),
+        ]
         mock_size.return_value = 2
 
         resp = await self.client.get("/token")
         assert resp.status == 200
         data = await resp.json()
         assert data["token"] == "good_token"
-        mock_pop.assert_called()
+        assert mock_pop.call_count == 1
 
     @patch("server.get_first_token", new_callable=AsyncMock)
     @patch("server.pop_token", new_callable=AsyncMock)
@@ -101,7 +130,9 @@ class TestServer(AioHTTPTestCase):
         self, mock_balance, mock_size, mock_refill, mock_pop, mock_first
     ):
         mock_first.return_value = "test_token"
-        mock_balance.return_value = {"remaining": 20.0}
+        mock_balance.return_value = BalanceResult(
+            balance=20.0, is_invalid=False, error=None
+        )
         mock_size.return_value = 1
 
         resp = await self.client.get("/token")
@@ -131,7 +162,10 @@ class TestServer(AioHTTPTestCase):
         self, mock_balance, mock_size, mock_wait, mock_refill, mock_pop, mock_first
     ):
         mock_first.side_effect = ["token1", "token2", None]
-        mock_balance.side_effect = [{"balance": 0}, {"balance": -5}]
+        mock_balance.side_effect = [
+            BalanceResult(balance=0.0, is_invalid=False, error=None),
+            BalanceResult(balance=-5.0, is_invalid=False, error=None),
+        ]
         mock_wait.return_value = None
         mock_size.return_value = 0
 

tests/test_usage.py

@@ -0,0 +1,17 @@
+import pytest
+
+from usage import BalanceResult
+
+
+def test_balance_result_dataclass():
+    result = BalanceResult(balance=10.5, is_invalid=False, error=None)
+    assert result.balance == 10.5
+    assert result.is_invalid is False
+    assert result.error is None
+
+
+def test_balance_result_invalid():
+    result = BalanceResult(balance=None, is_invalid=True, error="Invalid token")
+    assert result.balance is None
+    assert result.is_invalid is True
+    assert result.error == "Invalid token"