worker test

tests/test_worker_pool.py

@@ -1364,5 +1364,604 @@ class TestErrorPatternCoverage:
             assert result in [True, False]
 
 
+# =============================================================================
+# 15. INTEGRATION TESTS - Real Database
+# =============================================================================
+
+import os
+import tempfile
+from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
+from sqlalchemy import Column, Integer, String, DateTime, JSON, Text, Boolean, select, update
+from sqlalchemy.orm import declarative_base
+
+# Create a test-specific Base for integration tests
+IntegrationBase = declarative_base()
+
+
+class TestJobModel(IntegrationBase):
+    """Real job model for integration tests."""
+    __tablename__ = "test_jobs"
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    job_id = Column(String(100), unique=True, index=True, nullable=False)
+    user_id = Column(String(50), index=True, nullable=False)
+    job_type = Column(String(20), index=True, nullable=False)
+    status = Column(String(20), default="queued", index=True)
+    priority = Column(String(10), default="fast", index=True)
+    next_process_at = Column(DateTime, nullable=True, index=True)
+    retry_count = Column(Integer, default=0)
+    third_party_id = Column(String(255), nullable=True)
+    input_data = Column(JSON, nullable=True)
+    output_data = Column(JSON, nullable=True)
+    error_message = Column(Text, nullable=True)
+    created_at = Column(DateTime, nullable=False)
+    started_at = Column(DateTime, nullable=True)
+    completed_at = Column(DateTime, nullable=True)
+    credits_reserved = Column(Integer, default=0)
+    credits_refunded = Column(Boolean, default=False)
+
+
+class TestUserModel(IntegrationBase):
+    """Real user model for integration tests."""
+    __tablename__ = "test_users"
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    user_id = Column(String(50), unique=True, index=True, nullable=False)
+    email = Column(String(255), unique=True, nullable=False)
+    credits = Column(Integer, default=100)
+
+
+@pytest.fixture
+async def integration_db():
+    """Create a real SQLite database for integration tests."""
+    # Use a temp file for the test database
+    fd, db_path = tempfile.mkstemp(suffix=".db")
+    os.close(fd)
+    
+    db_url = f"sqlite+aiosqlite:///{db_path}"
+    engine = create_async_engine(db_url, echo=False)
+    
+    # Create tables
+    async with engine.begin() as conn:
+        await conn.run_sync(IntegrationBase.metadata.create_all)
+    
+    session_maker = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+    
+    yield {
+        "engine": engine,
+        "session_maker": session_maker,
+        "db_url": db_url,
+        "db_path": db_path
+    }
+    
+    # Cleanup
+    await engine.dispose()
+    if os.path.exists(db_path):
+        os.remove(db_path)
+
+
+class TestIntegrationRealDatabase:
+    """Integration tests using real SQLite database."""
+    
+    @pytest.mark.asyncio
+    async def test_create_and_query_job(self, integration_db):
+        """Test creating and querying a job in real database."""
+        session_maker = integration_db["session_maker"]
+        
+        async with session_maker() as session:
+            # Create a job
+            job = TestJobModel(
+                job_id="int-test-1",
+                user_id="user-123",
+                job_type="text",
+                status="queued",
+                priority="fast",
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+            
+            # Query it back
+            result = await session.execute(
+                select(TestJobModel).where(TestJobModel.job_id == "int-test-1")
+            )
+            fetched = result.scalar_one_or_none()
+            
+            assert fetched is not None
+            assert fetched.job_id == "int-test-1"
+            assert fetched.status == "queued"
+    
+    @pytest.mark.asyncio
+    async def test_atomic_update_with_where_clause(self, integration_db):
+        """Test atomic UPDATE with WHERE clause (real DB concurrency test)."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create a job
+        async with session_maker() as session:
+            job = TestJobModel(
+                job_id="atomic-test-1",
+                user_id="user-123",
+                job_type="text",
+                status="queued",
+                priority="fast",
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+        
+        # Simulate Worker 1 claiming the job
+        async with session_maker() as session1:
+            stmt = (
+                update(TestJobModel)
+                .where(
+                    TestJobModel.job_id == "atomic-test-1",
+                    TestJobModel.status == "queued"
+                )
+                .values(status="processing", started_at=datetime.utcnow())
+            )
+            result1 = await session1.execute(stmt)
+            await session1.commit()
+            
+            # Worker 1 should succeed
+            assert result1.rowcount == 1
+        
+        # Simulate Worker 2 trying to claim the same job
+        async with session_maker() as session2:
+            stmt = (
+                update(TestJobModel)
+                .where(
+                    TestJobModel.job_id == "atomic-test-1",
+                    TestJobModel.status == "queued"  # This won't match!
+                )
+                .values(status="processing", started_at=datetime.utcnow())
+            )
+            result2 = await session2.execute(stmt)
+            await session2.commit()
+            
+            # Worker 2 should fail (rowcount = 0)
+            assert result2.rowcount == 0, "Second worker should not be able to claim!"
+    
+    @pytest.mark.asyncio
+    async def test_job_ordering_by_created_at(self, integration_db):
+        """Test jobs are queried in FIFO order by created_at."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create jobs in non-chronological order
+        async with session_maker() as session:
+            job3 = TestJobModel(
+                job_id="order-3", user_id="user", job_type="text",
+                status="queued", priority="fast",
+                created_at=datetime(2024, 1, 1, 12, 0, 0)  # Latest
+            )
+            job1 = TestJobModel(
+                job_id="order-1", user_id="user", job_type="text",
+                status="queued", priority="fast",
+                created_at=datetime(2024, 1, 1, 10, 0, 0)  # Oldest
+            )
+            job2 = TestJobModel(
+                job_id="order-2", user_id="user", job_type="text",
+                status="queued", priority="fast",
+                created_at=datetime(2024, 1, 1, 11, 0, 0)  # Middle
+            )
+            session.add_all([job3, job1, job2])
+            await session.commit()
+        
+        # Query with ORDER BY created_at
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel)
+                .where(TestJobModel.status == "queued")
+                .order_by(TestJobModel.created_at)
+                .limit(3)
+            )
+            jobs = result.scalars().all()
+            
+            assert len(jobs) == 3
+            assert jobs[0].job_id == "order-1", "Oldest job should be first"
+            assert jobs[1].job_id == "order-2"
+            assert jobs[2].job_id == "order-3"
+    
+    @pytest.mark.asyncio
+    async def test_priority_filtering(self, integration_db):
+        """Test that priority filter works correctly in real DB."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create jobs with different priorities
+        async with session_maker() as session:
+            fast_job = TestJobModel(
+                job_id="prio-fast", user_id="user", job_type="text",
+                status="queued", priority="fast",
+                created_at=datetime.utcnow()
+            )
+            medium_job = TestJobModel(
+                job_id="prio-medium", user_id="user", job_type="image",
+                status="queued", priority="medium",
+                created_at=datetime.utcnow()
+            )
+            slow_job = TestJobModel(
+                job_id="prio-slow", user_id="user", job_type="video",
+                status="queued", priority="slow",
+                created_at=datetime.utcnow()
+            )
+            session.add_all([fast_job, medium_job, slow_job])
+            await session.commit()
+        
+        # Query only fast priority
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel).where(TestJobModel.priority == "fast")
+            )
+            fast_jobs = result.scalars().all()
+            
+            assert len(fast_jobs) == 1
+            assert fast_jobs[0].job_id == "prio-fast"
+
+
+class TestIntegrationConcurrency:
+    """Test concurrent access patterns."""
+    
+    @pytest.mark.asyncio
+    async def test_multiple_workers_different_jobs(self, integration_db):
+        """Multiple workers should process different jobs concurrently."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create multiple jobs
+        async with session_maker() as session:
+            for i in range(5):
+                job = TestJobModel(
+                    job_id=f"concurrent-{i}",
+                    user_id="user",
+                    job_type="text",
+                    status="queued",
+                    priority="fast",
+                    created_at=datetime.utcnow() + timedelta(seconds=i)
+                )
+                session.add(job)
+            await session.commit()
+        
+        # Simulate 3 workers claiming jobs concurrently
+        claimed_jobs = []
+        
+        async def worker_claim(worker_id):
+            async with session_maker() as session:
+                # SELECT the oldest unclaimed job
+                result = await session.execute(
+                    select(TestJobModel)
+                    .where(TestJobModel.status == "queued")
+                    .order_by(TestJobModel.created_at)
+                    .limit(1)
+                )
+                job = result.scalar_one_or_none()
+                
+                if job:
+                    # Try to claim it atomically
+                    stmt = (
+                        update(TestJobModel)
+                        .where(
+                            TestJobModel.job_id == job.job_id,
+                            TestJobModel.status == "queued"
+                        )
+                        .values(status="processing")
+                    )
+                    claim_result = await session.execute(stmt)
+                    await session.commit()
+                    
+                    if claim_result.rowcount == 1:
+                        claimed_jobs.append((worker_id, job.job_id))
+        
+        # Run workers concurrently
+        await asyncio.gather(
+            worker_claim(0),
+            worker_claim(1),
+            worker_claim(2)
+        )
+        
+        # Each worker should have claimed a different job (or failed gracefully)
+        claimed_job_ids = [job_id for _, job_id in claimed_jobs]
+        # No duplicates
+        assert len(claimed_job_ids) == len(set(claimed_job_ids)), "Same job claimed by multiple workers!"
+    
+    @pytest.mark.asyncio
+    async def test_next_process_at_prevents_duplicate_status_checks(self, integration_db):
+        """next_process_at should prevent multiple workers from checking same job."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create a processing job with future next_process_at
+        async with session_maker() as session:
+            job = TestJobModel(
+                job_id="processing-job",
+                user_id="user",
+                job_type="video",
+                status="processing",
+                priority="slow",
+                next_process_at=datetime.utcnow() + timedelta(minutes=5),  # Future
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+        
+        # Query should NOT return this job (next_process_at is in future)
+        async with session_maker() as session:
+            now = datetime.utcnow()
+            result = await session.execute(
+                select(TestJobModel).where(
+                    TestJobModel.status == "processing",
+                    TestJobModel.next_process_at <= now
+                )
+            )
+            jobs = result.scalars().all()
+            
+            assert len(jobs) == 0, "Job with future next_process_at should not be selected"
+
+
+class TestIntegrationEndToEnd:
+    """End-to-end job lifecycle tests."""
+    
+    @pytest.mark.asyncio
+    async def test_job_lifecycle_queued_to_completed(self, integration_db):
+        """Test full job lifecycle: queued → processing → completed."""
+        session_maker = integration_db["session_maker"]
+        
+        # 1. Create job
+        async with session_maker() as session:
+            job = TestJobModel(
+                job_id="lifecycle-1",
+                user_id="user-123",
+                job_type="text",
+                status="queued",
+                priority="fast",
+                input_data={"prompt": "Hello"},
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+        
+        # 2. Worker claims job (queued → processing)
+        async with session_maker() as session:
+            stmt = (
+                update(TestJobModel)
+                .where(
+                    TestJobModel.job_id == "lifecycle-1",
+                    TestJobModel.status == "queued"
+                )
+                .values(status="processing", started_at=datetime.utcnow())
+            )
+            result = await session.execute(stmt)
+            await session.commit()
+            assert result.rowcount == 1
+        
+        # 3. Worker completes job (processing → completed)
+        async with session_maker() as session:
+            stmt = (
+                update(TestJobModel)
+                .where(TestJobModel.job_id == "lifecycle-1")
+                .values(
+                    status="completed",
+                    output_data={"response": "Hello back!"},
+                    completed_at=datetime.utcnow()
+                )
+            )
+            await session.execute(stmt)
+            await session.commit()
+        
+        # 4. Verify final state
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel).where(TestJobModel.job_id == "lifecycle-1")
+            )
+            job = result.scalar_one()
+            
+            assert job.status == "completed"
+            assert job.output_data == {"response": "Hello back!"}
+            assert job.started_at is not None
+            assert job.completed_at is not None
+    
+    @pytest.mark.asyncio
+    async def test_job_lifecycle_queued_to_failed(self, integration_db):
+        """Test job failure lifecycle: queued → processing → failed."""
+        session_maker = integration_db["session_maker"]
+        
+        # 1. Create job
+        async with session_maker() as session:
+            job = TestJobModel(
+                job_id="lifecycle-fail",
+                user_id="user-123",
+                job_type="text",
+                status="queued",
+                priority="fast",
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+        
+        # 2. Start processing
+        async with session_maker() as session:
+            stmt = (
+                update(TestJobModel)
+                .where(TestJobModel.job_id == "lifecycle-fail")
+                .values(status="processing", started_at=datetime.utcnow())
+            )
+            await session.execute(stmt)
+            await session.commit()
+        
+        # 3. Job fails
+        async with session_maker() as session:
+            stmt = (
+                update(TestJobModel)
+                .where(TestJobModel.job_id == "lifecycle-fail")
+                .values(
+                    status="failed",
+                    error_message="Content blocked by safety filter",
+                    completed_at=datetime.utcnow()
+                )
+            )
+            await session.execute(stmt)
+            await session.commit()
+        
+        # 4. Verify
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel).where(TestJobModel.job_id == "lifecycle-fail")
+            )
+            job = result.scalar_one()
+            
+            assert job.status == "failed"
+            assert "safety" in job.error_message.lower()
+
+
+class TestIntegrationCreditSystem:
+    """Integration tests for credit system with real database."""
+    
+    @pytest.mark.asyncio
+    async def test_credit_reservation_and_refund(self, integration_db):
+        """Test credit reserve and refund with real database."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create user with credits
+        async with session_maker() as session:
+            user = TestUserModel(
+                user_id="credit-user",
+                email="test@example.com",
+                credits=100
+            )
+            session.add(user)
+            await session.commit()
+        
+        # Reserve credits
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestUserModel).where(TestUserModel.user_id == "credit-user")
+            )
+            user = result.scalar_one()
+            
+            # Deduct 10 credits
+            user.credits -= 10
+            await session.commit()
+        
+        # Verify deduction
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestUserModel).where(TestUserModel.user_id == "credit-user")
+            )
+            user = result.scalar_one()
+            assert user.credits == 90
+        
+        # Refund credits
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestUserModel).where(TestUserModel.user_id == "credit-user")
+            )
+            user = result.scalar_one()
+            user.credits += 10
+            await session.commit()
+        
+        # Verify refund
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestUserModel).where(TestUserModel.user_id == "credit-user")
+            )
+            user = result.scalar_one()
+            assert user.credits == 100
+
+
+class TestIntegrationOrphanedJobs:
+    """Test orphaned job handling (simulating server crash)."""
+    
+    @pytest.mark.asyncio
+    async def test_find_orphaned_processing_jobs(self, integration_db):
+        """Find jobs stuck in processing state (simulating crash recovery)."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create jobs in various states
+        async with session_maker() as session:
+            # This should NOT be found (queued)
+            job1 = TestJobModel(
+                job_id="orphan-queued", user_id="user", job_type="text",
+                status="queued", priority="fast",
+                credits_reserved=5, created_at=datetime.utcnow()
+            )
+            # This SHOULD be found (processing with credits)
+            job2 = TestJobModel(
+                job_id="orphan-processing", user_id="user", job_type="video",
+                status="processing", priority="slow",
+                credits_reserved=10, credits_refunded=False,
+                created_at=datetime.utcnow()
+            )
+            # This should NOT be found (already refunded)
+            job3 = TestJobModel(
+                job_id="orphan-refunded", user_id="user", job_type="video",
+                status="processing", priority="slow",
+                credits_reserved=0, credits_refunded=True,
+                created_at=datetime.utcnow()
+            )
+            session.add_all([job1, job2, job3])
+            await session.commit()
+        
+        # Query for orphaned jobs (as refund_orphaned_jobs does)
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel).where(
+                    TestJobModel.status == "processing",
+                    TestJobModel.credits_reserved > 0,
+                    TestJobModel.credits_refunded == False
+                )
+            )
+            orphaned = result.scalars().all()
+            
+            assert len(orphaned) == 1
+            assert orphaned[0].job_id == "orphan-processing"
+    
+    @pytest.mark.asyncio
+    async def test_mark_orphaned_job_as_failed(self, integration_db):
+        """Orphaned jobs should be marked as failed during recovery."""
+        session_maker = integration_db["session_maker"]
+        
+        # Create orphaned job
+        async with session_maker() as session:
+            job = TestJobModel(
+                job_id="crash-recovery",
+                user_id="user",
+                job_type="video",
+                status="processing",
+                priority="slow",
+                credits_reserved=5,
+                credits_refunded=False,
+                created_at=datetime.utcnow()
+            )
+            session.add(job)
+            await session.commit()
+        
+        # Simulate crash recovery - mark as failed
+        async with session_maker() as session:
+            stmt = (
+                update(TestJobModel)
+                .where(
+                    TestJobModel.job_id == "crash-recovery",
+                    TestJobModel.status == "processing"
+                )
+                .values(
+                    status="failed",
+                    error_message="Server shutdown during processing. Credits refunded.",
+                    credits_refunded=True,
+                    credits_reserved=0
+                )
+            )
+            await session.execute(stmt)
+            await session.commit()
+        
+        # Verify
+        async with session_maker() as session:
+            result = await session.execute(
+                select(TestJobModel).where(TestJobModel.job_id == "crash-recovery")
+            )
+            job = result.scalar_one()
+            
+            assert job.status == "failed"
+            assert job.credits_refunded == True
+            assert job.credits_reserved == 0
+            assert "shutdown" in job.error_message.lower()
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
+