atom-audio-engine 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

audio_engine/tests/test_personaplex/test_personaplex_message.py

@@ -0,0 +1,80 @@
+"""Tests for PersonaPlexMessage (Step 2)."""
+
+import pytest
+from pipelines.personaplex import PersonaPlexMessage, MessageType
+
+
+class TestPersonaPlexMessage:
+    """Test WebSocket message encoding and decoding."""
+
+    def test_encode_audio(self):
+        """Can we encode an audio message?"""
+        audio_data = b"fake_opus_audio"
+        msg = PersonaPlexMessage(type=MessageType.AUDIO, data=audio_data)
+        encoded = msg.encode()
+
+        assert encoded[0] == 0x01
+        assert encoded[1:] == audio_data
+
+    def test_encode_text(self):
+        """Can we encode a text message?"""
+        text = "Hello, world!"
+        msg = PersonaPlexMessage(type=MessageType.TEXT, data=text)
+        encoded = msg.encode()
+
+        assert encoded[0] == 0x02
+        assert encoded[1:] == text.encode("utf-8")
+
+    def test_decode_audio(self):
+        """Can we decode an audio message?"""
+        audio_data = b"fake_opus_audio"
+        raw = bytes([0x01]) + audio_data
+
+        msg = PersonaPlexMessage.decode(raw)
+        assert msg.type == MessageType.AUDIO
+        assert msg.data == audio_data
+
+    def test_decode_text(self):
+        """Can we decode a text message?"""
+        text = "Hello, assistant!"
+        raw = bytes([0x02]) + text.encode("utf-8")
+
+        msg = PersonaPlexMessage.decode(raw)
+        assert msg.type == MessageType.TEXT
+        assert msg.data == text
+
+    def test_decode_error(self):
+        """Can we decode an error message?"""
+        error_text = b"Connection timeout"
+        raw = bytes([0x05]) + error_text
+
+        msg = PersonaPlexMessage.decode(raw)
+        assert msg.type == MessageType.ERROR
+        assert msg.data == error_text
+
+    def test_roundtrip_audio(self):
+        """Audio message roundtrip: encode then decode?"""
+        original_data = b"\x00\x01\x02\x03\x04\x05"
+        msg1 = PersonaPlexMessage(type=MessageType.AUDIO, data=original_data)
+
+        encoded = msg1.encode()
+        msg2 = PersonaPlexMessage.decode(encoded)
+
+        assert msg2.type == msg1.type
+        assert msg2.data == original_data
+
+    def test_roundtrip_text(self):
+        """Text message roundtrip: encode then decode?"""
+        original_text = "This is a test message with émojis 🎉"
+        msg1 = PersonaPlexMessage(type=MessageType.TEXT, data=original_text)
+
+        encoded = msg1.encode()
+        msg2 = PersonaPlexMessage.decode(encoded)
+
+        assert msg2.type == msg1.type
+        assert msg2.data == original_text
+
+    def test_decode_too_short(self):
+        """Does decode reject empty message?"""
+        with pytest.raises(ValueError):
+            PersonaPlexMessage.decode(b"")

audio_engine/tests/test_personaplex/test_personaplex_pipeline.py

@@ -0,0 +1,226 @@
+"""Tests for PersonaPlexPipeline lifecycle and orchestration (Step 6-7)."""
+
+import pytest
+from unittest.mock import AsyncMock
+
+from pipelines.personaplex import (
+    PersonaPlexPipeline,
+    PersonaPlexConfig,
+    SessionData,
+    AudioChunk,
+    TextChunk,
+)
+
+
+class TestPersonaPlexPipelineInit:
+    """Test pipeline initialization."""
+
+    def test_pipeline_init_with_defaults(self):
+        """Can we create a pipeline with defaults?"""
+        pipeline = PersonaPlexPipeline()
+
+        assert pipeline.session_id is not None
+        assert pipeline.session_data is not None
+        assert pipeline.system_prompt is not None
+
+    def test_pipeline_init_with_custom_config(self):
+        """Can we create pipeline with custom config?"""
+        config = PersonaPlexConfig(voice_prompt="NATM0.pt")
+        pipeline = PersonaPlexPipeline(config=config)
+
+        assert pipeline.config.voice_prompt == "NATM0.pt"
+
+    def test_pipeline_init_creates_session_data(self):
+        """Does init create a SessionData object?"""
+        pipeline = PersonaPlexPipeline(system_prompt="Custom prompt")
+
+        assert isinstance(pipeline.session_data, SessionData)
+        assert pipeline.session_data.system_prompt == "Custom prompt"
+
+    def test_pipeline_is_not_running_initially(self):
+        """Is pipeline not running after init?"""
+        pipeline = PersonaPlexPipeline()
+
+        assert not pipeline._is_running
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexPipelineLifecycle:
+    """Test pipeline start/stop lifecycle."""
+
+    async def test_pipeline_start_connects_client(self):
+        """Does start() connect the client?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline.client = AsyncMock()
+
+        await pipeline.start()
+
+        assert pipeline._is_running
+        pipeline.client.connect.assert_called_once()
+
+    async def test_pipeline_start_creates_receive_task(self):
+        """Does start() create a receive task?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline.client = AsyncMock()
+        pipeline.client.stream_messages = AsyncMock(
+            return_value=AsyncMock(
+                __aiter__=AsyncMock(
+                    return_value=AsyncMock(
+                        __anext__=AsyncMock(side_effect=StopAsyncIteration)
+                    )
+                )
+            )
+        )
+
+        await pipeline.start()
+
+        assert pipeline._receive_task is not None
+
+    async def test_pipeline_stop_disconnects_client(self):
+        """Does stop() disconnect the client?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline.client = AsyncMock()
+        pipeline._is_running = True
+
+        result = await pipeline.stop()
+
+        pipeline.client.disconnect.assert_called_once()
+        assert not pipeline._is_running
+
+    async def test_pipeline_context_manager(self):
+        """Does context manager handle lifecycle?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline.client = AsyncMock()
+        pipeline.client.stream_messages = AsyncMock(
+            return_value=AsyncMock(
+                __aiter__=AsyncMock(
+                    return_value=AsyncMock(
+                        __anext__=AsyncMock(side_effect=StopAsyncIteration)
+                    )
+                )
+            )
+        )
+
+        async with pipeline:
+            assert pipeline._is_running
+
+        pipeline.client.disconnect.assert_called_once()
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexPipelineSendAudio:
+    """Test sending audio."""
+
+    async def test_pipeline_send_audio_when_running(self):
+        """Can we send audio when pipeline is running?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline.client = AsyncMock()
+        pipeline._is_running = True
+
+        audio = b"opus_audio_data"
+        await pipeline.send_audio(audio)
+
+        pipeline.client.send_audio.assert_called_once_with(audio)
+
+    async def test_pipeline_send_audio_fails_when_stopped(self):
+        """Does send_audio fail when not running?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline._is_running = False
+
+        with pytest.raises(RuntimeError):
+            await pipeline.send_audio(b"audio")
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexPipelineStreaming:
+    """Test streaming interface."""
+
+    async def test_pipeline_stream_requires_running(self):
+        """Does stream() require pipeline to be running?"""
+        pipeline = PersonaPlexPipeline()
+        pipeline._is_running = False
+
+        with pytest.raises(RuntimeError):
+            async for _ in pipeline.stream():
+                pass
+
+    async def test_pipeline_stream_yields_tuples(self):
+        """Does stream() yield (audio, text) tuples?"""
+        import asyncio
+
+        pipeline = PersonaPlexPipeline()
+        pipeline._is_running = True
+
+        # Pre-populate queues
+        audio_chunk = AudioChunk(data=b"audio", sample_rate=48000)
+        text_chunk = TextChunk(text="Hello")
+
+        await pipeline._audio_queue.put(audio_chunk)
+        await pipeline._text_queue.put(text_chunk)
+
+        # Stop after one iteration
+        async def stop_soon():
+            await asyncio.sleep(0.01)
+            pipeline._is_running = False
+
+        asyncio.create_task(stop_soon())
+
+        results = []
+        async for audio, text in pipeline.stream():
+            results.append((audio, text))
+            if len(results) >= 1:
+                break
+
+        # Should get at least one result
+        assert len(results) >= 1
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexPipelineTranscript:
+    """Test transcript management."""
+
+    def test_pipeline_session_data_stores_messages(self):
+        """Does session_data store conversation messages?"""
+        pipeline = PersonaPlexPipeline()
+
+        pipeline.session_data.add_message("user", "Hello")
+        pipeline.session_data.add_message("assistant", "Hi there!")
+
+        assert len(pipeline.session_data.messages) == 2
+        assert pipeline.session_data.messages[0].role == "user"
+        assert pipeline.session_data.messages[1].role == "assistant"
+
+    async def test_pipeline_stop_can_save_transcript(self):
+        """Can pipeline save transcript when configured?"""
+        import tempfile
+        from pathlib import Path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config = PersonaPlexConfig(transcript_path=tmpdir, save_transcripts=True)
+            pipeline = PersonaPlexPipeline(config=config)
+            pipeline.client = AsyncMock()
+            pipeline._is_running = True  # Pretend running so stop works
+
+            # Add a message to transcript
+            pipeline.session_data.add_message("user", "Test message")
+
+            # Stop (should save because save_transcripts=True)
+            result = await pipeline.stop()
+
+            # Check transcript was saved if result is not None
+            if result is not None:
+                transcript_files = list(Path(tmpdir).glob("*.json"))
+                assert len(transcript_files) > 0
+
+    def test_pipeline_session_data_to_dict(self):
+        """Can we convert session data to dict?"""
+        pipeline = PersonaPlexPipeline(system_prompt="Test")
+        pipeline.session_data.add_message("user", "Hello")
+
+        data = pipeline.session_data.to_dict()
+
+        assert "session_id" in data
+        assert "timestamp" in data
+        assert "system_prompt" in data
+        assert "messages" in data
+        assert len(data["messages"]) == 1

audio_engine/tests/test_personaplex/test_personaplex_session.py

@@ -0,0 +1,184 @@
+"""Tests for session management utilities (Step 4)."""
+
+from datetime import datetime, UTC
+import uuid
+
+from pipelines.personaplex import (
+    generate_session_id,
+    get_timestamp_iso,
+    list_transcripts,
+    format_transcript_for_display,
+)
+
+
+class TestSessionUtilities:
+    """Test session ID and timestamp generation."""
+
+    def test_generate_session_id_returns_string(self):
+        """Does generate_session_id return a string?"""
+        session_id = generate_session_id()
+
+        assert isinstance(session_id, str)
+        assert len(session_id) > 0
+
+    def test_generate_session_id_is_uuid4(self):
+        """Does it generate valid UUID4 format?"""
+        session_id = generate_session_id()
+
+        # Should be valid UUID
+        uuid.UUID(session_id)
+
+    def test_generate_session_id_unique(self):
+        """Do consecutive calls generate different IDs?"""
+        id1 = generate_session_id()
+        id2 = generate_session_id()
+
+        assert id1 != id2
+
+    def test_get_timestamp_iso_returns_string(self):
+        """Does get_timestamp_iso return a string?"""
+        timestamp = get_timestamp_iso()
+
+        assert isinstance(timestamp, str)
+        assert len(timestamp) > 0
+
+    def test_get_timestamp_iso_ends_with_z(self):
+        """Does timestamp end with Z suffix?"""
+        timestamp = get_timestamp_iso()
+
+        assert timestamp.endswith("Z")
+
+    def test_get_timestamp_iso_is_valid_iso8601(self):
+        """Is the timestamp in ISO 8601 format?"""
+        timestamp = get_timestamp_iso()
+
+        # Remove Z suffix and parse
+        iso_part = timestamp[:-1]
+        parsed = datetime.fromisoformat(iso_part)
+
+        assert isinstance(parsed, datetime)
+
+    def test_get_timestamp_iso_close_to_now(self):
+        """Is generated timestamp close to current time?"""
+        before = datetime.now(UTC)
+        timestamp = get_timestamp_iso()
+        after = datetime.now(UTC)
+
+        # Parse timestamp (remove Z and convert to UTC-aware)
+        iso_part = timestamp[:-1]
+        parsed = datetime.fromisoformat(iso_part).replace(tzinfo=UTC)
+
+        # Should be within 1 second of now
+        assert before <= parsed <= after
+
+
+class TestTranscriptUtilities:
+    """Test transcript management utilities."""
+
+    def test_format_transcript_for_display_empty(self):
+        """Can we format an empty transcript?"""
+        from pipelines.personaplex import SessionData
+
+        session = SessionData(
+            session_id="test-format",
+            timestamp="2025-01-01T00:00:00Z",
+            system_prompt="You are helpful",
+            voice_prompt="NATF0.pt",
+            messages=[],
+        )
+
+        formatted = format_transcript_for_display(session)
+
+        assert isinstance(formatted, str)
+
+    def test_format_transcript_for_display_with_messages(self):
+        """Does formatting show user/assistant messages?"""
+        from pipelines.personaplex import SessionData, TranscriptMessage
+
+        msg1 = TranscriptMessage(
+            role="user",
+            text="Hello",
+            timestamp="2025-01-01T00:00:01Z",
+        )
+        msg2 = TranscriptMessage(
+            role="assistant",
+            text="Hi there!",
+            timestamp="2025-01-01T00:00:02Z",
+        )
+        session = SessionData(
+            session_id="test-format-msg",
+            timestamp="2025-01-01T00:00:00Z",
+            system_prompt="You are helpful",
+            voice_prompt="NATF0.pt",
+            messages=[msg1, msg2],
+        )
+
+        formatted = format_transcript_for_display(session)
+
+        # Should contain message text
+        assert "Hello" in formatted
+        assert "Hi there!" in formatted
+
+    def test_format_transcript_for_display_labels_speakers(self):
+        """Does formatting include speaker labels?"""
+        from pipelines.personaplex import SessionData, TranscriptMessage
+
+        msg = TranscriptMessage(
+            role="user",
+            text="Test message",
+            timestamp="2025-01-01T00:00:01Z",
+        )
+        session = SessionData(
+            session_id="test-labels",
+            timestamp="2025-01-01T00:00:00Z",
+            system_prompt="You are helpful",
+            voice_prompt="NATF0.pt",
+            messages=[msg],
+        )
+
+        formatted = format_transcript_for_display(session)
+
+        # Should show role/speaker
+        assert "user" in formatted.lower() or "User" in formatted
+
+    def test_list_transcripts_returns_list(self):
+        """Does list_transcripts return a list?"""
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            transcripts = list_transcripts(tmpdir)
+
+            assert isinstance(transcripts, list)
+
+    def test_list_transcripts_empty_dir(self):
+        """Does it return empty list for empty directory?"""
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            transcripts = list_transcripts(tmpdir)
+
+            assert len(transcripts) == 0
+
+    def test_list_transcripts_finds_json_files(self):
+        """Does it find transcript JSON files?"""
+        import tempfile
+        from pathlib import Path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            from pipelines.personaplex import save_transcript, SessionData
+
+            # Create a transcript
+            session = SessionData(
+                session_id="list-test",
+                timestamp="2025-01-01T00:00:00Z",
+                system_prompt="You are helpful",
+                voice_prompt="NATF0.pt",
+                messages=[],
+            )
+            save_transcript(session, tmpdir)
+
+            # List should find it
+            transcripts = list_transcripts(tmpdir)
+
+            assert len(transcripts) >= 1
+            assert any("list-test" in str(t) for t in transcripts)

audio_engine/tests/test_personaplex/test_personaplex_transcript.py

@@ -0,0 +1,184 @@
+"""Tests for transcript persistence (Step 3)."""
+
+import json
+import tempfile
+from pathlib import Path
+
+from pipelines.personaplex import (
+    SessionData,
+    TranscriptMessage,
+    save_transcript,
+    load_transcript,
+)
+
+
+class TestTranscriptPersistence:
+    """Test transcript saving and loading."""
+
+    def test_save_transcript_creates_file(self):
+        """Does save_transcript create a JSON file?"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            session = SessionData(
+                session_id="test-123",
+                timestamp="2025-01-01T00:00:00Z",
+                system_prompt="You are helpful",
+                voice_prompt="NATF0.pt",
+                messages=[],
+            )
+
+            saved_path = save_transcript(session, Path(tmpdir))
+
+            assert saved_path.exists()
+            assert saved_path.suffix == ".json"
+
+    def test_save_transcript_contains_correct_structure(self):
+        """Does saved JSON have session_id, timestamp, and messages?"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            msg = TranscriptMessage(
+                role="user",
+                text="Hello",
+                timestamp="2025-01-01T00:00:01Z",
+            )
+            session = SessionData(
+                session_id="test-456",
+                timestamp="2025-01-01T00:00:00Z",
+                system_prompt="You are helpful",
+                voice_prompt="NATF0.pt",
+                messages=[msg],
+            )
+
+            saved_path = save_transcript(session, Path(tmpdir))
+
+            with open(saved_path, "r") as f:
+                data = json.load(f)
+
+            assert data["session_id"] == "test-456"
+            assert data["timestamp"] == "2025-01-01T00:00:00Z"
+            assert len(data["messages"]) == 1
+            assert data["messages"][0]["role"] == "user"
+            assert data["messages"][0]["text"] == "Hello"
+
+    def test_load_transcript_reads_file(self):
+        """Can we load a transcript from JSON?"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create and save a session
+            msg = TranscriptMessage(
+                role="assistant",
+                text="Hi there!",
+                timestamp="2025-01-01T00:00:02Z",
+            )
+            original = SessionData(
+                session_id="test-789",
+                timestamp="2025-01-01T00:00:00Z",
+                system_prompt="You are helpful",
+                voice_prompt="NATF0.pt",
+                messages=[msg],
+            )
+
+            saved_path = save_transcript(original, Path(tmpdir))
+
+            # Load it back
+            loaded = load_transcript(saved_path)
+
+            assert loaded.session_id == "test-789"
+            assert loaded.timestamp == "2025-01-01T00:00:00Z"
+            assert len(loaded.messages) == 1
+            assert loaded.messages[0].role == "assistant"
+            assert loaded.messages[0].text == "Hi there!"
+
+    def test_transcript_roundtrip(self):
+        """Session → save → load → Session equals original?"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            msg1 = TranscriptMessage(
+                role="user",
+                text="First message",
+                timestamp="2025-01-01T00:00:01Z",
+            )
+            msg2 = TranscriptMessage(
+                role="assistant",
+                text="Response",
+                timestamp="2025-01-01T00:00:02Z",
+            )
+            original = SessionData(
+                session_id="roundtrip-123",
+                timestamp="2025-01-01T00:00:00Z",
+                system_prompt="You are helpful",
+                voice_prompt="NATF0.pt",
+                messages=[msg1, msg2],
+            )
+
+            saved_path = save_transcript(original, Path(tmpdir))
+            loaded = load_transcript(saved_path)
+
+            # Compare all fields
+            assert loaded.session_id == original.session_id
+            assert loaded.timestamp == original.timestamp
+            assert loaded.system_prompt == original.system_prompt
+            assert loaded.voice_prompt == original.voice_prompt
+            assert len(loaded.messages) == len(original.messages)
+            for i, msg in enumerate(original.messages):
+                assert loaded.messages[i].role == msg.role
+                assert loaded.messages[i].text == msg.text
+                assert loaded.messages[i].timestamp == msg.timestamp
+
+    def test_transcript_add_message(self):
+        """Can we add messages to a SessionData?"""
+        session = SessionData(
+            session_id="add-test",
+            timestamp="2025-01-01T00:00:00Z",
+            system_prompt="You are helpful",
+            voice_prompt="NATF0.pt",
+            messages=[],
+        )
+
+        assert len(session.messages) == 0
+
+        session.add_message("user", "Test")
+
+        assert len(session.messages) == 1
+        assert session.messages[0].text == "Test"
+
+    def test_transcript_to_dict(self):
+        """Can we convert SessionData to dict?"""
+        msg = TranscriptMessage(
+            role="user",
+            text="Convert me",
+            timestamp="2025-01-01T00:00:01Z",
+        )
+        session = SessionData(
+            session_id="dict-test",
+            timestamp="2025-01-01T00:00:00Z",
+            system_prompt="You are helpful",
+            voice_prompt="NATF0.pt",
+            messages=[msg],
+        )
+
+        data = session.to_dict()
+
+        assert isinstance(data, dict)
+        assert data["session_id"] == "dict-test"
+        assert data["timestamp"] == "2025-01-01T00:00:00Z"
+        assert len(data["messages"]) == 1
+
+    def test_transcript_from_dict(self):
+        """Can we create SessionData from dict?"""
+        data = {
+            "session_id": "from-dict-test",
+            "timestamp": "2025-01-01T00:00:00Z",
+            "system_prompt": "You are helpful",
+            "voice_prompt": "NATF0.pt",
+            "messages": [
+                {
+                    "role": "user",
+                    "text": "Recreate me",
+                    "timestamp": "2025-01-01T00:00:01Z",
+                }
+            ],
+        }
+
+        session = SessionData.from_dict(data)
+
+        assert session.session_id == "from-dict-test"
+        assert session.system_prompt == "You are helpful"
+        assert len(session.messages) == 1
+        assert session.messages[0].text == "Recreate me"

audio_engine/tests/test_traditional_pipeline/__init__.py

@@ -0,0 +1 @@
+"""Tests for the audio engine."""