diff --git a/README.md b/README.md index 06bb77c..076f306 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,11 @@ This is the worker backend API for the WebBuddhist application. +It handles background work such as TTS/audio generation and notification +dispatch. Plan/content data owned by the main backend is accessed **only via +HTTP APIs** — the worker does not share or query the backend Postgres schema +for subtasks, plan days, or related domain tables. + ## Installation Follow these steps to set up the project on your local machine: @@ -19,30 +24,37 @@ Follow these steps to set up the project on your local machine: poetry install ``` -## Database Setup +## Configuration -This worker backend shares the same databases as `app-pecha-backend`. Ensure the databases are running: +Set at least: -1. Navigate to the app-pecha-backend local setup directory: - ```sh - cd ../app-pecha-backend/local_setup - ``` -2. Start the databases using Docker (if not already running): - ```sh - docker-compose up -d - ``` -3. Return to webuddhist-worker directory: - ```sh - cd ../../webuddhist-worker - ``` -4. Apply database migrations: - ```sh - poetry run alembic upgrade head - ``` +| Variable | Purpose | +|----------|---------| +| `BACKEND_API_URL` | Base URL of the main backend API (default `http://127.0.0.1:8000/api/v1`) | +| `NOTIFICATION_DISPATCH_SECRET_TOKEN` | Shared secret sent as `X-Dispatch-Token` on internal backend calls | +| `DATABASE_URL` | Worker-owned Postgres (reminders / notification tables only) | +| `AUDIO_SQS_QUEUE_URL` | SQS queue for audio jobs produced by the backend | + +The worker talks to the backend for audio job status, generation payloads +(day/subtask content), and persisting generation results. Do **not** point +`DATABASE_URL` at the backend database for plan/subtask data. + +## Database Setup (worker-owned DB only) + +Start local infra from the backend repo if needed (Postgres/Mongo/Redis/etc.), +then apply **worker** migrations against the worker database: + +```sh +poetry run alembic upgrade head +``` + +These migrations cover worker tables (e.g. upcoming reminders), not backend +plan tables. ## Running the Application -1. Start the FastAPI development server: +1. Ensure the main backend is running (default `http://127.0.0.1:8000`). +2. Start the worker: ```sh poetry run uvicorn worker_api.app:api --port 8001 --reload ``` @@ -65,17 +77,17 @@ To check the coverage: poetry run pytest --cov=worker_api ``` ```sh -poetry run coverage html +poetry run coverage html ``` Open the coverage report: ```sh -open htmlcov/index.html +open htmlcov/index.html ``` ## Alembic Commands -Alembic is used for handling database migrations. Here are some common commands: +Alembic is used for handling **worker-owned** database migrations: 1. Create a new migration: ```sh @@ -104,23 +116,20 @@ Alembic is used for handling database migrations. Here are some common commands: ## Shared Infrastructure -This worker backend shares the following infrastructure with `app-pecha-backend`: -- PostgreSQL database (port 5434) -- MongoDB database (port 27017) -- Redis/Dragonfly cache (port 6379) -- Elasticsearch (port 9200) +This worker can reuse the same local Docker services as `WeBuddhist-Backend` +(Postgres, MongoDB, Redis, Elasticsearch), but: -Both backends can run simultaneously: -- `app-pecha-backend`: http://127.0.0.1:8000 -- `webuddhist-worker`: http://127.0.0.1:8001 +- Backend plan/content data is fetched and updated through backend internal APIs +- Worker Postgres should only hold worker-owned tables + +Typical local ports: -## Transferring Endpoints +- PostgreSQL: 5434 +- MongoDB: 27017 +- Redis/Dragonfly: 6379 +- Elasticsearch: 9200 -When transferring endpoints from `app-pecha-backend`: -1. Copy the relevant models, services, repositories, and views -2. Update imports to use `worker_api` instead of `pecha_api` -3. Add model imports to `migrations/env.py` for Alembic -4. Add document models to `worker_api/db/mongo_database.py` for Beanie -5. Include routers in `worker_api/app.py` -6. Run migrations if needed -7. Update tests accordingly +Both apps can run simultaneously: + +- `WeBuddhist-Backend`: http://127.0.0.1:8000 +- `webuddhist-worker`: http://127.0.0.1:8001 diff --git a/tests/audio/test_audio_generate_service.py b/tests/audio/test_audio_generate_service.py index 73df96f..a3e2f5c 100644 --- a/tests/audio/test_audio_generate_service.py +++ b/tests/audio/test_audio_generate_service.py @@ -4,15 +4,13 @@ import pytest from uuid import uuid4 from unittest.mock import patch, MagicMock, AsyncMock -from io import BytesIO from worker_api.audio.enums import ContentType, PlanAudioType, MonlamVoiceName from worker_api.audio.services.audio_generate_service import ( generate_plan_audio_service, _generate_audio_segments, - _update_subtask_timestamps, + _build_subtask_timestamps, _build_combined_wav, - _upload_and_persist_audio, ) @@ -20,200 +18,155 @@ class TestGeneratePlanAudioService: """Tests for generate_plan_audio_service function.""" @pytest.mark.asyncio - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - @patch("worker_api.audio.services.audio_generate_service.get_plan_day_by_id_any_plan") + @patch("worker_api.audio.services.audio_generate_service.get_day_audio_generation_payload", new_callable=AsyncMock) @patch("worker_api.audio.services.audio_generate_service._generate_audio_segments") - @patch("worker_api.audio.services.audio_generate_service._update_subtask_timestamps") + @patch("worker_api.audio.services.audio_generate_service._build_subtask_timestamps") @patch("worker_api.audio.services.audio_generate_service._build_combined_wav") - @patch("worker_api.audio.services.audio_generate_service._upload_and_persist_audio") + @patch("worker_api.audio.services.audio_generate_service._upload_day_audio") + @patch("worker_api.audio.services.audio_generate_service.apply_day_audio_generation_result", new_callable=AsyncMock) @patch("worker_api.audio.services.audio_generate_service.generate_presigned_access_url") async def test_generate_day_audio_success( self, mock_presigned_url, + mock_apply_result, mock_upload, mock_build_wav, - mock_update_timestamps, + mock_build_timestamps, mock_generate_segments, - mock_get_day, - mock_session_local, + mock_get_day_payload, ): """Test successful audio generation for a plan day.""" day_id = uuid4() plan_id = uuid4() - - # Mock database session - mock_db = MagicMock() - mock_session_local.return_value.__enter__.return_value = mock_db - - # Mock plan item - mock_plan_item = MagicMock() - mock_plan_item.id = day_id - mock_plan_item.plan_id = plan_id - mock_plan_item.tasks = [] - mock_get_day.return_value = mock_plan_item - - # Mock audio segments - mock_generate_segments.return_value = ([b"audio_data"], [MagicMock()]) - mock_update_timestamps.return_value = 45000 + sub_task_id = uuid4() + + mock_get_day_payload.return_value = { + "id": str(day_id), + "plan_id": str(plan_id), + "subtasks": [{"id": str(sub_task_id), "content_type": "TEXT", "content": "Hello"}], + } + mock_generate_segments.return_value = ([b"audio_data"], [{"id": str(sub_task_id)}]) + mock_build_timestamps.return_value = ( + 45000, + [{"sub_task_id": str(sub_task_id), "start_ms": 0, "end_ms": 45000}], + ) mock_build_wav.return_value = (b"wav_data", 1000) - - # Mock audio row - mock_audio_row = MagicMock() - mock_audio_row.audio_key = "audio/test.wav" - mock_audio_row.duration_ms = 45000 - mock_upload.return_value = mock_audio_row - + mock_upload.return_value = "audio/test.wav" mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - - # Execute + result = await generate_plan_audio_service( language="en", day_id=day_id, audio_type=PlanAudioType.TEXT_READING, ) - - # Assert + assert result["audio_url"] == "https://s3.example.com/audio.wav" assert result["audio_duration_ms"] == 45000 assert result["s3_key"] == "audio/test.wav" - - mock_get_day.assert_called_once_with(db=mock_db, day_id=day_id) - mock_generate_segments.assert_called_once() + mock_get_day_payload.assert_awaited_once_with(day_id=day_id) + mock_apply_result.assert_awaited_once() mock_upload.assert_called_once() @pytest.mark.asyncio - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - @patch("worker_api.audio.services.audio_generate_service.get_plan_day_by_id_any_plan") + @patch("worker_api.audio.services.audio_generate_service.get_day_audio_generation_payload", new_callable=AsyncMock) @patch("worker_api.audio.services.audio_generate_service._generate_audio_segments") async def test_generate_day_audio_no_segments( self, mock_generate_segments, - mock_get_day, - mock_session_local, + mock_get_day_payload, ): """Test audio generation returns empty when no segments are generated.""" day_id = uuid4() - - mock_db = MagicMock() - mock_session_local.return_value.__enter__.return_value = mock_db - - mock_plan_item = MagicMock() - mock_plan_item.tasks = [] - mock_get_day.return_value = mock_plan_item - - # No audio segments + mock_get_day_payload.return_value = { + "id": str(day_id), + "plan_id": str(uuid4()), + "subtasks": [], + } mock_generate_segments.return_value = ([], []) - + result = await generate_plan_audio_service( language="en", day_id=day_id, ) - + assert result == [] @pytest.mark.asyncio - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - @patch("worker_api.audio.services.audio_generate_service.get_sub_task_by_subtask_id") + @patch("worker_api.audio.services.audio_generate_service.get_sub_task_audio_generation_payload", new_callable=AsyncMock) @patch("worker_api.audio.services.audio_generate_service.generate_tts_audio") @patch("worker_api.audio.services.audio_generate_service.upload_bytes") - @patch("worker_api.audio.services.audio_generate_service.upsert_sub_task_timestamp") + @patch("worker_api.audio.services.audio_generate_service.apply_sub_task_audio_generation_result", new_callable=AsyncMock) @patch("worker_api.audio.services.audio_generate_service.generate_presigned_access_url") async def test_generate_subtask_audio_success( self, mock_presigned_url, - mock_upsert_timestamp, + mock_apply_result, mock_upload, mock_tts, mock_get_subtask, - mock_session_local, ): """Test successful audio generation for a single subtask.""" sub_task_id = uuid4() task_id = uuid4() - - mock_db = MagicMock() - mock_session_local.return_value.__enter__.return_value = mock_db - - # Mock subtask - mock_subtask = MagicMock() - mock_subtask.id = sub_task_id - mock_subtask.task_id = task_id - mock_subtask.content = "Test content" - mock_subtask.content_type = ContentType.TEXT - mock_get_subtask.return_value = mock_subtask - - # Mock TTS audio (44 byte header + audio data) + + mock_get_subtask.return_value = { + "id": str(sub_task_id), + "task_id": str(task_id), + "content": "Test content", + "content_type": "TEXT", + } + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 1000 mock_tts.return_value = wav_header + audio_data - mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - - # Execute + result = await generate_plan_audio_service( language="en", sub_task_id=sub_task_id, audio_type=PlanAudioType.TEXT_READING, ) - - # Assert + assert "audio_url" in result assert "audio_duration_ms" in result assert "s3_key" in result - - mock_get_subtask.assert_called_once_with(db=mock_db, id=sub_task_id) + mock_get_subtask.assert_awaited_once_with(sub_task_id=sub_task_id) mock_tts.assert_called_once() mock_upload.assert_called_once() - mock_upsert_timestamp.assert_called_once() + mock_apply_result.assert_awaited_once() @pytest.mark.asyncio - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - @patch("worker_api.audio.services.audio_generate_service.get_sub_task_by_subtask_id") - async def test_generate_subtask_audio_not_found( - self, - mock_get_subtask, - mock_session_local, - ): - """Test error when subtask is not found.""" + @patch("worker_api.audio.services.audio_generate_service.get_sub_task_audio_generation_payload", new_callable=AsyncMock) + async def test_generate_subtask_audio_not_found(self, mock_get_subtask): + """Test error when subtask is not found on backend.""" + from fastapi import HTTPException + sub_task_id = uuid4() - - mock_db = MagicMock() - mock_session_local.return_value.__enter__.return_value = mock_db - - mock_get_subtask.return_value = None - - with pytest.raises(Exception) as exc_info: + mock_get_subtask.side_effect = HTTPException(status_code=404, detail="not found") + + with pytest.raises(HTTPException) as exc_info: await generate_plan_audio_service( language="en", sub_task_id=sub_task_id, ) - + assert exc_info.value.status_code == 404 @pytest.mark.asyncio - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - @patch("worker_api.audio.services.audio_generate_service.get_sub_task_by_subtask_id") - async def test_generate_subtask_audio_invalid_content_type( - self, - mock_get_subtask, - mock_session_local, - ): - """Test error when subtask has invalid content type.""" + @patch("worker_api.audio.services.audio_generate_service.get_sub_task_audio_generation_payload", new_callable=AsyncMock) + async def test_generate_subtask_audio_invalid_content_type(self, mock_get_subtask): + """Test error when backend rejects invalid content type.""" + from fastapi import HTTPException + sub_task_id = uuid4() - - mock_db = MagicMock() - mock_session_local.return_value.__enter__.return_value = mock_db - - mock_subtask = MagicMock() - mock_subtask.content_type = ContentType.VIDEO # Invalid for audio - mock_get_subtask.return_value = mock_subtask - - with pytest.raises(Exception) as exc_info: + mock_get_subtask.side_effect = HTTPException(status_code=400, detail="invalid type") + + with pytest.raises(HTTPException) as exc_info: await generate_plan_audio_service( language="en", sub_task_id=sub_task_id, ) - + assert exc_info.value.status_code == 400 @pytest.mark.asyncio @@ -228,17 +181,17 @@ async def test_generate_audio_with_text_routes_correctly( "audio_duration_ms": 5000, "s3_key": "audio/generated/test.wav" } - + result = await generate_plan_audio_service( text="Test text input", language="en", audio_type=PlanAudioType.TEXT_READING, voice_name=MonlamVoiceName.DOLKAR_LHASA_FEMALE, ) - + assert result["audio_url"] == "https://s3.example.com/audio.wav" assert result["audio_duration_ms"] == 5000 - + mock_generate_from_text.assert_called_once_with( text="Test text input", language="en", @@ -259,40 +212,40 @@ async def test_generate_audio_with_text_and_prefix( "audio_duration_ms": 3000, "s3_key": "custom/prefix/test.wav" } - + result = await generate_plan_audio_service( text="Custom prefix text", language="bo", s3_key_prefix="custom/prefix", ) - + assert result["s3_key"] == "custom/prefix/test.wav" - + mock_generate_from_text.assert_called_once() call_kwargs = mock_generate_from_text.call_args.kwargs assert call_kwargs["s3_key_prefix"] == "custom/prefix" @pytest.mark.asyncio @patch("worker_api.audio.services.audio_generate_service._generate_audio_from_text") - @patch("worker_api.audio.services.audio_generate_service.SessionLocal") - async def test_generate_audio_text_bypasses_database( + @patch("worker_api.audio.services.audio_generate_service.get_day_audio_generation_payload", new_callable=AsyncMock) + async def test_generate_audio_text_bypasses_backend_fetch( self, - mock_session_local, + mock_get_day_payload, mock_generate_from_text, ): - """Test that text input bypasses database operations.""" + """Test that text input bypasses backend content fetch.""" mock_generate_from_text.return_value = { "audio_url": "https://s3.example.com/audio.wav", "audio_duration_ms": 4000, "s3_key": "audio/generated/test.wav" } - + await generate_plan_audio_service( - text="No database needed", + text="No backend content needed", language="en", ) - - mock_session_local.assert_not_called() + + mock_get_day_payload.assert_not_called() mock_generate_from_text.assert_called_once() @@ -304,24 +257,23 @@ def test_generate_segments_with_text_content(self, mock_tts): """Test generating audio segments from text content.""" wav_data = b"RIFF" + b"\x00" * 40 + b"audio_data" mock_tts.return_value = wav_data - - mock_subtask = MagicMock() - mock_subtask.content = "Test content" - mock_subtask.content_type = ContentType.TEXT - mock_subtask.audio_url = None - - mock_task = MagicMock() - mock_task.sub_tasks = [mock_subtask] - + + subtask = { + "id": str(uuid4()), + "content": "Test content", + "content_type": ContentType.TEXT, + "audio_url": None, + } + segments, refs = _generate_audio_segments( - [mock_task], + [subtask], PlanAudioType.TEXT_READING, "en", ) - + assert len(segments) == 1 assert len(refs) == 1 - assert refs[0] == mock_subtask + assert refs[0] == subtask mock_tts.assert_called_once() @patch("worker_api.audio.services.audio_generate_service.download_bytes") @@ -329,37 +281,35 @@ def test_generate_segments_with_existing_audio(self, mock_download): """Test reusing existing audio from subtask.""" wav_data = b"RIFF" + b"\x00" * 40 + b"existing_audio" mock_download.return_value = wav_data - - mock_subtask = MagicMock() - mock_subtask.content_type = ContentType.TEXT - mock_subtask.audio_url = "audio/existing.wav" - - mock_task = MagicMock() - mock_task.sub_tasks = [mock_subtask] - + + subtask = { + "id": str(uuid4()), + "content_type": "TEXT", + "audio_url": "audio/existing.wav", + } + segments, refs = _generate_audio_segments( - [mock_task], + [subtask], PlanAudioType.TEXT_READING, "en", ) - + assert len(segments) == 1 mock_download.assert_called_once_with(key="audio/existing.wav") def test_generate_segments_skips_invalid_content_types(self): """Test that non-text/source_reference subtasks are skipped.""" - mock_subtask = MagicMock() - mock_subtask.content_type = ContentType.VIDEO - - mock_task = MagicMock() - mock_task.sub_tasks = [mock_subtask] - + subtask = { + "id": str(uuid4()), + "content_type": ContentType.VIDEO, + } + segments, refs = _generate_audio_segments( - [mock_task], + [subtask], PlanAudioType.TEXT_READING, "en", ) - + assert len(segments) == 0 assert len(refs) == 0 @@ -370,10 +320,10 @@ class TestBuildCombinedWav: def test_build_wav_single_segment(self): """Test building WAV file from single audio segment.""" audio_data = b"\x00" * 1000 - + wav, size = _build_combined_wav([audio_data]) - - assert len(wav) > len(audio_data) # Header + data + + assert len(wav) > len(audio_data) assert wav[:4] == b"RIFF" assert b"WAVE" in wav assert size == len(audio_data) @@ -382,49 +332,55 @@ def test_build_wav_multiple_segments(self): """Test building WAV file from multiple audio segments.""" segment1 = b"\x00" * 500 segment2 = b"\x01" * 500 - + wav, size = _build_combined_wav([segment1, segment2]) - + assert wav[:4] == b"RIFF" assert size == len(segment1) + len(segment2) def test_build_wav_empty_segments(self): """Test building WAV file with no segments.""" wav, size = _build_combined_wav([]) - + assert wav[:4] == b"RIFF" assert size == 0 -class TestUpdateSubtaskTimestamps: - """Tests for _update_subtask_timestamps helper function.""" +class TestBuildSubtaskTimestamps: + """Tests for _build_subtask_timestamps helper function.""" - @patch("worker_api.audio.services.audio_generate_service.upsert_sub_task_timestamp") - def test_update_timestamps(self, mock_upsert): - """Test updating subtask timestamps.""" - mock_db = MagicMock() - - # 1000 bytes at 2 bytes per sample = 500 samples - # 500 samples / 24000 Hz = 0.020833 seconds = 20.833 ms + def test_build_timestamps(self): + """Test building subtask timestamps without DB writes.""" audio_segment = b"\x00" * 1000 - - mock_subtask = MagicMock() - mock_subtask.id = uuid4() - - duration = _update_subtask_timestamps( - mock_db, + sub_task_id = uuid4() + + duration, timestamps = _build_subtask_timestamps( [audio_segment], - [mock_subtask], - 24000, # sample_rate - 2, # bytes_per_sample + [{"id": str(sub_task_id)}], + 24000, + 2, ) - + assert duration > 0 - mock_upsert.assert_called_once() - call_kwargs = mock_upsert.call_args.kwargs - assert call_kwargs["sub_task_id"] == mock_subtask.id - assert call_kwargs["start_ms"] == 0 - assert call_kwargs["end_ms"] > 0 + assert len(timestamps) == 1 + assert timestamps[0]["sub_task_id"] == str(sub_task_id) + assert timestamps[0]["start_ms"] == 0 + assert timestamps[0]["end_ms"] > 0 + + +class TestUploadDayAudio: + @patch("worker_api.audio.services.audio_generate_service.upload_bytes") + def test_upload_day_audio_builds_key(self, mock_upload): + from worker_api.audio.services.audio_generate_service import _upload_day_audio + + plan_id = uuid4() + plan_item_id = uuid4() + + key = _upload_day_audio(b"wav", plan_id, plan_item_id) + + assert key.startswith(f"audio/plan_days/{plan_id}/{plan_item_id}/") + assert key.endswith(".wav") + mock_upload.assert_called_once() class TestGenerateAudioFromText: @@ -442,24 +398,24 @@ async def test_generate_audio_from_text_success( ): """Test successful audio generation from text.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 1000 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - + result = await _generate_audio_from_text( text="Test text for audio generation", language="en", audio_type=PlanAudioType.TEXT_READING, ) - + assert result["audio_url"] == "https://s3.example.com/audio.wav" assert result["audio_duration_ms"] > 0 assert "s3_key" in result assert result["s3_key"].startswith("audio/generated/") assert result["s3_key"].endswith(".wav") - + mock_tts.assert_called_once_with( "Test text for audio generation", PlanAudioType.TEXT_READING, @@ -480,21 +436,21 @@ async def test_generate_audio_from_text_with_custom_prefix( ): """Test audio generation with custom S3 key prefix.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 500 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/custom/audio.wav" - + result = await _generate_audio_from_text( text="Custom prefix test", language="en", s3_key_prefix="custom/prefix", ) - + assert result["s3_key"].startswith("custom/prefix/") assert result["s3_key"].endswith(".wav") - + upload_call = mock_upload.call_args uploaded_key = upload_call.kwargs["key"] assert uploaded_key.startswith("custom/prefix/") @@ -511,18 +467,18 @@ async def test_generate_audio_from_text_with_voice_name( ): """Test audio generation with specific voice name.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 800 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - + result = await _generate_audio_from_text( text="Voice test", language="bo", voice_name=MonlamVoiceName.YANGCHEN_LHASA_FEMALE, ) - + assert "audio_url" in result mock_tts.assert_called_once_with( "Voice test", @@ -543,18 +499,18 @@ async def test_generate_audio_from_text_recitation_type( ): """Test audio generation with RECITATION type.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 1200 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - + result = await _generate_audio_from_text( text="Recitation text", language="bo", audio_type=PlanAudioType.RECITATION, ) - + assert result["audio_duration_ms"] > 0 mock_tts.assert_called_once_with( "Recitation text", @@ -575,18 +531,18 @@ async def test_generate_audio_from_text_instruction_type( ): """Test audio generation with INSTRUCTION type.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 600 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - + result = await _generate_audio_from_text( text="Instruction text", language="en", audio_type=PlanAudioType.INSTRUCTION, ) - + assert "s3_key" in result mock_tts.assert_called_once_with( "Instruction text", @@ -607,16 +563,15 @@ async def test_generate_audio_from_text_duration_calculation( ): """Test correct duration calculation from audio data.""" from worker_api.audio.services.audio_generate_service import _generate_audio_from_text - + wav_header = b"RIFF" + b"\x00" * 40 audio_data = b"\x00" * 48000 mock_tts.return_value = wav_header + audio_data mock_presigned_url.return_value = "https://s3.example.com/audio.wav" - + result = await _generate_audio_from_text( text="Duration test", language="en", ) - - expected_duration_ms = int((48000 / 2 / 24000) * 1000) - assert result["audio_duration_ms"] == expected_duration_ms + + assert result["audio_duration_ms"] == 1000 diff --git a/tests/audio/test_audio_job_consumer.py b/tests/audio/test_audio_job_consumer.py new file mode 100644 index 0000000..e4cad94 --- /dev/null +++ b/tests/audio/test_audio_job_consumer.py @@ -0,0 +1,300 @@ +"""Tests for SQS audio job consumer.""" +import asyncio +import json +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from worker_api.audio.enums import AudioJobStatus, MonlamVoiceName, PlanAudioType +from worker_api.audio.services.audio_job_consumer import ( + _error_detail, + _normalize_result, + _parse_audio_type, + _parse_uuid, + _parse_voice_name, + process_audio_job_message, + run_audio_sqs_consumer, +) + + +class TestParseHelpers: + def test_parse_uuid_none_and_empty(self): + assert _parse_uuid(None) is None + assert _parse_uuid("") is None + + def test_parse_uuid_valid(self): + value = uuid4() + assert _parse_uuid(str(value)) == value + + def test_parse_audio_type_defaults_and_enum(self): + assert _parse_audio_type(None) == PlanAudioType.TEXT_READING + assert _parse_audio_type("") == PlanAudioType.TEXT_READING + assert _parse_audio_type(PlanAudioType.RECITATION) == PlanAudioType.RECITATION + assert _parse_audio_type("INSTRUCTION") == PlanAudioType.INSTRUCTION + + def test_parse_voice_name_defaults_and_enum(self): + assert _parse_voice_name(None) == MonlamVoiceName.DOLKAR_LHASA_FEMALE + assert _parse_voice_name("") == MonlamVoiceName.DOLKAR_LHASA_FEMALE + assert _parse_voice_name(MonlamVoiceName.YANGCHEN_LHASA_FEMALE) == MonlamVoiceName.YANGCHEN_LHASA_FEMALE + assert _parse_voice_name("yangchen_lhasa_female") == MonlamVoiceName.YANGCHEN_LHASA_FEMALE + + def test_normalize_result_success(self): + assert _normalize_result({"s3_key": "a.wav", "audio_url": "u", "audio_duration_ms": 1}) == { + "audio_url": "u", + "audio_duration_ms": 1, + "s3_key": "a.wav", + } + + def test_normalize_result_rejects_non_dict(self): + with pytest.raises(ValueError, match="no result"): + _normalize_result([]) + + def test_normalize_result_rejects_empty(self): + with pytest.raises(ValueError, match="empty result"): + _normalize_result({"audio_duration_ms": 1}) + + def test_error_detail_http_dict_and_string(self): + assert _error_detail(HTTPException(status_code=400, detail={"message": "bad"})) == "bad" + assert _error_detail(HTTPException(status_code=400, detail={"code": "x"})) == "{'code': 'x'}" + assert _error_detail(HTTPException(status_code=400, detail="plain")) == "plain" + assert _error_detail(ValueError("oops")) == "oops" + + +class TestProcessAudioJobMessage: + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.generate_plan_audio_service", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.update_audio_job_status", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_processes_day_job_successfully( + self, + mock_get_job, + mock_update_status, + mock_generate, + mock_delete, + ): + job_id = uuid4() + day_id = uuid4() + mock_get_job.return_value = {"job_id": str(job_id), "status": AudioJobStatus.PENDING.value} + mock_generate.return_value = { + "audio_url": "https://example.com/a.wav", + "audio_duration_ms": 1200, + "s3_key": "audio/plan_days/a.wav", + } + + message = { + "ReceiptHandle": "abc", + "Body": json.dumps( + { + "job_id": str(job_id), + "day_id": str(day_id), + "sub_task_id": None, + "language": "bo", + "type": "TEXT_READING", + "voice_name": "dolkar_lhasa_female", + } + ), + } + + await process_audio_job_message(message) + + assert mock_update_status.await_count == 2 + assert mock_update_status.await_args_list[0].kwargs["status"] == AudioJobStatus.PROCESSING + assert mock_update_status.await_args_list[1].kwargs["status"] == AudioJobStatus.COMPLETED + mock_generate.assert_awaited_once() + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.generate_plan_audio_service", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.update_audio_job_status", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_marks_failed_on_generation_error( + self, + mock_get_job, + mock_update_status, + mock_generate, + mock_delete, + ): + job_id = uuid4() + mock_get_job.return_value = {"job_id": str(job_id), "status": AudioJobStatus.PENDING.value} + mock_generate.side_effect = HTTPException(status_code=404, detail={"message": "Sub task not found"}) + + message = { + "ReceiptHandle": "abc", + "Body": json.dumps( + { + "job_id": str(job_id), + "day_id": None, + "sub_task_id": str(uuid4()), + "language": "en", + "type": "TEXT_READING", + "voice_name": "dolkar_lhasa_female", + } + ), + } + + await process_audio_job_message(message) + + assert mock_update_status.await_args_list[-1].kwargs["status"] == AudioJobStatus.FAILED + assert "Sub task not found" in mock_update_status.await_args_list[-1].kwargs["error_message"] + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.generate_plan_audio_service", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.update_audio_job_status", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_skips_completed_job( + self, + mock_get_job, + mock_update_status, + mock_generate, + mock_delete, + ): + job_id = uuid4() + mock_get_job.return_value = {"job_id": str(job_id), "status": AudioJobStatus.COMPLETED.value} + + message = { + "ReceiptHandle": "abc", + "Body": json.dumps({"job_id": str(job_id), "language": "en"}), + } + + await process_audio_job_message(message) + + mock_generate.assert_not_called() + mock_update_status.assert_not_awaited() + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.parse_audio_job_message_body", return_value=None) + async def test_deletes_invalid_body(self, _parse, mock_delete): + await process_audio_job_message({"ReceiptHandle": "abc", "Body": "bad"}) + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.parse_audio_job_message_body") + async def test_deletes_invalid_job_id(self, mock_parse, mock_delete): + mock_parse.return_value = {"job_id": ""} + await process_audio_job_message({"ReceiptHandle": "abc", "Body": "{}"}) + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_deletes_when_job_missing(self, mock_get_job, mock_delete): + job_id = uuid4() + mock_get_job.return_value = None + message = { + "ReceiptHandle": "abc", + "Body": json.dumps({"job_id": str(job_id), "language": "en"}), + } + await process_audio_job_message(message) + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.generate_plan_audio_service", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.update_audio_job_status", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_fails_when_language_missing( + self, + mock_get_job, + mock_update_status, + mock_generate, + mock_delete, + ): + job_id = uuid4() + mock_get_job.return_value = {"job_id": str(job_id), "status": AudioJobStatus.PENDING.value} + message = { + "ReceiptHandle": "abc", + "Body": json.dumps( + { + "job_id": str(job_id), + "day_id": str(uuid4()), + "language": "", + "type": "TEXT_READING", + } + ), + } + + await process_audio_job_message(message) + + assert mock_update_status.await_args_list[-1].kwargs["status"] == AudioJobStatus.FAILED + assert "language is required" in mock_update_status.await_args_list[-1].kwargs["error_message"] + mock_generate.assert_not_called() + mock_delete.assert_called_once_with("abc") + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.delete_audio_job_message") + @patch("worker_api.audio.services.audio_job_consumer.generate_plan_audio_service", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.update_audio_job_status", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.get_audio_job_status", new_callable=AsyncMock) + async def test_fails_when_target_ids_missing( + self, + mock_get_job, + mock_update_status, + mock_generate, + mock_delete, + ): + job_id = uuid4() + mock_get_job.return_value = {"job_id": str(job_id), "status": AudioJobStatus.PENDING.value} + message = { + "ReceiptHandle": "abc", + "Body": json.dumps({"job_id": str(job_id), "language": "en"}), + } + + await process_audio_job_message(message) + + assert mock_update_status.await_args_list[-1].kwargs["status"] == AudioJobStatus.FAILED + assert "Exactly one of day_id or sub_task_id" in mock_update_status.await_args_list[-1].kwargs["error_message"] + mock_generate.assert_not_called() + + +class TestRunAudioSqsConsumer: + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer._POLL_IDLE_SECONDS", 0.01) + @patch("worker_api.audio.services.audio_job_consumer.is_audio_sqs_poll_enabled", return_value=False) + async def test_idle_loop_until_stopped(self, _poll_enabled): + stop_event = asyncio.Event() + + async def stop_soon(): + await asyncio.sleep(0.02) + stop_event.set() + + asyncio.create_task(stop_soon()) + await run_audio_sqs_consumer(stop_event) + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer.process_audio_job_message", new_callable=AsyncMock) + @patch("worker_api.audio.services.audio_job_consumer.receive_audio_job_messages") + @patch("worker_api.audio.services.audio_job_consumer.is_audio_sqs_poll_enabled", return_value=True) + async def test_processes_received_messages(self, _poll_enabled, mock_receive, mock_process): + stop_event = asyncio.Event() + mock_receive.side_effect = [[], [{"ReceiptHandle": "1"}, {"ReceiptHandle": "2"}]] + + async def process_then_stop(_message): + stop_event.set() + + mock_process.side_effect = process_then_stop + await run_audio_sqs_consumer(stop_event) + mock_process.assert_awaited_once() + assert mock_receive.call_count >= 2 + + @pytest.mark.asyncio + @patch("worker_api.audio.services.audio_job_consumer._POLL_ERROR_SECONDS", 0.01) + @patch("worker_api.audio.services.audio_job_consumer.receive_audio_job_messages", side_effect=RuntimeError("boom")) + @patch("worker_api.audio.services.audio_job_consumer.is_audio_sqs_poll_enabled", return_value=True) + async def test_error_loop_until_stopped(self, _poll_enabled, _receive): + stop_event = asyncio.Event() + + async def stop_soon(): + await asyncio.sleep(0.02) + stop_event.set() + + asyncio.create_task(stop_soon()) + await run_audio_sqs_consumer(stop_event) diff --git a/tests/audio/test_backend_client.py b/tests/audio/test_backend_client.py new file mode 100644 index 0000000..843f840 --- /dev/null +++ b/tests/audio/test_backend_client.py @@ -0,0 +1,243 @@ +"""Tests for audio backend HTTP client.""" +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from worker_api.audio.enums import AudioJobStatus +from worker_api.audio.services.backend_client import ( + _backend_url, + _dispatch_headers, + _raise_for_backend_response, + apply_day_audio_generation_result, + apply_sub_task_audio_generation_result, + get_audio_job_status, + get_day_audio_generation_payload, + get_sub_task_audio_generation_payload, + update_audio_job_status, +) + + +class TestBackendUrl: + @patch("worker_api.audio.services.backend_client.get") + def test_returns_stripped_url(self, mock_get): + mock_get.return_value = "http://backend.example/api/v1/" + assert _backend_url() == "http://backend.example/api/v1" + + @patch("worker_api.audio.services.backend_client.get") + def test_raises_when_missing(self, mock_get): + mock_get.return_value = "" + with pytest.raises(RuntimeError, match="BACKEND_API_URL"): + _backend_url() + + +class TestDispatchHeaders: + @patch("worker_api.audio.services.backend_client.get") + def test_returns_headers(self, mock_get): + mock_get.return_value = "secret-token" + assert _dispatch_headers() == {"X-Dispatch-Token": "secret-token"} + + @patch("worker_api.audio.services.backend_client.get") + def test_raises_when_missing(self, mock_get): + mock_get.return_value = "" + with pytest.raises(RuntimeError, match="NOTIFICATION_DISPATCH_SECRET_TOKEN"): + _dispatch_headers() + + +class TestRaiseForBackendResponse: + def test_success_noop(self): + response = MagicMock() + response.is_success = True + _raise_for_backend_response(response) + + def test_raises_with_json_detail(self): + response = MagicMock() + response.is_success = False + response.status_code = 400 + response.json.return_value = {"message": "bad request"} + + with pytest.raises(HTTPException) as exc_info: + _raise_for_backend_response(response) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"message": "bad request"} + + def test_raises_with_text_when_json_invalid(self): + response = MagicMock() + response.is_success = False + response.status_code = 500 + response.json.side_effect = ValueError("not json") + response.text = "server error" + + with pytest.raises(HTTPException) as exc_info: + _raise_for_backend_response(response) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "server error" + + +def _mock_async_client(response: MagicMock): + client = AsyncMock() + client.get = AsyncMock(return_value=response) + client.patch = AsyncMock(return_value=response) + client.post = AsyncMock(return_value=response) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + return client + + +class TestGetAudioJobStatus: + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_returns_json_on_success(self, mock_client_cls, _url, _headers): + job_id = uuid4() + response = MagicMock() + response.status_code = 200 + response.is_success = True + response.json.return_value = {"job_id": str(job_id), "status": "PENDING"} + mock_client_cls.return_value = _mock_async_client(response) + + result = await get_audio_job_status(job_id) + + assert result["status"] == "PENDING" + + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_returns_none_on_404(self, mock_client_cls, _url, _headers): + response = MagicMock() + response.status_code = 404 + response.is_success = False + mock_client_cls.return_value = _mock_async_client(response) + + result = await get_audio_job_status(uuid4()) + + assert result is None + + +class TestUpdateAudioJobStatus: + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_sends_status_payload(self, mock_client_cls, _url, _headers): + job_id = uuid4() + response = MagicMock() + response.is_success = True + response.json.return_value = {"status": "COMPLETED"} + client = _mock_async_client(response) + mock_client_cls.return_value = client + + result = await update_audio_job_status( + job_id=job_id, + status=AudioJobStatus.COMPLETED, + result={"s3_key": "a.wav"}, + error_message=None, + ) + + assert result["status"] == "COMPLETED" + payload = client.patch.await_args.kwargs["json"] + assert payload == {"status": "completed", "result": {"s3_key": "a.wav"}} + + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_includes_error_message(self, mock_client_cls, _url, _headers): + response = MagicMock() + response.is_success = True + response.json.return_value = {"status": "FAILED"} + client = _mock_async_client(response) + mock_client_cls.return_value = client + + await update_audio_job_status( + job_id=uuid4(), + status=AudioJobStatus.FAILED, + error_message="boom", + ) + + payload = client.patch.await_args.kwargs["json"] + assert payload == {"status": "failed", "error_message": "boom"} + + +class TestGenerationPayloads: + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_get_day_payload(self, mock_client_cls, _url, _headers): + day_id = uuid4() + response = MagicMock() + response.is_success = True + response.json.return_value = {"id": str(day_id), "subtasks": []} + mock_client_cls.return_value = _mock_async_client(response) + + result = await get_day_audio_generation_payload(day_id) + + assert result["id"] == str(day_id) + + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_get_subtask_payload(self, mock_client_cls, _url, _headers): + sub_task_id = uuid4() + response = MagicMock() + response.is_success = True + response.json.return_value = {"id": str(sub_task_id), "content": "hi"} + mock_client_cls.return_value = _mock_async_client(response) + + result = await get_sub_task_audio_generation_payload(sub_task_id) + + assert result["content"] == "hi" + + +class TestApplyGenerationResults: + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_apply_day_result(self, mock_client_cls, _url, _headers): + day_id = uuid4() + response = MagicMock() + response.is_success = True + client = _mock_async_client(response) + mock_client_cls.return_value = client + + await apply_day_audio_generation_result( + day_id=day_id, + audio_key="audio/day.wav", + duration_ms=1000, + file_size_bytes=2000, + timestamps=[{"sub_task_id": "1", "start_ms": 0, "end_ms": 1000}], + ) + + client.post.assert_awaited_once() + payload = client.post.await_args.kwargs["json"] + assert payload["audio_key"] == "audio/day.wav" + assert payload["duration_ms"] == 1000 + assert payload["file_size_bytes"] == 2000 + + @pytest.mark.asyncio + @patch("worker_api.audio.services.backend_client._dispatch_headers", return_value={"X-Dispatch-Token": "t"}) + @patch("worker_api.audio.services.backend_client._backend_url", return_value="http://backend") + @patch("worker_api.audio.services.backend_client.httpx.AsyncClient") + async def test_apply_subtask_result(self, mock_client_cls, _url, _headers): + sub_task_id = uuid4() + response = MagicMock() + response.is_success = True + client = _mock_async_client(response) + mock_client_cls.return_value = client + + await apply_sub_task_audio_generation_result( + sub_task_id=sub_task_id, + audio_key="audio/sub.wav", + duration_ms=500, + ) + + payload = client.post.await_args.kwargs["json"] + assert payload == {"audio_key": "audio/sub.wav", "duration_ms": 500} diff --git a/tests/audio/test_sqs_client.py b/tests/audio/test_sqs_client.py new file mode 100644 index 0000000..5ad506c --- /dev/null +++ b/tests/audio/test_sqs_client.py @@ -0,0 +1,143 @@ +"""Tests for audio SQS client helpers.""" +import json +from unittest.mock import MagicMock, patch + +from botocore.exceptions import ClientError + +from worker_api.audio.sqs_client import ( + _get_sqs_client, + delete_audio_job_message, + get_audio_sqs_queue_url, + is_audio_sqs_configured, + is_audio_sqs_poll_enabled, + parse_audio_job_message_body, + receive_audio_job_messages, +) +import worker_api.audio.sqs_client as sqs_module + + +class TestQueueConfig: + @patch("worker_api.audio.sqs_client.get", return_value=" https://sqs.example/queue ") + def test_get_queue_url_strips(self, _get): + assert get_audio_sqs_queue_url() == "https://sqs.example/queue" + + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_is_configured_true(self, _url): + assert is_audio_sqs_configured() is True + + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="") + def test_is_configured_false(self, _url): + assert is_audio_sqs_configured() is False + + @patch("worker_api.audio.sqs_client.is_audio_sqs_configured", return_value=True) + @patch("worker_api.audio.sqs_client.get_bool", return_value=True) + def test_poll_enabled(self, _bool, _configured): + assert is_audio_sqs_poll_enabled() is True + + @patch("worker_api.audio.sqs_client.is_audio_sqs_configured", return_value=False) + @patch("worker_api.audio.sqs_client.get_bool", return_value=True) + def test_poll_disabled_when_unconfigured(self, _bool, _configured): + assert is_audio_sqs_poll_enabled() is False + + +class TestGetSqsClient: + def setup_method(self): + sqs_module._sqs_client = None + + def teardown_method(self): + sqs_module._sqs_client = None + + @patch("worker_api.audio.sqs_client.boto3.client") + @patch("worker_api.audio.sqs_client.get", side_effect=lambda key: f"value-{key}") + def test_creates_client_once(self, _get, mock_boto_client): + mock_boto_client.return_value = MagicMock(name="sqs") + + first = _get_sqs_client() + second = _get_sqs_client() + + assert first is second + mock_boto_client.assert_called_once_with( + "sqs", + aws_access_key_id="value-AWS_ACCESS_KEY", + aws_secret_access_key="value-AWS_SECRET_KEY", + region_name="value-AWS_REGION", + ) + + +class TestReceiveMessages: + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="") + def test_returns_empty_without_queue(self, _url): + assert receive_audio_job_messages() == [] + + @patch("worker_api.audio.sqs_client.get_int", side_effect=lambda key: 1) + @patch("worker_api.audio.sqs_client._get_sqs_client") + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_returns_messages(self, _url, mock_get_client, _get_int): + client = MagicMock() + client.receive_message.return_value = {"Messages": [{"MessageId": "1"}]} + mock_get_client.return_value = client + + assert receive_audio_job_messages() == [{"MessageId": "1"}] + + @patch("worker_api.audio.sqs_client.get_int", side_effect=lambda key: 1) + @patch("worker_api.audio.sqs_client._get_sqs_client") + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_returns_empty_on_client_error(self, _url, mock_get_client, _get_int): + client = MagicMock() + client.receive_message.side_effect = ClientError( + {"Error": {"Code": "AccessDenied", "Message": "denied"}}, + "ReceiveMessage", + ) + mock_get_client.return_value = client + + assert receive_audio_job_messages() == [] + + +class TestDeleteMessage: + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="") + def test_noop_without_queue(self, _url): + delete_audio_job_message("receipt") + + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_noop_without_receipt(self, _url): + delete_audio_job_message("") + + @patch("worker_api.audio.sqs_client._get_sqs_client") + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_deletes_message(self, _url, mock_get_client): + client = MagicMock() + mock_get_client.return_value = client + + delete_audio_job_message("receipt-1") + + client.delete_message.assert_called_once_with( + QueueUrl="https://sqs.example/queue", + ReceiptHandle="receipt-1", + ) + + @patch("worker_api.audio.sqs_client._get_sqs_client") + @patch("worker_api.audio.sqs_client.get_audio_sqs_queue_url", return_value="https://sqs.example/queue") + def test_swallows_client_error(self, _url, mock_get_client): + client = MagicMock() + client.delete_message.side_effect = ClientError( + {"Error": {"Code": "ReceiptHandleIsInvalid", "Message": "bad"}}, + "DeleteMessage", + ) + mock_get_client.return_value = client + + delete_audio_job_message("receipt-1") + + +class TestParseMessageBody: + def test_parses_valid_body(self): + body = parse_audio_job_message_body(json.dumps({"job_id": "abc", "language": "en"})) + assert body["job_id"] == "abc" + + def test_invalid_json(self): + assert parse_audio_job_message_body("not-json") is None + + def test_missing_job_id(self): + assert parse_audio_job_message_body(json.dumps({"language": "en"})) is None + + def test_non_dict_body(self): + assert parse_audio_job_message_body(json.dumps(["job"])) is None diff --git a/tests/audio/test_tibetan_chunker.py b/tests/audio/test_tibetan_chunker.py new file mode 100644 index 0000000..4753f89 --- /dev/null +++ b/tests/audio/test_tibetan_chunker.py @@ -0,0 +1,82 @@ +"""Tests for Tibetan text chunking.""" +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from worker_api.audio.services.tibetan_chunker import ( + _get_tokenizer, + chunk_tibetan_text, +) +import worker_api.audio.services.tibetan_chunker as chunker_module + + +def _token(text: str, syllables: int | None): + return SimpleNamespace( + text=text, + syls=["x"] * syllables if syllables is not None else None, + ) + + +class TestGetTokenizer: + def setup_method(self): + chunker_module._tokenizer = None + + def teardown_method(self): + chunker_module._tokenizer = None + + @patch("worker_api.audio.services.tibetan_chunker.WordTokenizer") + def test_creates_tokenizer_once(self, mock_word_tokenizer): + instance = MagicMock() + mock_word_tokenizer.return_value = instance + + first = _get_tokenizer() + second = _get_tokenizer() + + assert first is second + mock_word_tokenizer.assert_called_once() + + +class TestChunkTibetanText: + def test_empty_text_returns_empty(self): + assert chunk_tibetan_text("") == [] + assert chunk_tibetan_text(" ") == [] + + @patch("worker_api.audio.services.tibetan_chunker._get_tokenizer") + def test_single_chunk_under_limit(self, mock_get_tokenizer): + tokenizer = MagicMock() + tokenizer.tokenize.return_value = [ + _token("ཁྱེད", 2), + _token("རང", 1), + ] + mock_get_tokenizer.return_value = tokenizer + + chunks = chunk_tibetan_text("ཁྱེད་རང", max_syllables=15) + + assert chunks == ["ཁྱེདརང"] + tokenizer.tokenize.assert_called_once_with("ཁྱེད་རང", split_affixes=False) + + @patch("worker_api.audio.services.tibetan_chunker._get_tokenizer") + def test_splits_when_exceeding_max_syllables(self, mock_get_tokenizer): + tokenizer = MagicMock() + tokenizer.tokenize.return_value = [ + _token("aaa", 3), + _token("bbb", 3), + _token("ccc", 3), + ] + mock_get_tokenizer.return_value = tokenizer + + chunks = chunk_tibetan_text("aaa bbb ccc", max_syllables=6) + + assert chunks == ["aaabbb", "ccc"] + + @patch("worker_api.audio.services.tibetan_chunker._get_tokenizer") + def test_missing_syls_counts_as_one(self, mock_get_tokenizer): + tokenizer = MagicMock() + tokenizer.tokenize.return_value = [ + _token("x", None), + _token("y", None), + ] + mock_get_tokenizer.return_value = tokenizer + + chunks = chunk_tibetan_text("xy", max_syllables=1) + + assert chunks == ["x", "y"] diff --git a/worker_api/audio/enums.py b/worker_api/audio/enums.py index b0d435a..dde6876 100644 --- a/worker_api/audio/enums.py +++ b/worker_api/audio/enums.py @@ -33,5 +33,12 @@ class MonlamVoiceName(str, enum.Enum): WANGDONTSO_KHAM_FEMALE = "wangdontso_kham_female" +class AudioJobStatus(str, enum.Enum): + PENDING = "pending" + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + + ContentTypeEnum = Enum(ContentType) PlanAudioTypeEnum = Enum(PlanAudioType) diff --git a/worker_api/audio/services/audio_generate_service.py b/worker_api/audio/services/audio_generate_service.py index e2a3072..588c1ad 100644 --- a/worker_api/audio/services/audio_generate_service.py +++ b/worker_api/audio/services/audio_generate_service.py @@ -1,81 +1,80 @@ import struct from io import BytesIO -from typing import Optional, List +from typing import Any, Dict, List, Optional from uuid import UUID, uuid4 -from fastapi import HTTPException -from sqlalchemy.orm import Session -from starlette import status - from worker_api.audio.enums import ContentType, PlanAudioType, MonlamVoiceName -from worker_api.audio.models.plan_item_audio_models import PlanItemAudio -from worker_api.audio.models.plan_items_models import PlanItem -from worker_api.audio.models.plan_sub_tasks_models import PlanSubTask -from worker_api.audio.repositories.plan_items_repository import get_plan_day_by_id_any_plan -from worker_api.audio.repositories.plan_sub_tasks_repository import get_sub_task_by_subtask_id -from worker_api.audio.repositories.plan_item_audio_repository import upsert_plan_item_audio -from worker_api.audio.repositories.sub_task_timestamps_repository import upsert_sub_task_timestamp +from worker_api.audio.services.backend_client import ( + apply_day_audio_generation_result, + apply_sub_task_audio_generation_result, + get_day_audio_generation_payload, + get_sub_task_audio_generation_payload, +) from worker_api.audio.services.tts_service import generate_tts_audio -from worker_api.config import get -from worker_api.db.database import SessionLocal from worker_api.uploads.S3_utils import upload_bytes, download_bytes, generate_presigned_access_url WAV_CONTENT_TYPE = "audio/wav" +def _is_tts_content_type(content_type: Any) -> bool: + if isinstance(content_type, ContentType): + return content_type in {ContentType.TEXT, ContentType.SOURCE_REFERENCE} + return str(content_type) in {"TEXT", "SOURCE_REFERENCE"} + + def _generate_audio_segments( - tasks, + subtasks: List[Dict[str, Any]], audio_type: PlanAudioType, language: str, voice_name: Optional[str] = None, -) -> tuple[List[bytes], list]: +) -> tuple[List[bytes], List[Dict[str, Any]]]: wav_header_size = 44 audio_segments: List[bytes] = [] - subtask_refs = [] - allowed_types = {ContentType.TEXT, ContentType.SOURCE_REFERENCE} - - for task in tasks: - for subtask in task.sub_tasks: - if subtask.content_type not in allowed_types: - continue - - if subtask.audio_url: - existing_wav = download_bytes( - key=subtask.audio_url, - ) - raw_pcm = existing_wav[wav_header_size:] - else: - wav_bytes = generate_tts_audio( - subtask.content, audio_type, language, voice_name=voice_name - ) - raw_pcm = wav_bytes[wav_header_size:] - - audio_segments.append(raw_pcm) - subtask_refs.append(subtask) + subtask_refs: List[Dict[str, Any]] = [] + + for subtask in subtasks: + if not _is_tts_content_type(subtask.get("content_type")): + continue + + audio_url = subtask.get("audio_url") + if audio_url: + existing_wav = download_bytes(key=audio_url) + raw_pcm = existing_wav[wav_header_size:] + else: + wav_bytes = generate_tts_audio( + subtask.get("content") or "", + audio_type, + language, + voice_name=voice_name, + ) + raw_pcm = wav_bytes[wav_header_size:] + + audio_segments.append(raw_pcm) + subtask_refs.append(subtask) return audio_segments, subtask_refs -def _update_subtask_timestamps( - db: Session, +def _build_subtask_timestamps( audio_segments: List[bytes], - subtask_refs: list, + subtask_refs: List[Dict[str, Any]], sample_rate: int, bytes_per_sample: int, -) -> int: +) -> tuple[int, List[Dict[str, Any]]]: current_offset_ms = 0 + timestamps: List[Dict[str, Any]] = [] for i, raw_pcm in enumerate(audio_segments): segment_samples = len(raw_pcm) // bytes_per_sample segment_duration_ms = int((segment_samples / sample_rate) * 1000) - upsert_sub_task_timestamp( - db=db, - sub_task_id=subtask_refs[i].id, - start_ms=current_offset_ms, - end_ms=current_offset_ms + segment_duration_ms, - created_by="system", + timestamps.append( + { + "sub_task_id": str(subtask_refs[i]["id"]), + "start_ms": current_offset_ms, + "end_ms": current_offset_ms + segment_duration_ms, + } ) current_offset_ms += segment_duration_ms - return current_offset_ms + return current_offset_ms, timestamps def _build_combined_wav(audio_segments: List[bytes]) -> tuple[bytes, int]: @@ -100,30 +99,18 @@ def _build_combined_wav(audio_segments: List[bytes]) -> tuple[bytes, int]: return wav_header + combined_pcm, data_size -def _upload_and_persist_audio( - db: Session, +def _upload_day_audio( combined_wav: bytes, - duration_ms: int, plan_id: UUID, plan_item_id: UUID, -) -> PlanItemAudio: +) -> str: s3_key = f"audio/plan_days/{plan_id}/{plan_item_id}/{uuid4()}.wav" upload_bytes( file_bytes=BytesIO(combined_wav), key=s3_key, content_type=WAV_CONTENT_TYPE, ) - return upsert_plan_item_audio( - db=db, - plan_item_audio=PlanItemAudio( - plan_item_id=plan_item_id, - audio_key=s3_key, - duration_ms=duration_ms, - mime_type=WAV_CONTENT_TYPE, - file_size_bytes=len(combined_wav), - created_by="system", - ), - ) + return s3_key async def _generate_audio_from_text( @@ -151,7 +138,7 @@ async def _generate_audio_from_text( s3_key = f"{s3_key_prefix}/{uuid4()}.wav" else: s3_key = f"audio/generated/{uuid4()}.wav" - + upload_bytes( file_bytes=BytesIO(combined_wav), key=s3_key, @@ -198,41 +185,44 @@ async def generate_plan_audio_service( SAMPLE_RATE = 24000 BYTES_PER_SAMPLE = 2 - with SessionLocal() as db: - plan_item: PlanItem = get_plan_day_by_id_any_plan(db=db, day_id=day_id) - - audio_segments, subtask_refs = _generate_audio_segments( - plan_item.tasks, audio_type, language, voice_name - ) - if not audio_segments: - return [] - - duration_ms = _update_subtask_timestamps( - db=db, - audio_segments=audio_segments, - subtask_refs=subtask_refs, - sample_rate=SAMPLE_RATE, - bytes_per_sample=BYTES_PER_SAMPLE, - ) - - combined_wav, _ = _build_combined_wav(audio_segments) + day_payload = await get_day_audio_generation_payload(day_id=day_id) + audio_segments, subtask_refs = _generate_audio_segments( + day_payload.get("subtasks") or [], + audio_type, + language, + voice_name, + ) + if not audio_segments: + return [] + + duration_ms, timestamps = _build_subtask_timestamps( + audio_segments=audio_segments, + subtask_refs=subtask_refs, + sample_rate=SAMPLE_RATE, + bytes_per_sample=BYTES_PER_SAMPLE, + ) - audio_row = _upload_and_persist_audio( - db=db, - combined_wav=combined_wav, - duration_ms=duration_ms, - plan_id=plan_item.plan_id, - plan_item_id=plan_item.id, - ) + combined_wav, _ = _build_combined_wav(audio_segments) + s3_key = _upload_day_audio( + combined_wav=combined_wav, + plan_id=UUID(str(day_payload["plan_id"])), + plan_item_id=UUID(str(day_payload["id"])), + ) - audio_url = generate_presigned_access_url( - key=audio_row.audio_key, + await apply_day_audio_generation_result( + day_id=UUID(str(day_payload["id"])), + audio_key=s3_key, + duration_ms=duration_ms, + file_size_bytes=len(combined_wav), + timestamps=timestamps, + mime_type=WAV_CONTENT_TYPE, ) + audio_url = generate_presigned_access_url(key=s3_key) return { "audio_url": audio_url, - "audio_duration_ms": audio_row.duration_ms, - "s3_key": audio_row.audio_key, + "audio_duration_ms": duration_ms, + "s3_key": s3_key, } @@ -246,57 +236,36 @@ async def _generate_subtask_audio( BYTES_PER_SAMPLE = 2 WAV_HEADER_SIZE = 44 - with SessionLocal() as db: - subtask: PlanSubTask = get_sub_task_by_subtask_id(db=db, id=sub_task_id) - if not subtask: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": "BAD_REQUEST", "message": "Sub task not found"}, - ) - - allowed_types = {ContentType.TEXT, ContentType.SOURCE_REFERENCE} - if subtask.content_type not in allowed_types: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "BAD_REQUEST", - "message": "Sub task content type must be TEXT or SOURCE_REFERENCE for audio generation", - }, - ) - - wav_bytes = generate_tts_audio( - subtask.content, audio_type, language, voice_name=voice_name - ) - raw_pcm = wav_bytes[WAV_HEADER_SIZE:] - - segment_samples = len(raw_pcm) // BYTES_PER_SAMPLE - duration_ms = int((segment_samples / SAMPLE_RATE) * 1000) - - combined_wav, _ = _build_combined_wav([raw_pcm]) + subtask = await get_sub_task_audio_generation_payload(sub_task_id=sub_task_id) - s3_key = f"audio/plan_subtasks/{subtask.task_id}/{sub_task_id}/{uuid4()}.wav" - upload_bytes( - file_bytes=BytesIO(combined_wav), - key=s3_key, - content_type=WAV_CONTENT_TYPE, - ) + wav_bytes = generate_tts_audio( + subtask.get("content") or "", + audio_type, + language, + voice_name=voice_name, + ) + raw_pcm = wav_bytes[WAV_HEADER_SIZE:] - subtask.audio_url = s3_key - subtask.duration = str(duration_ms) - db.commit() + segment_samples = len(raw_pcm) // BYTES_PER_SAMPLE + duration_ms = int((segment_samples / SAMPLE_RATE) * 1000) - upsert_sub_task_timestamp( - db=db, - sub_task_id=sub_task_id, - start_ms=0, - end_ms=duration_ms, - created_by="system", - ) + combined_wav, _ = _build_combined_wav([raw_pcm]) - audio_url = generate_presigned_access_url( + task_id = UUID(str(subtask["task_id"])) + s3_key = f"audio/plan_subtasks/{task_id}/{sub_task_id}/{uuid4()}.wav" + upload_bytes( + file_bytes=BytesIO(combined_wav), key=s3_key, + content_type=WAV_CONTENT_TYPE, + ) + + await apply_sub_task_audio_generation_result( + sub_task_id=sub_task_id, + audio_key=s3_key, + duration_ms=duration_ms, ) + audio_url = generate_presigned_access_url(key=s3_key) return { "audio_url": audio_url, "audio_duration_ms": duration_ms, diff --git a/worker_api/audio/services/audio_job_consumer.py b/worker_api/audio/services/audio_job_consumer.py new file mode 100644 index 0000000..20423f9 --- /dev/null +++ b/worker_api/audio/services/audio_job_consumer.py @@ -0,0 +1,166 @@ +import asyncio +import logging +from typing import Any, Dict, Optional +from uuid import UUID + +from fastapi import HTTPException + +from worker_api.audio.enums import AudioJobStatus, MonlamVoiceName, PlanAudioType +from worker_api.audio.services.audio_generate_service import generate_plan_audio_service +from worker_api.audio.services.backend_client import ( + get_audio_job_status, + update_audio_job_status, +) +from worker_api.audio.sqs_client import ( + delete_audio_job_message, + is_audio_sqs_poll_enabled, + parse_audio_job_message_body, + receive_audio_job_messages, +) + +logger = logging.getLogger(__name__) + +_POLL_IDLE_SECONDS = 5 +_POLL_ERROR_SECONDS = 10 +_TERMINAL_STATUSES = { + AudioJobStatus.COMPLETED.value, + AudioJobStatus.FAILED.value, +} + + +def _parse_uuid(value: Any) -> Optional[UUID]: + if value is None or value == "": + return None + return UUID(str(value)) + + +def _parse_audio_type(value: Any) -> PlanAudioType: + if isinstance(value, PlanAudioType): + return value + if value is None or value == "": + return PlanAudioType.TEXT_READING + return PlanAudioType(str(value)) + + +def _parse_voice_name(value: Any) -> MonlamVoiceName: + if isinstance(value, MonlamVoiceName): + return value + if value is None or value == "": + return MonlamVoiceName.DOLKAR_LHASA_FEMALE + return MonlamVoiceName(str(value)) + + +def _normalize_result(result: Any) -> Dict[str, Any]: + if not isinstance(result, dict): + raise ValueError("Audio generation returned no result") + if not result.get("s3_key") and not result.get("audio_url"): + raise ValueError("Audio generation returned an empty result") + return { + "audio_url": result.get("audio_url"), + "audio_duration_ms": result.get("audio_duration_ms"), + "s3_key": result.get("s3_key"), + } + + +def _error_detail(exc: Exception) -> str: + if isinstance(exc, HTTPException): + detail = exc.detail + if isinstance(detail, dict): + return str(detail.get("message") or detail) + return str(detail) + return str(exc) + + +async def process_audio_job_message(message: Dict[str, Any]) -> None: + receipt_handle = message.get("ReceiptHandle") + body = parse_audio_job_message_body(message.get("Body", "")) + if not body: + if receipt_handle: + delete_audio_job_message(receipt_handle) + return + + job_id = _parse_uuid(body.get("job_id")) + if not job_id: + logger.error("Invalid job_id in audio SQS message: %s", body) + if receipt_handle: + delete_audio_job_message(receipt_handle) + return + + existing = await get_audio_job_status(job_id=job_id) + if not existing: + logger.error("Audio job not found on backend: %s", job_id) + if receipt_handle: + delete_audio_job_message(receipt_handle) + return + + existing_status = str(existing.get("status") or "") + if existing_status in _TERMINAL_STATUSES: + logger.info("Skipping already finished audio job %s (%s)", job_id, existing_status) + if receipt_handle: + delete_audio_job_message(receipt_handle) + return + + await update_audio_job_status(job_id=job_id, status=AudioJobStatus.PROCESSING) + + try: + day_id = _parse_uuid(body.get("day_id")) + sub_task_id = _parse_uuid(body.get("sub_task_id")) + language = str(body.get("language") or "") + if not language: + raise ValueError("language is required") + if not day_id and not sub_task_id: + raise ValueError("Exactly one of day_id or sub_task_id is required") + + result = await generate_plan_audio_service( + language=language, + day_id=day_id, + sub_task_id=sub_task_id, + audio_type=_parse_audio_type(body.get("type")), + voice_name=_parse_voice_name(body.get("voice_name")), + ) + normalized = _normalize_result(result) + + await update_audio_job_status( + job_id=job_id, + status=AudioJobStatus.COMPLETED, + result=normalized, + ) + logger.info("Completed audio job %s", job_id) + except Exception as exc: + logger.exception("Failed audio job %s", job_id) + await update_audio_job_status( + job_id=job_id, + status=AudioJobStatus.FAILED, + error_message=_error_detail(exc), + ) + + if receipt_handle: + delete_audio_job_message(receipt_handle) + + +async def run_audio_sqs_consumer(stop_event: asyncio.Event) -> None: + logger.info("Audio SQS consumer started") + while not stop_event.is_set(): + if not is_audio_sqs_poll_enabled(): + try: + await asyncio.wait_for(stop_event.wait(), timeout=_POLL_IDLE_SECONDS) + except asyncio.TimeoutError: + pass + continue + + try: + messages = await asyncio.to_thread(receive_audio_job_messages) + if not messages: + continue + for message in messages: + if stop_event.is_set(): + break + await process_audio_job_message(message) + except Exception: + logger.exception("Audio SQS consumer loop error") + try: + await asyncio.wait_for(stop_event.wait(), timeout=_POLL_ERROR_SECONDS) + except asyncio.TimeoutError: + pass + + logger.info("Audio SQS consumer stopped") diff --git a/worker_api/audio/services/backend_client.py b/worker_api/audio/services/backend_client.py new file mode 100644 index 0000000..f09cec6 --- /dev/null +++ b/worker_api/audio/services/backend_client.py @@ -0,0 +1,132 @@ +from typing import Any, Dict, List, Optional +from uuid import UUID + +import httpx +from fastapi import HTTPException + +from worker_api.audio.enums import AudioJobStatus +from worker_api.config import get + + +def _backend_url() -> str: + backend_url = get("BACKEND_API_URL").rstrip("/") + if not backend_url: + raise RuntimeError("BACKEND_API_URL is not configured") + return backend_url + + +def _dispatch_headers() -> Dict[str, str]: + dispatch_token = get("NOTIFICATION_DISPATCH_SECRET_TOKEN") + if not dispatch_token: + raise RuntimeError("NOTIFICATION_DISPATCH_SECRET_TOKEN is not configured") + return {"X-Dispatch-Token": dispatch_token} + + +def _raise_for_backend_response(response: httpx.Response) -> None: + if response.is_success: + return + detail: Any + try: + detail = response.json() + except ValueError: + detail = response.text or f"Backend request failed with status {response.status_code}" + raise HTTPException(status_code=response.status_code, detail=detail) + + +async def get_audio_job_status(job_id: UUID) -> Optional[Dict[str, Any]]: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get( + f"{_backend_url()}/internal/audio/jobs/{job_id}", + headers=_dispatch_headers(), + ) + if response.status_code == 404: + return None + _raise_for_backend_response(response) + return response.json() + + +async def update_audio_job_status( + *, + job_id: UUID, + status: AudioJobStatus, + result: Optional[Dict[str, Any]] = None, + error_message: Optional[str] = None, +) -> Dict[str, Any]: + payload: Dict[str, Any] = {"status": status.value} + if result is not None: + payload["result"] = result + if error_message is not None: + payload["error_message"] = error_message + + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.patch( + f"{_backend_url()}/internal/audio/jobs/{job_id}", + headers=_dispatch_headers(), + json=payload, + ) + _raise_for_backend_response(response) + return response.json() + + +async def get_day_audio_generation_payload(day_id: UUID) -> Dict[str, Any]: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get( + f"{_backend_url()}/internal/audio/days/{day_id}/generation-payload", + headers=_dispatch_headers(), + ) + _raise_for_backend_response(response) + return response.json() + + +async def get_sub_task_audio_generation_payload(sub_task_id: UUID) -> Dict[str, Any]: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get( + f"{_backend_url()}/internal/audio/sub-tasks/{sub_task_id}/generation-payload", + headers=_dispatch_headers(), + ) + _raise_for_backend_response(response) + return response.json() + + +async def apply_day_audio_generation_result( + *, + day_id: UUID, + audio_key: str, + duration_ms: int, + file_size_bytes: int, + timestamps: List[Dict[str, Any]], + mime_type: str = "audio/wav", +) -> None: + payload = { + "audio_key": audio_key, + "duration_ms": duration_ms, + "mime_type": mime_type, + "file_size_bytes": file_size_bytes, + "timestamps": timestamps, + } + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{_backend_url()}/internal/audio/days/{day_id}/generation-result", + headers=_dispatch_headers(), + json=payload, + ) + _raise_for_backend_response(response) + + +async def apply_sub_task_audio_generation_result( + *, + sub_task_id: UUID, + audio_key: str, + duration_ms: int, +) -> None: + payload = { + "audio_key": audio_key, + "duration_ms": duration_ms, + } + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{_backend_url()}/internal/audio/sub-tasks/{sub_task_id}/generation-result", + headers=_dispatch_headers(), + json=payload, + ) + _raise_for_backend_response(response) diff --git a/worker_api/audio/sqs_client.py b/worker_api/audio/sqs_client.py new file mode 100644 index 0000000..461244e --- /dev/null +++ b/worker_api/audio/sqs_client.py @@ -0,0 +1,82 @@ +import json +import logging +from typing import Any, Dict, List, Optional + +import boto3 +from botocore.exceptions import ClientError + +from worker_api.config import get, get_bool, get_int + +logger = logging.getLogger(__name__) + +_sqs_client = None + + +def _get_sqs_client(): + global _sqs_client + if _sqs_client is None: + _sqs_client = boto3.client( + "sqs", + aws_access_key_id=get("AWS_ACCESS_KEY"), + aws_secret_access_key=get("AWS_SECRET_KEY"), + region_name=get("AWS_REGION"), + ) + return _sqs_client + + +def get_audio_sqs_queue_url() -> str: + return get("AUDIO_SQS_QUEUE_URL").strip() + + +def is_audio_sqs_configured() -> bool: + return bool(get_audio_sqs_queue_url()) + + +def is_audio_sqs_poll_enabled() -> bool: + return get_bool("AUDIO_SQS_POLL_ENABLED") and is_audio_sqs_configured() + + +def receive_audio_job_messages() -> List[Dict[str, Any]]: + queue_url = get_audio_sqs_queue_url() + if not queue_url: + return [] + + try: + response = _get_sqs_client().receive_message( + QueueUrl=queue_url, + MaxNumberOfMessages=get_int("AUDIO_SQS_MAX_MESSAGES"), + WaitTimeSeconds=get_int("AUDIO_SQS_WAIT_TIME_SECONDS"), + VisibilityTimeout=get_int("AUDIO_SQS_VISIBILITY_TIMEOUT_SECONDS"), + MessageAttributeNames=["All"], + ) + return response.get("Messages", []) + except ClientError as e: + logger.error("Failed to receive audio SQS messages: %s", e) + return [] + + +def delete_audio_job_message(receipt_handle: str) -> None: + queue_url = get_audio_sqs_queue_url() + if not queue_url or not receipt_handle: + return + + try: + _get_sqs_client().delete_message( + QueueUrl=queue_url, + ReceiptHandle=receipt_handle, + ) + except ClientError as e: + logger.error("Failed to delete audio SQS message: %s", e) + + +def parse_audio_job_message_body(raw_body: str) -> Optional[Dict[str, Any]]: + try: + body = json.loads(raw_body) + except (TypeError, json.JSONDecodeError) as e: + logger.error("Invalid audio SQS message body: %s", e) + return None + + if not isinstance(body, dict) or not body.get("job_id"): + logger.error("Audio SQS message missing job_id: %s", body) + return None + return body diff --git a/worker_api/config.py b/worker_api/config.py index 085e008..c5e2614 100644 --- a/worker_api/config.py +++ b/worker_api/config.py @@ -29,6 +29,13 @@ MONLAM_TTS_MODEL_NAME="", MONLAM_TTS_VOICE_NAME="", + # Audio job SQS consumer (backend producer → worker consumer) + AUDIO_SQS_QUEUE_URL="", + AUDIO_SQS_WAIT_TIME_SECONDS=20, + AUDIO_SQS_VISIBILITY_TIMEOUT_SECONDS=900, + AUDIO_SQS_MAX_MESSAGES=1, + AUDIO_SQS_POLL_ENABLED="true", + # Notification dispatch (Cloud Scheduler -> worker) NOTIFICATION_DISPATCH_SECRET_TOKEN="Dispatch", NOTIFICATION_DISPATCH_BATCH_SIZE=100, diff --git a/worker_api/db/mongo_database.py b/worker_api/db/mongo_database.py index 2e4b184..08fd8ff 100644 --- a/worker_api/db/mongo_database.py +++ b/worker_api/db/mongo_database.py @@ -1,3 +1,4 @@ +import asyncio import logging from contextlib import asynccontextmanager @@ -6,6 +7,7 @@ from motor.motor_asyncio import AsyncIOMotorClient from ..config import get +from worker_api.audio.services.audio_job_consumer import run_audio_sqs_consumer mongodb_client = None mongodb = None @@ -29,7 +31,20 @@ async def lifespan(api: FastAPI): logging.error(f"Error during collection initialization: {e}") raise + stop_event = asyncio.Event() + consumer_task = asyncio.create_task(run_audio_sqs_consumer(stop_event)) + yield + stop_event.set() + try: + await asyncio.wait_for(consumer_task, timeout=5) + except asyncio.TimeoutError: + consumer_task.cancel() + try: + await consumer_task + except asyncio.CancelledError: + pass + if mongodb_client: mongodb_client.close()