File size: 17,821 Bytes
76f9cd2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
"""
Unit tests for storage configuration system
Tests the new storage configuration functionality including:
- Storage config management
- Environment detection
- Path generation
- Storage tools
"""

import pytest
import os
import tempfile
import shutil
from pathlib import Path
from unittest.mock import patch, MagicMock
import asyncio

import sys
sys.path.append(str(Path(__file__).parent.parent / "src"))

from src.utils.storage_config import StorageConfig, get_storage_config
from src.tools.storage_tools import get_storage_info_tool


class TestStorageConfig:
    """Test cases for StorageConfig class"""
    
    def setup_method(self):
        """Setup for each test method"""
        # Create temporary directory for testing
        self.temp_dir = Path(tempfile.mkdtemp())
        self.config_file = self.temp_dir / "test_config.env"
        
        # Create test config file
        config_content = """
DOWNLOADS_DIR=./test_downloads
TRANSCRIPTS_DIR=./test_transcripts
CACHE_DIR=./test_cache
DEFAULT_MODEL_SIZE=base
DEFAULT_OUTPUT_FORMAT=srt
USE_PARALLEL_PROCESSING=true
CHUNK_DURATION=30
"""
        with open(self.config_file, 'w') as f:
            f.write(config_content)
    
    def teardown_method(self):
        """Cleanup after each test method"""
        # Remove temporary directory
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
    
    def test_local_environment_detection(self):
        """Test local environment detection and configuration loading"""
        # Mock environment to ensure local detection
        with patch.dict(os.environ, {}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            assert not storage_config.is_modal_env
            assert storage_config.default_model_size == "base"
            assert storage_config.default_output_format == "srt"
            assert storage_config.use_parallel_processing == True
            assert storage_config.chunk_duration == 30
    
    def test_modal_environment_detection(self):
        """Test Modal environment detection"""
        # Mock Modal environment variables
        with patch.dict(os.environ, {"MODAL_TASK_ID": "test-task-123"}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            assert storage_config.is_modal_env
            assert str(storage_config.downloads_dir) == "/root/downloads"
            assert str(storage_config.transcripts_dir) == "/root/transcripts"
            assert str(storage_config.cache_dir) == "/root/cache"
    
    def test_modal_environment_detection_deployment_mode(self):
        """Test Modal environment detection via DEPLOYMENT_MODE"""
        with patch.dict(os.environ, {"DEPLOYMENT_MODE": "modal"}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            assert storage_config.is_modal_env
    
    def test_modal_environment_detection_container_var(self):
        """Test Modal environment detection via MODAL_IS_INSIDE_CONTAINER"""
        with patch.dict(os.environ, {"MODAL_IS_INSIDE_CONTAINER": "true"}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            assert storage_config.is_modal_env
    
    def test_path_generation(self):
        """Test path generation methods"""
        with patch.dict(os.environ, {}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            # Test download path
            download_path = storage_config.get_download_path("test.mp3")
            assert download_path.name == "test.mp3"
            assert "test_downloads" in str(download_path)
            
            # Test transcript paths
            txt_path = storage_config.get_transcript_path("test.mp3", "txt")
            assert txt_path.name == "test.txt"
            assert "test_transcripts" in str(txt_path)
            
            srt_path = storage_config.get_transcript_path("test.mp3", "srt")
            assert srt_path.name == "test.srt"
            
            # Test default format
            default_path = storage_config.get_transcript_path("test.mp3")
            assert default_path.name == "test.srt"  # Should use default format
            
            # Test cache path
            cache_path = storage_config.get_cache_path("temp.dat")
            assert cache_path.name == "temp.dat"
            assert "test_cache" in str(cache_path)
    
    def test_audio_files_listing(self):
        """Test audio files listing functionality"""
        with patch.dict(os.environ, {}, clear=True):
            # Create a separate test directory for this specific test
            test_dir = self.temp_dir / "audio_test"
            test_config_file = test_dir / "config.env"
            test_dir.mkdir(exist_ok=True)
            
            # Create isolated config file
            config_content = """
DOWNLOADS_DIR=./audio_test_downloads
TRANSCRIPTS_DIR=./audio_test_transcripts
CACHE_DIR=./audio_test_cache
"""
            with open(test_config_file, 'w') as f:
                f.write(config_content)
            
            storage_config = StorageConfig(config_file=str(test_config_file))
            
            # Create test audio files
            storage_config.downloads_dir.mkdir(parents=True, exist_ok=True)
            
            test_files = ["test1.mp3", "test2.wav", "test3.m4a", "not_audio.txt"]
            for filename in test_files:
                (storage_config.downloads_dir / filename).touch()
            
            audio_files = storage_config.get_audio_files()
            audio_names = [f.name for f in audio_files]
            
            assert "test1.mp3" in audio_names
            assert "test2.wav" in audio_names
            assert "test3.m4a" in audio_names
            assert "not_audio.txt" not in audio_names
            assert len(audio_files) == 3
    
    def test_transcript_files_mapping(self):
        """Test transcript files mapping functionality"""
        with patch.dict(os.environ, {}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            # Test specific audio file mapping
            transcript_files = storage_config.get_transcript_files("episode123.mp3")
            
            assert "txt" in transcript_files
            assert "srt" in transcript_files
            assert "json" in transcript_files
            
            assert transcript_files["txt"].name == "episode123.txt"
            assert transcript_files["srt"].name == "episode123.srt"
            assert transcript_files["json"].name == "episode123.json"
    
    def test_storage_info_generation(self):
        """Test storage information generation"""
        with patch.dict(os.environ, {}, clear=True):
            # Create a separate test directory for this specific test
            test_dir = self.temp_dir / "info_test"
            test_config_file = test_dir / "config.env"
            test_dir.mkdir(exist_ok=True)
            
            # Create isolated config file
            config_content = """
DOWNLOADS_DIR=./info_test_downloads
TRANSCRIPTS_DIR=./info_test_transcripts
CACHE_DIR=./info_test_cache
"""
            with open(test_config_file, 'w') as f:
                f.write(config_content)
            
            storage_config = StorageConfig(config_file=str(test_config_file))
            
            # Create some test files
            storage_config.downloads_dir.mkdir(parents=True, exist_ok=True)
            storage_config.transcripts_dir.mkdir(parents=True, exist_ok=True)
            
            # Create test audio file
            test_audio = storage_config.downloads_dir / "test.mp3"
            test_audio.write_bytes(b"fake audio data" * 100)
            
            # Create test transcript files
            (storage_config.transcripts_dir / "test.txt").write_text("transcript text")
            (storage_config.transcripts_dir / "test.srt").write_text("srt content")
            
            storage_info = storage_config.get_storage_info()
            
            assert storage_info["environment"] == "local"
            assert storage_info["audio_files_count"] == 1
            assert storage_info["transcript_txt_count"] == 1
            assert storage_info["transcript_srt_count"] == 1
            assert storage_info["transcript_json_count"] == 0
            # Check that sizes are calculated (should be greater than 0 due to our test files)
            assert storage_info["downloads_size_mb"] >= 0
            assert storage_info["transcripts_size_mb"] >= 0
    
    def test_cleanup_temp_files(self):
        """Test temporary files cleanup"""
        with patch.dict(os.environ, {}, clear=True):
            storage_config = StorageConfig(config_file=str(self.config_file))
            
            # Create cache directory and temp files
            storage_config.cache_dir.mkdir(parents=True, exist_ok=True)
            
            temp_file1 = storage_config.cache_dir / "temp_file1.dat"
            temp_file2 = storage_config.cache_dir / "temp_file2.dat"
            normal_file = storage_config.cache_dir / "normal_file.dat"
            
            temp_file1.touch()
            temp_file2.touch()
            normal_file.touch()
            
            # Test cleanup
            storage_config.cleanup_temp_files("temp_*")
            
            assert not temp_file1.exists()
            assert not temp_file2.exists()
            assert normal_file.exists()  # Should not be deleted
    
    def test_config_file_not_exists(self):
        """Test behavior when config file doesn't exist"""
        non_existent_config = self.temp_dir / "non_existent.env"
        
        with patch.dict(os.environ, {}, clear=True):
            storage_config = StorageConfig(config_file=str(non_existent_config))
            
            # Should use defaults
            assert not storage_config.is_modal_env
            assert storage_config.default_model_size == "turbo"
            assert storage_config.default_output_format == "srt"


class TestStorageConfigGlobalInstance:
    """Test cases for global storage config instance management"""
    
    def test_global_instance_singleton(self):
        """Test that get_storage_config returns singleton instance"""
        # Clear any existing global instance
        import src.utils.storage_config as storage_module
        storage_module._storage_config = None
        
        with patch.dict(os.environ, {}, clear=True):
            config1 = get_storage_config()
            config2 = get_storage_config()
            
            assert config1 is config2  # Should be the same instance
    
    def test_global_instance_reset(self):
        """Test resetting global instance"""
        import src.utils.storage_config as storage_module
        
        with patch.dict(os.environ, {}, clear=True):
            config1 = get_storage_config()
            
            # Reset global instance
            storage_module._storage_config = None
            
            config2 = get_storage_config()
            
            assert config1 is not config2  # Should be different instances


class TestStorageTools:
    """Test cases for storage management tools"""
    
    def setup_method(self):
        """Setup for each test method"""
        self.temp_dir = Path(tempfile.mkdtemp())
        
        # Mock storage config to use temp directory
        self.mock_config = MagicMock()
        self.mock_config.downloads_dir = self.temp_dir / "downloads"
        self.mock_config.transcripts_dir = self.temp_dir / "transcripts"
        self.mock_config.cache_dir = self.temp_dir / "cache"
        self.mock_config.is_modal_env = False
        
        # Create directories
        for directory in [self.mock_config.downloads_dir, 
                         self.mock_config.transcripts_dir, 
                         self.mock_config.cache_dir]:
            directory.mkdir(parents=True, exist_ok=True)
    
    def teardown_method(self):
        """Cleanup after each test method"""
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
    
    @pytest.mark.asyncio
    async def test_get_storage_info_tool_success(self):
        """Test get_storage_info_tool with successful execution"""
        
        # Create test files
        (self.mock_config.downloads_dir / "test.mp3").write_bytes(b"audio data")
        (self.mock_config.transcripts_dir / "test.txt").write_text("transcript")
        
        # Mock storage config
        mock_storage_info = {
            "environment": "local",
            "downloads_dir": str(self.mock_config.downloads_dir),
            "transcripts_dir": str(self.mock_config.transcripts_dir),
            "cache_dir": str(self.mock_config.cache_dir),
            "audio_files_count": 1,
            "transcript_txt_count": 1,
            "transcript_srt_count": 0,
            "transcript_json_count": 0,
            "downloads_size_mb": 0.01,
            "transcripts_size_mb": 0.01,
            "cache_size_mb": 0.0
        }
        
        self.mock_config.get_storage_info.return_value = mock_storage_info
        
        with patch('src.tools.storage_tools.get_storage_config', return_value=self.mock_config):
            result = await get_storage_info_tool()
            
            assert result["status"] == "success"
            assert result["environment"] == "local"
            assert result["audio_files_count"] == 1
            assert result["transcript_txt_count"] == 1
    
    @pytest.mark.asyncio
    async def test_get_storage_info_tool_failure(self):
        """Test get_storage_info_tool with exception handling"""
        
        # Mock config that raises exception
        self.mock_config.get_storage_info.side_effect = Exception("Test error")
        
        with patch('src.tools.storage_tools.get_storage_config', return_value=self.mock_config):
            result = await get_storage_info_tool()
            
            assert result["status"] == "failed"
            assert "Test error" in result["error_message"]


class TestDistributedTranscriptionFixes:
    """Test cases for distributed transcription speaker information collection fixes"""
    
    def test_collect_speaker_information_string_speakers(self):
        """Test handling of string format speakers_detected"""
        from src.services.distributed_transcription_service import DistributedTranscriptionService
        
        service = DistributedTranscriptionService()
        
        # Test with string format speakers_detected
        successful_chunks = [
            {
                "speakers_detected": "SPEAKER_01",  # String instead of list
                "speaker_summary": {
                    "SPEAKER_01": {
                        "total_duration": 120.5,
                        "segment_count": 5
                    }
                }
            },
            {
                "speakers_detected": ["SPEAKER_02"],  # Normal list format
                "speaker_summary": {
                    "SPEAKER_02": {
                        "total_duration": 95.3,
                        "segment_count": 3
                    }
                }
            }
        ]
        
        result = service._collect_speaker_information(successful_chunks, True)
        
        assert result["global_speaker_count"] == 2
        assert "SPEAKER_01" in result["speakers_detected"]
        assert "SPEAKER_02" in result["speakers_detected"]
        assert result["speaker_summary"]["SPEAKER_01"]["total_duration"] == 120.5
        assert result["speaker_summary"]["SPEAKER_02"]["total_duration"] == 95.3
    
    def test_collect_speaker_information_invalid_data(self):
        """Test handling of invalid speaker data"""
        from src.services.distributed_transcription_service import DistributedTranscriptionService
        
        service = DistributedTranscriptionService()
        
        # Test with invalid data formats
        successful_chunks = [
            {
                "speakers_detected": 123,  # Invalid type (number)
                "speaker_summary": "invalid"  # Invalid type (string)
            },
            {
                "speakers_detected": None,  # None value
                "speaker_summary": {
                    "SPEAKER_01": "invalid_info"  # Invalid speaker info format
                }
            },
            {
                "speakers_detected": ["SPEAKER_02"],  # Valid
                "speaker_summary": {
                    "SPEAKER_02": {
                        "total_duration": 50.0,
                        "segment_count": 2
                    }
                }
            }
        ]
        
        result = service._collect_speaker_information(successful_chunks, True)
        
        # Should handle invalid data gracefully and only process valid chunk
        assert result["global_speaker_count"] == 1
        assert result["speakers_detected"] == ["SPEAKER_02"]
        assert result["speaker_summary"]["SPEAKER_02"]["total_duration"] == 50.0
    
    def test_collect_speaker_information_disabled(self):
        """Test when speaker diarization is disabled"""
        from src.services.distributed_transcription_service import DistributedTranscriptionService
        
        service = DistributedTranscriptionService()
        
        successful_chunks = [{"speakers_detected": ["SPEAKER_01"]}]
        
        result = service._collect_speaker_information(successful_chunks, False)
        
        # Should return empty result when disabled
        assert result == {}


if __name__ == "__main__":
    # Run tests with pytest
    pytest.main([__file__, "-v"])