add tests

Author: Josephrp

modules/jax/testcontainers/whisper-diarization/__init__.py

@@ -2,8 +2,8 @@
 import tempfile
 from typing import Optional
 
-from testcontainers.core.container import DockerContainer
-from testcontainers.core.waiting_utils import wait_container_is_ready
+from core.testcontainers.core.container import DockerContainer
+from core.testcontainers.core.waiting_utils import wait_container_is_ready
 from urllib.error import URLError
 
 class JAXWhisperDiarizationContainer(DockerContainer):

modules/jax/tests/test_jax.py

@@ -0,0 +1,28 @@
+import pytest
+from testcontainers.jax import JAXContainer
+
+def test_jax_container():
+    with JAXContainer() as jax_container:
+        jax_container.connect()
+        
+        # Test running a simple JAX computation
+        result = jax_container.run_jax_command("import jax; print(jax.numpy.add(1, 1))")
+        assert "2" in result.output.decode()
+
+def test_jax_container_gpu_support():
+    with JAXContainer() as jax_container:
+        jax_container.connect()
+        
+        # Test GPU availability
+        result = jax_container.run_jax_command(
+            "import jax; print(jax.devices())"
+        )
+        assert "gpu" in result.output.decode().lower()
+
+def test_jax_container_jupyter():
+    with JAXContainer() as jax_container:
+        jax_container.connect()
+        
+        jupyter_url = jax_container.get_jupyter_url()
+        assert jupyter_url.startswith("http://")
+        assert ":8888" in jupyter_url

modules/jax/tests/test_whisper.py

@@ -0,0 +1,30 @@
+import pytest
+from testcontainers.whisper_jax import WhisperJAXContainer
+
+@pytest.mark.parametrize("model_name", ["openai/whisper-tiny", "openai/whisper-base"])
+def test_whisper_jax_container(model_name):
+    with WhisperJAXContainer(model_name) as whisper:
+        whisper.connect()
+        
+        # Test file transcription
+        result = whisper.transcribe_file("/path/to/test/audio.wav")
+        assert isinstance(result, dict)
+        assert 'text' in result
+        assert isinstance(result['text'], str)
+        
+        # Test YouTube transcription
+        result = whisper.transcribe_youtube("https://www.youtube.com/watch?v=dQw4w9WgXcQ")
+        assert isinstance(result, dict)
+        assert 'text' in result
+        assert isinstance(result['text'], str)
+
+def test_whisper_jax_container_with_timestamps():
+    with WhisperJAXContainer() as whisper:
+        whisper.connect()
+        
+        result = whisper.transcribe_file("/path/to/test/audio.wav", return_timestamps=True)
+        assert isinstance(result, dict)
+        assert 'text' in result
+        assert 'chunks' in result
+        assert isinstance(result['chunks'], list)
+        assert all('timestamp' in chunk for chunk in result['chunks'])

modules/jax/tests/test_whisper_diarization.py

@@ -0,0 +1,31 @@
+import pytest
+from testcontainers.jax_whisper_diarization import JAXWhisperDiarizationContainer
+
+@pytest.fixture(scope="module")
+def hf_token():
+    return "your_huggingface_token_here"  # Replace with a valid token or use an environment variable
+
+def test_jax_whisper_diarization_container(hf_token):
+    with JAXWhisperDiarizationContainer(hf_token=hf_token) as whisper_diarization:
+        whisper_diarization.connect()
+        
+        # Test file transcription and diarization
+        result = whisper_diarization.transcribe_and_diarize_file("/path/to/test/audio.wav")
+        assert isinstance(result, list)
+        assert all(isinstance(item, dict) for item in result)
+        assert all('speaker' in item and 'text' in item and 'timestamp' in item for item in result)
+        
+        # Test YouTube transcription and diarization
+        result = whisper_diarization.transcribe_and_diarize_youtube("https://www.youtube.com/watch?v=dQw4w9WgXcQ")
+        assert isinstance(result, list)
+        assert all(isinstance(item, dict) for item in result)
+        assert all('speaker' in item and 'text' in item and 'timestamp' in item for item in result)
+
+def test_jax_whisper_diarization_container_without_grouping(hf_token):
+    with JAXWhisperDiarizationContainer(hf_token=hf_token) as whisper_diarization:
+        whisper_diarization.connect()
+        
+        result = whisper_diarization.transcribe_and_diarize_file("/path/to/test/audio.wav", group_by_speaker=False)
+        assert isinstance(result, list)
+        assert all(isinstance(item, dict) for item in result)
+        assert all('speaker' in item and 'text' in item and 'timestamp' in item for item in result)