Spaces:
Running
Running
| import unittest | |
| from sqlalchemy import create_engine | |
| from sqlalchemy.orm import sessionmaker | |
| from app.database import Base | |
| from app.models.user import User | |
| from app.models.process import Worker | |
| from app.routers.workers import get_my_worker | |
| class TestWorkerTokenGeneration(unittest.TestCase): | |
| def setUpClass(cls): | |
| cls.engine = create_engine("sqlite:///:memory:") | |
| Base.metadata.create_all(cls.engine) | |
| cls.Session = sessionmaker(bind=cls.engine) | |
| def test_new_user_gets_generated_token(self): | |
| db = self.Session() | |
| user = User(username="pedro_guitar", email="pedro@test.com", credits=10) | |
| db.add(user) | |
| db.commit() | |
| db.refresh(user) | |
| res = get_my_worker(db=db, current_user=user) | |
| self.assertIsNotNone(res["token"]) | |
| self.assertTrue(res["token"].startswith("melodix_wk_")) | |
| self.assertEqual(res["name"], f"Worker de {user.username}") | |
| db.close() | |
| def test_existing_worker_without_token_gets_token_assigned(self): | |
| db = self.Session() | |
| user = User(username="juan_piano", email="juan@test.com", credits=10) | |
| db.add(user) | |
| db.commit() | |
| db.refresh(user) | |
| # Simular worker creado en versión antigua con token = None | |
| legacy_worker = Worker( | |
| hostname="legacy_host_123", | |
| name="Worker Juan", | |
| status="offline", | |
| user_id=user.id, | |
| token=None | |
| ) | |
| db.add(legacy_worker) | |
| db.commit() | |
| res = get_my_worker(db=db, current_user=user) | |
| self.assertIsNotNone(res["token"]) | |
| self.assertTrue(res["token"].startswith("melodix_wk_")) | |
| db.close() | |
| def test_multiple_users_no_hostname_collision(self): | |
| db = self.Session() | |
| user1 = User(username="user1", email="u1@test.com", credits=5) | |
| user2 = User(username="user2", email="u2@test.com", credits=5) | |
| db.add_all([user1, user2]) | |
| db.commit() | |
| db.refresh(user1) | |
| db.refresh(user2) | |
| res1 = get_my_worker(db=db, current_user=user1) | |
| res2 = get_my_worker(db=db, current_user=user2) | |
| self.assertNotEqual(res1["token"], res2["token"]) | |
| self.assertNotEqual(res1["hostname"], res2["hostname"]) | |
| db.close() | |
| if __name__ == "__main__": | |
| unittest.main() | |