This comprehensive session management system tracks and automatically cleans up resources associated with user sessions, preventing resource exhaustion and disk space issues. Key features: - Automatic tracking of all session resources (audio files, temp files, streams) - Per-session resource limits (100 files max, 100MB storage max) - Automatic cleanup of idle sessions (15 minutes) and expired sessions (1 hour) - Background cleanup thread runs every minute - Real-time monitoring via admin endpoints - CLI commands for manual management - Integration with Flask request lifecycle Implementation details: - SessionManager class manages lifecycle of UserSession objects - Each session tracks resources with metadata (type, size, creation time) - Automatic resource eviction when limits are reached (LRU policy) - Orphaned file detection and cleanup - Thread-safe operations with proper locking - Comprehensive metrics and statistics export - Admin API endpoints for monitoring and control Security considerations: - Sessions tied to IP address and user agent - Admin endpoints require authentication - Secure file path handling - Resource limits prevent DoS attacks This addresses the critical issue of temporary file accumulation that could lead to disk exhaustion in production environments. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
264 lines
9.0 KiB
Python
264 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Unit tests for session management system
|
|
"""
|
|
import unittest
|
|
import tempfile
|
|
import shutil
|
|
import time
|
|
import os
|
|
from session_manager import SessionManager, UserSession, SessionResource
|
|
from flask import Flask, g, session
|
|
|
|
class TestSessionManager(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test fixtures"""
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.config = {
|
|
'max_session_duration': 3600,
|
|
'max_idle_time': 900,
|
|
'max_resources_per_session': 5, # Small limit for testing
|
|
'max_bytes_per_session': 1024 * 1024, # 1MB for testing
|
|
'cleanup_interval': 1, # 1 second for faster testing
|
|
'session_storage_path': self.temp_dir
|
|
}
|
|
self.manager = SessionManager(self.config)
|
|
|
|
def tearDown(self):
|
|
"""Clean up test fixtures"""
|
|
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
|
|
|
def test_create_session(self):
|
|
"""Test session creation"""
|
|
session = self.manager.create_session(
|
|
session_id='test-123',
|
|
user_id='user-1',
|
|
ip_address='127.0.0.1',
|
|
user_agent='Test Agent'
|
|
)
|
|
|
|
self.assertEqual(session.session_id, 'test-123')
|
|
self.assertEqual(session.user_id, 'user-1')
|
|
self.assertEqual(session.ip_address, '127.0.0.1')
|
|
self.assertEqual(session.user_agent, 'Test Agent')
|
|
self.assertEqual(len(session.resources), 0)
|
|
|
|
def test_get_session(self):
|
|
"""Test session retrieval"""
|
|
self.manager.create_session(session_id='test-456')
|
|
session = self.manager.get_session('test-456')
|
|
|
|
self.assertIsNotNone(session)
|
|
self.assertEqual(session.session_id, 'test-456')
|
|
|
|
# Non-existent session
|
|
session = self.manager.get_session('non-existent')
|
|
self.assertIsNone(session)
|
|
|
|
def test_add_resource(self):
|
|
"""Test adding resources to session"""
|
|
self.manager.create_session(session_id='test-789')
|
|
|
|
# Add a resource
|
|
resource = self.manager.add_resource(
|
|
session_id='test-789',
|
|
resource_type='audio_file',
|
|
resource_id='audio-1',
|
|
path='/tmp/test.wav',
|
|
size_bytes=1024,
|
|
metadata={'format': 'wav'}
|
|
)
|
|
|
|
self.assertIsNotNone(resource)
|
|
self.assertEqual(resource.resource_id, 'audio-1')
|
|
self.assertEqual(resource.resource_type, 'audio_file')
|
|
self.assertEqual(resource.size_bytes, 1024)
|
|
|
|
# Check session updated
|
|
session = self.manager.get_session('test-789')
|
|
self.assertEqual(len(session.resources), 1)
|
|
self.assertEqual(session.total_bytes_used, 1024)
|
|
|
|
def test_resource_limits(self):
|
|
"""Test resource limit enforcement"""
|
|
self.manager.create_session(session_id='test-limits')
|
|
|
|
# Add resources up to limit
|
|
for i in range(5):
|
|
self.manager.add_resource(
|
|
session_id='test-limits',
|
|
resource_type='temp_file',
|
|
resource_id=f'file-{i}',
|
|
size_bytes=100
|
|
)
|
|
|
|
session = self.manager.get_session('test-limits')
|
|
self.assertEqual(len(session.resources), 5)
|
|
|
|
# Add one more - should remove oldest
|
|
self.manager.add_resource(
|
|
session_id='test-limits',
|
|
resource_type='temp_file',
|
|
resource_id='file-new',
|
|
size_bytes=100
|
|
)
|
|
|
|
session = self.manager.get_session('test-limits')
|
|
self.assertEqual(len(session.resources), 5) # Still 5
|
|
self.assertNotIn('file-0', session.resources) # Oldest removed
|
|
self.assertIn('file-new', session.resources) # New one added
|
|
|
|
def test_size_limits(self):
|
|
"""Test size limit enforcement"""
|
|
self.manager.create_session(session_id='test-size')
|
|
|
|
# Add a large resource
|
|
self.manager.add_resource(
|
|
session_id='test-size',
|
|
resource_type='audio_file',
|
|
resource_id='large-1',
|
|
size_bytes=500 * 1024 # 500KB
|
|
)
|
|
|
|
# Add another large resource
|
|
self.manager.add_resource(
|
|
session_id='test-size',
|
|
resource_type='audio_file',
|
|
resource_id='large-2',
|
|
size_bytes=600 * 1024 # 600KB - would exceed 1MB limit
|
|
)
|
|
|
|
session = self.manager.get_session('test-size')
|
|
# First resource should be removed to make space
|
|
self.assertNotIn('large-1', session.resources)
|
|
self.assertIn('large-2', session.resources)
|
|
self.assertLessEqual(session.total_bytes_used, 1024 * 1024)
|
|
|
|
def test_remove_resource(self):
|
|
"""Test resource removal"""
|
|
self.manager.create_session(session_id='test-remove')
|
|
self.manager.add_resource(
|
|
session_id='test-remove',
|
|
resource_type='temp_file',
|
|
resource_id='to-remove',
|
|
size_bytes=1000
|
|
)
|
|
|
|
# Remove resource
|
|
success = self.manager.remove_resource('test-remove', 'to-remove')
|
|
self.assertTrue(success)
|
|
|
|
# Check it's gone
|
|
session = self.manager.get_session('test-remove')
|
|
self.assertEqual(len(session.resources), 0)
|
|
self.assertEqual(session.total_bytes_used, 0)
|
|
|
|
def test_cleanup_session(self):
|
|
"""Test session cleanup"""
|
|
# Create session with resources
|
|
self.manager.create_session(session_id='test-cleanup')
|
|
|
|
# Create actual temp file
|
|
temp_file = os.path.join(self.temp_dir, 'test-file.txt')
|
|
with open(temp_file, 'w') as f:
|
|
f.write('test content')
|
|
|
|
self.manager.add_resource(
|
|
session_id='test-cleanup',
|
|
resource_type='temp_file',
|
|
path=temp_file,
|
|
size_bytes=12
|
|
)
|
|
|
|
# Cleanup session
|
|
success = self.manager.cleanup_session('test-cleanup')
|
|
self.assertTrue(success)
|
|
|
|
# Check session is gone
|
|
session = self.manager.get_session('test-cleanup')
|
|
self.assertIsNone(session)
|
|
|
|
# Check file is deleted
|
|
self.assertFalse(os.path.exists(temp_file))
|
|
|
|
def test_session_info(self):
|
|
"""Test session info retrieval"""
|
|
self.manager.create_session(
|
|
session_id='test-info',
|
|
ip_address='192.168.1.1'
|
|
)
|
|
|
|
self.manager.add_resource(
|
|
session_id='test-info',
|
|
resource_type='audio_file',
|
|
size_bytes=2048
|
|
)
|
|
|
|
info = self.manager.get_session_info('test-info')
|
|
self.assertIsNotNone(info)
|
|
self.assertEqual(info['session_id'], 'test-info')
|
|
self.assertEqual(info['ip_address'], '192.168.1.1')
|
|
self.assertEqual(info['resource_count'], 1)
|
|
self.assertEqual(info['total_bytes_used'], 2048)
|
|
|
|
def test_stats(self):
|
|
"""Test statistics calculation"""
|
|
# Create multiple sessions
|
|
for i in range(3):
|
|
self.manager.create_session(session_id=f'test-stats-{i}')
|
|
self.manager.add_resource(
|
|
session_id=f'test-stats-{i}',
|
|
resource_type='temp_file',
|
|
size_bytes=1000
|
|
)
|
|
|
|
stats = self.manager.get_stats()
|
|
self.assertEqual(stats['active_sessions'], 3)
|
|
self.assertEqual(stats['active_resources'], 3)
|
|
self.assertEqual(stats['active_bytes'], 3000)
|
|
self.assertEqual(stats['total_sessions_created'], 3)
|
|
|
|
def test_metrics_export(self):
|
|
"""Test metrics export"""
|
|
self.manager.create_session(session_id='test-metrics')
|
|
metrics = self.manager.export_metrics()
|
|
|
|
self.assertIn('sessions', metrics)
|
|
self.assertIn('resources', metrics)
|
|
self.assertIn('limits', metrics)
|
|
self.assertEqual(metrics['sessions']['active'], 1)
|
|
|
|
class TestFlaskIntegration(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up Flask app for testing"""
|
|
self.app = Flask(__name__)
|
|
self.app.config['TESTING'] = True
|
|
self.app.config['SECRET_KEY'] = 'test-secret'
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.app.config['UPLOAD_FOLDER'] = self.temp_dir
|
|
|
|
# Initialize session manager
|
|
from session_manager import init_app
|
|
init_app(self.app)
|
|
|
|
self.client = self.app.test_client()
|
|
self.ctx = self.app.test_request_context()
|
|
self.ctx.push()
|
|
|
|
def tearDown(self):
|
|
"""Clean up"""
|
|
self.ctx.pop()
|
|
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
|
|
|
def test_before_request_handler(self):
|
|
"""Test Flask before_request integration"""
|
|
with self.client:
|
|
# Make a request
|
|
response = self.client.get('/')
|
|
|
|
# Session should be created
|
|
with self.client.session_transaction() as sess:
|
|
self.assertIn('session_id', sess)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |