test_space_app.py

"""Test suite for ZeroGPU Space app and handler integration."""

import os
import sys
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

# Ensure local space directory is in sys.path
space_dir = Path(__file__).resolve().parent
if str(space_dir) not in sys.path:
    sys.path.insert(0, str(space_dir))

# Provide lightweight mock for gradio if not installed in host environment
try:
    import gradio as gr
except ImportError:
    from unittest.mock import MagicMock
    gr = MagicMock()
    gr.Error = Exception
    sys.modules["gradio"] = gr

# Load space app via importlib to avoid namespace collision with backend/app
import importlib.util
app_path = space_dir / "app.py"
spec = importlib.util.spec_from_file_location("space_app", str(app_path))
space_app = importlib.util.module_from_spec(spec)
spec.loader.exec_module(space_app)

EndpointHandler = space_app.EndpointHandler


class MockWord:
    def __init__(self, word: str, start: float, end: float):
        self.word = word
        self.start = start
        self.end = end


class MockSegment:
    def __init__(self, text: str, start: float, end: float, words=None):
        self.text = text
        self.start = start
        self.end = end
        self.words = words or []


class TestSpaceApp(unittest.TestCase):
    def test_lazy_initialization_does_not_load_models_at_startup(self):
        """Verifies that EndpointHandler(lazy=True) does not load heavy models at startup."""
        handler = EndpointHandler(lazy=True)
        self.assertIsNone(handler.whisper_model)
        self.assertIsNone(handler.diarization_pipeline)

    def test_app_transcribe_validation(self):
        """Verifies that transcribe raises an error if audio path is missing or invalid."""
        import gradio as gr

        with self.assertRaises(gr.Error):
            space_app._run_transcription(audio_path=None)

        with self.assertRaises(gr.Error):
            space_app._run_transcription(audio_path="non_existent_meeting_file.wav")

    @patch("handler.EndpointHandler._load_models")
    @patch("handler.EndpointHandler._prepare_audio")
    def test_app_transcribe_end_to_end_payload(self, mock_prepare, mock_load):
        """Tests that Gradio inputs are formatted correctly into EndpointHandler payload and return structured JSON."""
        mock_prepare.return_value = ("/tmp/mock_audio.wav", 10.0)

        # Mock faster-whisper output
        mock_whisper = MagicMock()
        mock_whisper.transcribe.return_value = (
            [
                MockSegment(
                    text=" Good morning team, let's review the sprint deliverables.",
                    start=0.0,
                    end=4.0,
                    words=[
                        MockWord(" Good", 0.0, 0.5),
                        MockWord(" morning", 0.5, 1.0),
                        MockWord(" team,", 1.0, 1.5),
                        MockWord(" let's", 1.6, 2.0),
                        MockWord(" review", 2.0, 2.7),
                        MockWord(" the", 2.7, 3.0),
                        MockWord(" sprint", 3.0, 3.5),
                        MockWord(" deliverables.", 3.5, 4.0),
                    ],
                )
            ],
            MagicMock(language="en", duration=4.0),
        )

        space_app.handler.whisper_model = mock_whisper

        # Mock diarization output
        mock_diarize = MagicMock()
        mock_diarize.itertracks.return_value = [
            (MagicMock(start=0.0, end=4.0), None, "SPEAKER_00")
        ]
        space_app.handler.diarization_pipeline = MagicMock(return_value=mock_diarize)

        # Create a dummy audio file
        dummy_file = Path("test_dummy.wav")
        dummy_file.write_bytes(b"RIFFdummybytes")

        try:
            result = space_app._run_transcription(
                audio_path=str(dummy_file),
                min_speakers=1,
                max_speakers=2,
                language="en",
            )

            self.assertIn("segments", result)
            self.assertIn("language", result)
            self.assertIn("duration", result)
            self.assertEqual(len(result["segments"]), 1)
            self.assertEqual(result["segments"][0]["speaker"], "Speaker 1")
            self.assertIn("Good morning team", result["segments"][0]["text"])

        finally:
            if dummy_file.exists():
                dummy_file.unlink()

    def test_hftoken_secret_detection_and_propagation(self):
        """Verifies that HFTOKEN in environment is detected and propagated to HF_TOKEN and HUGGING_FACE_HUB_TOKEN."""
        with patch.dict(os.environ, {"HFTOKEN": "hf_test_secret_token_abc"}, clear=True):
            handler = EndpointHandler(lazy=True)
            self.assertEqual(handler.hf_token, "hf_test_secret_token_abc")
            self.assertEqual(os.environ.get("HF_TOKEN"), "hf_test_secret_token_abc")
            self.assertEqual(os.environ.get("HUGGING_FACE_HUB_TOKEN"), "hf_test_secret_token_abc")

    def test_pyannote_explicit_token_passing(self):
        """Verifies that PyAnnote Pipeline.from_pretrained receives the token explicitly."""
        mock_pipeline_cls = MagicMock()
        mock_pipeline_cls.from_pretrained.return_value = MagicMock()
        
        import handler as handler_module
        orig_pipeline = handler_module.Pipeline
        handler_module.Pipeline = mock_pipeline_cls
        try:
            with patch.dict(os.environ, {"HFTOKEN": "hf_valid_token_xyz"}, clear=True):
                handler = EndpointHandler(lazy=True)
                handler._load_models()
                mock_pipeline_cls.from_pretrained.assert_called_with(
                    "pyannote/speaker-diarization-community-1",
                    token="hf_valid_token_xyz",
                )
        finally:
            handler_module.Pipeline = orig_pipeline


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