fix: automatic disk cleanup on upload and document deletion

.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:

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
 

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)

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()