New version

mediaflow_proxy/speedtest/models.py

@@ -1,8 +1,7 @@
-from datetime import datetime
 from enum import Enum
 from typing import Dict, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, HttpUrl
 
 
 class SpeedTestProvider(str, Enum):
@@ -21,26 +20,20 @@ class UserInfo(BaseModel):
     country: Optional[str] = None
 
 
-class SpeedTestResult(BaseModel):
-    speed_mbps: float = Field(..., description="Speed in Mbps")
-    duration: float = Field(..., description="Test duration in seconds")
-    data_transferred: int = Field(..., description="Data transferred in bytes")
-    timestamp: datetime = Field(default_factory=datetime.utcnow)
+class MediaFlowServer(BaseModel):
+    url: HttpUrl
+    api_password: Optional[str] = None
+    name: Optional[str] = None
 
 
-class LocationResult(BaseModel):
-    result: Optional[SpeedTestResult] = None
-    error: Optional[str] = None
-    server_name: str
-    server_url: str
-
-
-class SpeedTestTask(BaseModel):
-    task_id: str
+class BrowserSpeedTestConfig(BaseModel):
     provider: SpeedTestProvider
-    results: Dict[str, LocationResult] = {}
-    started_at: datetime
-    completed_at: Optional[datetime] = None
-    status: str = "running"
+    test_urls: Dict[str, str]
+    test_duration: int = 10
     user_info: Optional[UserInfo] = None
-    current_location: Optional[str] = None
+
+
+class BrowserSpeedTestRequest(BaseModel):
+    provider: SpeedTestProvider
+    api_key: Optional[str] = None
+    current_api_password: Optional[str] = None

mediaflow_proxy/speedtest/service.py

@@ -1,20 +1,13 @@
-import logging
-import time
-from datetime import datetime, timezone
 from typing import Dict, Optional, Type
 
-from mediaflow_proxy.utils.cache_utils import get_cached_speedtest, set_cache_speedtest
-from mediaflow_proxy.utils.http_utils import Streamer, create_httpx_client
-from .models import SpeedTestTask, LocationResult, SpeedTestResult, SpeedTestProvider
+from .models import SpeedTestProvider
 from .providers.all_debrid import AllDebridSpeedTest
 from .providers.base import BaseSpeedTestProvider
 from .providers.real_debrid import RealDebridSpeedTest
 
-logger = logging.getLogger(__name__)
-
 
 class SpeedTestService:
-    """Service for managing speed tests across different providers."""
+    """Service for managing speed test provider configurations."""
 
     def __init__(self):
         # Provider mapping
@@ -23,7 +16,7 @@ class SpeedTestService:
             SpeedTestProvider.ALL_DEBRID: AllDebridSpeedTest,
         }
 
-    def _get_provider(self, provider: SpeedTestProvider, api_key: Optional[str] = None) -> BaseSpeedTestProvider:
+    def get_provider(self, provider: SpeedTestProvider, api_key: Optional[str] = None) -> BaseSpeedTestProvider:
         """Get the appropriate provider implementation."""
         provider_class = self._providers.get(provider)
         if not provider_class:
@@ -33,97 +26,3 @@ class SpeedTestService:
             raise ValueError("API key required for AllDebrid")
 
         return provider_class(api_key) if provider == SpeedTestProvider.ALL_DEBRID else provider_class()
-
-    async def create_test(
-        self, task_id: str, provider: SpeedTestProvider, api_key: Optional[str] = None
-    ) -> SpeedTestTask:
-        """Create a new speed test task."""
-        provider_impl = self._get_provider(provider, api_key)
-
-        # Get initial URLs and user info
-        urls, user_info = await provider_impl.get_test_urls()
-
-        task = SpeedTestTask(
-            task_id=task_id, provider=provider, started_at=datetime.now(tz=timezone.utc), user_info=user_info
-        )
-
-        await set_cache_speedtest(task_id, task)
-        return task
-
-    @staticmethod
-    async def get_test_results(task_id: str) -> Optional[SpeedTestTask]:
-        """Get results for a specific task."""
-        return await get_cached_speedtest(task_id)
-
-    async def run_speedtest(self, task_id: str, provider: SpeedTestProvider, api_key: Optional[str] = None):
-        """Run the speed test with real-time updates."""
-        try:
-            task = await get_cached_speedtest(task_id)
-            if not task:
-                raise ValueError(f"Task {task_id} not found")
-
-            provider_impl = self._get_provider(provider, api_key)
-            config = await provider_impl.get_config()
-
-            async with create_httpx_client() as client:
-                streamer = Streamer(client)
-
-                for location, url in config.test_urls.items():
-                    try:
-                        task.current_location = location
-                        await set_cache_speedtest(task_id, task)
-                        result = await self._test_location(location, url, streamer, config.test_duration, provider_impl)
-                        task.results[location] = result
-                        await set_cache_speedtest(task_id, task)
-                    except Exception as e:
-                        logger.error(f"Error testing {location}: {str(e)}")
-                        task.results[location] = LocationResult(
-                            error=str(e), server_name=location, server_url=config.test_urls[location]
-                        )
-                        await set_cache_speedtest(task_id, task)
-
-            # Mark task as completed
-            task.completed_at = datetime.now(tz=timezone.utc)
-            task.status = "completed"
-            task.current_location = None
-            await set_cache_speedtest(task_id, task)
-
-        except Exception as e:
-            logger.error(f"Error in speed test task {task_id}: {str(e)}")
-            if task := await get_cached_speedtest(task_id):
-                task.status = "failed"
-                await set_cache_speedtest(task_id, task)
-
-    async def _test_location(
-        self, location: str, url: str, streamer: Streamer, test_duration: int, provider: BaseSpeedTestProvider
-    ) -> LocationResult:
-        """Test speed for a specific location."""
-        try:
-            start_time = time.time()
-            total_bytes = 0
-
-            await streamer.create_streaming_response(url, headers={})
-
-            async for chunk in streamer.stream_content():
-                if time.time() - start_time >= test_duration:
-                    break
-                total_bytes += len(chunk)
-
-            duration = time.time() - start_time
-            speed_mbps = (total_bytes * 8) / (duration * 1_000_000)
-
-            # Get server info if available (for AllDebrid)
-            server_info = getattr(provider, "servers", {}).get(location)
-            server_url = server_info.url if server_info else url
-
-            return LocationResult(
-                result=SpeedTestResult(
-                    speed_mbps=round(speed_mbps, 2), duration=round(duration, 2), data_transferred=total_bytes
-                ),
-                server_name=location,
-                server_url=server_url,
-            )
-
-        except Exception as e:
-            logger.error(f"Error testing {location}: {str(e)}")
-            raise  # Re-raise to be handled by run_speedtest