pdf-quiz-generator/backend/tests/test_article_ai.py

"""AI authoring endpoints on disposable SQLite; no network, no model calls."""
import json
import re
import sys
import unittest
from unittest.mock import Mock, patch

import test_quiz_builder as fixtures
from app.models.article import Article
from app.models.flashcard import Flashcard, FlashcardDeck, FlashcardArticleLink
from app.models.question import Question
from app.routers import articles
from app.tasks.quiz_tasks import generate_article_draft, generate_article_cards
from sqlalchemy.orm import sessionmaker

DRAFT_RESPONSE = {
    "title": "AI draft topic", "slug": "ai-draft-topic", "summary": "Draft summary",
    "content": "Draft intro", "sections": [
        {"id": "a" * 32, "slug": "first", "title": "First", "content": "Body"},
        {"id": "b" * 32, "slug": "second", "title": "Second", "content": "More"},
    ],
}


class ArticleAiTests(unittest.TestCase):
    def setUp(self):
        self.bank = fixtures.BuilderTests()
        self.bank.setUp()
        self.client = self.bank.client
        self.client.app.include_router(articles.router, prefix='/articles')
        self.redis = Mock()
        self.redis.from_url.return_value = self.redis
        self.redis.get.return_value = None
        self.redis.lrange.return_value = []
        self.redis.lpush.return_value = 1
        self.redis.expire.return_value = True
        self.redis.incr.side_effect = [1, 21]
        self.delay = patch.dict(sys.modules, {'redis': self.redis})
        self.delay.start()
        self.queue = patch('app.tasks.quiz_tasks.generate_article_draft.delay').start()
        self.card_queue = patch('app.tasks.quiz_tasks.generate_article_cards.delay').start()

    def tearDown(self):
        patch.stopall()
        self.bank.tearDown()

    def publish(self, article_id):
        return self.client.post(f'/articles/{article_id}/publish', json={'published': True})

    def test_ai_endpoints_queue_and_poll(self):
        self.bank.user = self.bank.mod
        response = self.client.post('/articles/ai-draft', json={'topic': 'Neonatal jaundice', 'instructions': 'Two sections'})
        self.assertEqual(response.status_code, 200, response.text)
        job_id = response.json()['job_id']
        self.queue.assert_called_once()
        self.redis.get.return_value = 'completed'
        poll = self.client.get(f'/articles/job/{job_id}').json()
        self.assertEqual(poll['status'], 'completed')
        self.bank.user = self.bank.owner
        self.assertEqual(self.client.get(f'/articles/job/{job_id}').status_code, 404)
        self.bank.user = self.bank.mod
        article = self.client.post('/articles/', json={
            "title": "Refine me", "slug": "refine-me", "content": "Intro", "sections": []}).json()
        self.assertEqual(self.client.post(f"/articles/{article['id']}/ai-refine", json={'instructions': 'Shorten'}).status_code, 200)
        self.assertEqual(self.client.post(f"/articles/{article['id']}/ai-cards").status_code, 200)
        self.card_queue.assert_called_once()
        self.bank.user = self.bank.owner
        self.assertEqual(self.client.post('/articles/ai-draft', json={'topic': 'Learner'}).status_code, 403)
        self.assertEqual(self.client.post(f"/articles/{article['id']}/ai-refine", json={'instructions': 'x'}).status_code, 403)

    def test_draft_task_creates_and_refines_without_publishing(self):
        self.redis.set.return_value = True
        ai = patch('app.services.ai_service.chat').start()
        ai.return_value = json.dumps(DRAFT_RESPONSE)
        with patch('app.tasks.quiz_tasks.SessionLocal', sessionmaker(bind=self.bank.engine)), \
             patch('app.services.ai_service.get_model_for_task', return_value=('synthetic', None)):
            generate_article_draft('job-1', 3, 'Neonatal jaundice', '')
            article = self.bank.db.query(Article).filter_by(slug='ai-draft-topic').one()
            self.assertEqual(article.status, 'draft')
            self.assertEqual([s['id'] for s in article.sections], ['a' * 32, 'b' * 32])
            self.redis.set.assert_any_call('extraction:status:job-1', 'completed', ex=3600)
            # Refine must send the existing body to the model, not just the title.
            article.content = 'Unique draft body to preserve'
            article.sections = [{'id': 'c' * 32, 'slug': 'kept', 'title': 'Kept', 'content': 'Kept body'}]
            self.bank.db.commit()
            generate_article_draft('job-2', 3, 'Neonatal jaundice', 'Shorten', article.id)
            prompt = ai.call_args.kwargs['messages'][0]['content']
            self.assertIn('Unique draft body to preserve', prompt)
            self.assertIn('Kept body', prompt)
            refreshed = self.bank.db.query(Article).filter_by(id=article.id).one()
            self.assertEqual(refreshed.status, 'draft')
            self.assertEqual(refreshed.slug, 'ai-draft-topic')
            # Invalid model section ids are replaced with valid hex ids.
            bad = dict(DRAFT_RESPONSE, sections=[{'id': 'bad-id', 'slug': 'bad', 'title': 'Bad', 'content': ''}])
            ai.return_value = json.dumps(bad)
            generate_article_draft('job-3', 3, 'Neonatal jaundice', 'Again', article.id)
            fixed = self.bank.db.query(Article).filter_by(id=article.id).one()
            self.assertTrue(all(re.fullmatch(r'[0-9a-f]{32}', s['id']) for s in fixed.sections), fixed.sections)
        patch.stopall()
        self.redis.reset_mock()

    def test_cards_task_builds_private_linked_deck(self):
        self.redis.set.return_value = True
        self.bank.db.add(Article(slug='cards-source', title='Cards source', content='Body',
                                 sections=[{'id': 'c' * 32, 'slug': 's', 'title': 'S', 'content': 'More'}],
                                 user_id=3, status='published'))
        self.bank.db.commit()
        article = self.bank.db.query(Article).filter_by(slug='cards-source').one()
        with patch('app.tasks.quiz_tasks.SessionLocal', sessionmaker(bind=self.bank.engine)), \
             patch('app.services.extraction_modes.generate_flashcards',
                   return_value=[{'front': 'F1', 'back': 'B1'}, {'front': 'F2', 'back': 'B2', 'page_reference': 3}]), \
             patch('app.services.ai_service.get_model_for_task', return_value=('synthetic', None)):
            generate_article_cards('job-3', 3, article.id)
            deck = self.bank.db.query(FlashcardDeck).filter_by(title='Cards: Cards source').one()
            self.assertEqual(deck.is_shared, 0)
            cards = self.bank.db.query(Flashcard).filter_by(deck_id=deck.id).all()
            self.assertEqual(len(cards), 2)
            self.assertEqual(deck.card_count, 2)
            linked = self.bank.db.query(FlashcardArticleLink).filter_by(article_id=article.id).count()
            self.assertEqual(linked, 2)
            self.redis.set.assert_any_call('extraction:status:job-3', 'completed', ex=3600)
        patch.stopall()
        self.redis.reset_mock()


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