diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 113256c..9716a80 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,9 @@ name: CI on: push: - branches: [main] + branches: [main, dev] pull_request: - branches: [main] + branches: [main, dev] jobs: python: diff --git a/src/awaaz/api.py b/src/awaaz/api.py index 423a14a..c445ea8 100644 --- a/src/awaaz/api.py +++ b/src/awaaz/api.py @@ -159,31 +159,46 @@ async def upload_document( filename = Path((file.filename or "upload").replace("\\", "/")).name try: extension = validate_upload(filename) - stored = settings.uploads_dir / f"{uuid.uuid4()}{extension}" + except AwaazError as error: + raise _bad_request(error) from error + + stored = settings.uploads_dir / f"{uuid.uuid4()}{extension}" + cover_target = settings.uploads_dir / f"{uuid.uuid4()}.cover.jpg" + try: await save_upload(file, stored, settings.max_upload_bytes) text = await extract_text(stored) metadata = await extract_metadata(stored) cover_path = None - cover_target = settings.uploads_dir / f"{uuid.uuid4()}.cover.jpg" if await extract_cover(stored, cover_target): cover_path = str(cover_target) - except AwaazError as error: - raise _bad_request(error) from error - document = Document( - title=(title or metadata.get("title") or Path(filename).stem).strip(), - source_filename=filename, - text=text.strip(), - author=_join_authors(metadata.get("authors")), - series=metadata.get("series"), - tags=", ".join(metadata.get("tags", [])) if metadata.get("tags") else None, - cover_path=cover_path, - metadata_json=metadata, - word_count=len(text.split()), - ) - session.add(document) - await session.commit() - await session.refresh(document, ["collections"]) - return document + + document = Document( + title=(title or metadata.get("title") or Path(filename).stem).strip(), + source_filename=filename, + text=text.strip(), + author=_join_authors(metadata.get("authors")), + series=metadata.get("series"), + tags=", ".join(metadata.get("tags", [])) if metadata.get("tags") else None, + cover_path=cover_path, + metadata_json=metadata, + word_count=len(text.split()), + ) + session.add(document) + await session.commit() + await session.refresh(document, ["collections"]) + return document + except Exception as error: + if cover_target.exists(): + cover_target.unlink(missing_ok=True) + if isinstance(error, AwaazError): + raise _bad_request(error) from error + raise error + finally: + if stored.exists(): + stored.unlink(missing_ok=True) + extracted = stored.with_suffix(".extracted.txt") + if extracted.exists(): + extracted.unlink(missing_ok=True) def _join_authors(authors: list[str] | None) -> str | None: @@ -249,9 +264,13 @@ async def download_cover( @router.delete("/documents/{document_id}", status_code=status.HTTP_204_NO_CONTENT) -async def remove_document(document_id: str, session: AsyncSession = Depends(get_session)) -> None: +async def remove_document( + document_id: str, + session: AsyncSession = Depends(get_session), + settings: Settings = Depends(get_settings), +) -> None: try: - await delete_document(session, document_id) + await delete_document(session, document_id, settings) except AwaazError as error: raise _bad_request(error) from error diff --git a/src/awaaz/services/documents.py b/src/awaaz/services/documents.py index c55ec20..02fcd92 100644 --- a/src/awaaz/services/documents.py +++ b/src/awaaz/services/documents.py @@ -1,8 +1,12 @@ +import shutil +from pathlib import Path + from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload +from awaaz.config import Settings from awaaz.domain.exceptions import DocumentError from awaaz.models import Collection, Document, Job, document_collections, utc_now @@ -113,7 +117,7 @@ async def delete_collection(session: AsyncSession, collection_id: str) -> None: await session.commit() -async def delete_document(session: AsyncSession, document_id: str) -> None: +async def delete_document(session: AsyncSession, document_id: str, settings: Settings) -> None: document = await get_document(session, document_id) active = await session.scalar( select(Job.id).where( @@ -123,5 +127,24 @@ async def delete_document(session: AsyncSession, document_id: str) -> None: ) if active is not None: raise DocumentError("cancel active jobs before deleting") + + # Get associated job IDs to delete their audio files + jobs = await session.scalars(select(Job).where(Job.document_id == document_id)) + job_ids = [job.id for job in jobs] + + cover_path = document.cover_path + await session.delete(document) await session.commit() + + # Delete cover image if it exists + if cover_path: + cover_file = Path(cover_path) + if cover_file.is_file(): + cover_file.unlink(missing_ok=True) + + # Delete all generated job audio folders + for job_id in job_ids: + job_audio_dir = settings.audio_dir / job_id + if job_audio_dir.is_dir(): + shutil.rmtree(job_audio_dir, ignore_errors=True) diff --git a/tests/test_api_cleanup.py b/tests/test_api_cleanup.py new file mode 100644 index 0000000..dc1d3b3 --- /dev/null +++ b/tests/test_api_cleanup.py @@ -0,0 +1,88 @@ +from pathlib import Path +from unittest.mock import patch + +from fastapi.testclient import TestClient + +from awaaz.config import Settings + + +def test_document_upload_and_delete_cleanup(client: TestClient) -> None: + settings = Settings() + + # Pre-create dirs + settings.uploads_dir.mkdir(parents=True, exist_ok=True) + settings.audio_dir.mkdir(parents=True, exist_ok=True) + + async def mock_extract_text(source: Path) -> str: + return "This is my book text." + + async def mock_extract_metadata(source: Path) -> dict: + return {"title": "Test Title", "authors": ["John Doe"]} + + async def mock_extract_cover(source: Path, target: Path) -> bool: + target.write_bytes(b"dummy image data") + return True + + # 1. Upload Document (EPUB) + with ( + patch("awaaz.api.extract_text", new=mock_extract_text), + patch("awaaz.api.extract_metadata", new=mock_extract_metadata), + patch("awaaz.api.extract_cover", new=mock_extract_cover), + ): + response = client.post( + "/api/v1/documents/upload", + files={"file": ("test_book.epub", b"epub content bytes")}, + ) + assert response.status_code == 201 + doc_data = response.json() + doc_id = doc_data["id"] + cover_path = doc_data["cover_path"] + assert cover_path is not None + + # Verify temporary files are deleted, cover image exists + assert Path(cover_path).is_file() + remaining_files = list(settings.uploads_dir.iterdir()) + assert len(remaining_files) == 1 + assert remaining_files[0] == Path(cover_path) + + # 2. Create a Job for the document + response = client.post( + f"/api/v1/documents/{doc_id}/jobs", + json={ + "backend": "kokoro", + "model": "kokoro", + "voice": "af_bella", + "speed": 1.0, + "chunking_mode": "paragraph", + "character_limit": 1000, + }, + ) + assert response.status_code == 201 + job_data = response.json() + job_id = job_data["id"] + + # 3. Simulate generated audio files on disk + job_audio_dir = settings.audio_dir / job_id + job_audio_dir.mkdir(parents=True, exist_ok=True) + (job_audio_dir / "audiobook.mp3").write_bytes(b"mp3 data") + chunks_dir = job_audio_dir / "chunks" + chunks_dir.mkdir(parents=True, exist_ok=True) + (chunks_dir / "00000000.wav").write_bytes(b"wav chunk data") + + assert job_audio_dir.is_dir() + assert (job_audio_dir / "audiobook.mp3").is_file() + + # 4. Cancel the Job (so the document can be deleted) + response = client.post(f"/api/v1/jobs/{job_id}/cancel") + assert response.status_code == 200 + assert response.json()["status"] == "cancelled" + + # 5. Delete the Document + response = client.delete(f"/api/v1/documents/{doc_id}") + assert response.status_code == 204 + + # 5. Verify cover image is deleted + assert not Path(cover_path).exists() + + # 6. Verify job audio directory is deleted + assert not job_audio_dir.exists()