Merge pull request #7 from codeSamuraii/new-test-client

New "live" integration test client + misc

Author: codeSamuraii

.github/workflows/unit-tests.yml

@@ -27,5 +27,8 @@ jobs:
           pip install -r requirements.txt
           pip install pytest fakeredis
 
+      - name: Redis Server in GitHub Actions
+        uses: supercharge/[email protected]
+
       - name: Run tests
-        run: pytest --disable-pytest-warnings -v tests/
+        run: pytest -s -v --tb=short tests/

.gitignore

@@ -14,3 +14,6 @@ Dockerfile
 fly.toml
 .pytest_cache/
 *.instructions.md
+CLAUDE.md
+poetry.lock
+pyproject.toml

lib/callbacks.py

@@ -16,13 +16,16 @@ async def _send_error_and_close(error: Exception | str) -> None:
     return _send_error_and_close
 
 
+class StreamTerminated(Exception):
+    """Raised to terminate a stream when an error occurs after the response has started."""
+    pass
+
 def raise_http_exception(request: Request) -> Callable[[Exception | str], Awaitable[None]]:
     """Callback to raise an HTTPException with a specific status code."""
 
     async def _raise_http_exception(error: Exception | str) -> None:
         message = str(error) if isinstance(error, Exception) else error
         code = error.status_code if isinstance(error, HTTPException) else 400
-        if not await request.is_disconnected():
-            raise HTTPException(status_code=code, detail=message)
+        raise StreamTerminated(f"{code}: {message}") from error
 
     return _raise_http_exception

lib/logging.py

@@ -1,6 +1,8 @@
 import logging
 
 
+# ---- FORMATTING ----
+
 class ColoredFormatter(logging.Formatter):
     """Custom formatter with colored output and specific formatting for transfers."""
 
@@ -59,12 +61,13 @@ def format(self, record: logging.LogRecord) -> str:
 
         return result + exc_text
 
-
 class HealthCheckFilter(logging.Filter):
     def filter(self, record):
         return '"GET /health HTTP/1.1" 200' not in record.getMessage()
 
 
+# ---- LOGGING SETUP ----
+
 def setup_logging() -> logging.Logger:
     """Configure all loggers to use our custom ColoredFormatter."""
     formatter = ColoredFormatter(
@@ -96,7 +99,6 @@ def setup_logging() -> logging.Logger:
 
     return root_logger
 
-
 def get_logger(logger_name: str) -> logging.Logger:
     console_handler = logging.StreamHandler()
     console_handler.setLevel(logging.DEBUG)
@@ -114,3 +116,38 @@ def get_logger(logger_name: str) -> logging.Logger:
     logger.propagate = False
 
     return logger
+
+
+# ---- PATCHING ----
+
+class HasLogging(type):
+    """Metaclass that automatically adds logging methods and a logger property."""
+
+    def __new__(mcs, name, bases, namespace, **kwargs):
+        name_from = kwargs.get('name_from', 'name')
+
+        @property
+        def logger(self):
+            if not hasattr(self, '_logger'):
+                class_name = self.__class__.__name__
+                fallback_name = class_name + str(id(self))[-4:]
+                logger_name = getattr(self, name_from, fallback_name)
+                self._logger = get_logger(logger_name)
+                if not hasattr(self, name_from):
+                    self._logger.warning(
+                        f"Object {class_name} does not have attribute '{name_from}', "
+                        f"using default name: {logger_name}"
+                    )
+            return self._logger
+
+        namespace['logger'] = logger
+
+        def make_log_method(level):
+            def log_method(self, msg, *args, **kwargs):
+                getattr(self.logger, level)(msg, *args, **kwargs)
+            return log_method
+
+        for level in {'debug', 'info', 'warning', 'error', 'exception', 'critical'}:
+            namespace[level] = make_log_method(level)
+
+        return super().__new__(mcs, name, bases, namespace)

lib/store.py

@@ -3,10 +3,10 @@
 import redis.asyncio as redis
 from typing import Optional, Annotated
 
-from lib.logging import get_logger
+from lib.logging import HasLogging, get_logger
 
 
-class Store:
+class Store(metaclass=HasLogging, name_from='transfer_id'):
     """
     Redis-based store for file transfer queues and events.
     Handles data queuing and event signaling for transfer coordination.
@@ -17,7 +17,6 @@ class Store:
     def __init__(self, transfer_id: str):
         self.transfer_id = transfer_id
         self.redis = self.get_redis()
-        self.log = get_logger(transfer_id)
 
         self._k_queue = self.key('queue')
         self._k_meta = self.key('metadata')
@@ -42,13 +41,13 @@ async def _wait_for_queue_space(self, maxsize: int) -> None:
         while await self.redis.llen(self._k_queue) >= maxsize:
             await asyncio.sleep(0.5)
 
-    async def put_in_queue(self, data: bytes, maxsize: int = 16, timeout: float = 10.0) -> None:
+    async def put_in_queue(self, data: bytes, maxsize: int = 16, timeout: float = 20.0) -> None:
         """Add data to the transfer queue with backpressure control."""
         async with asyncio.timeout(timeout):
             await self._wait_for_queue_space(maxsize)
         await self.redis.lpush(self._k_queue, data)
 
-    async def get_from_queue(self, timeout: float = 10.0) -> bytes:
+    async def get_from_queue(self, timeout: float = 20.0) -> bytes:
         """Get data from the transfer queue with timeout."""
         result = await self.redis.brpop([self._k_queue], timeout=timeout)
         if not result:
@@ -77,12 +76,12 @@ async def wait_for_event(self, event_name: str, timeout: float = 300.0) -> None:
         async def _poll_marker():
             while not await self.redis.exists(event_marker_key):
                 await asyncio.sleep(1)
-            self.log.debug(f">> POLL: Event '{event_name}' fired.")
+            self.debug(f">> POLL: Event '{event_name}' fired.")
 
         async def _listen_for_message():
             async for message in pubsub.listen():
                 if message and message['type'] == 'message':
-                    self.log.debug(f">> SUB : Received message for event '{event_name}'.")
+                    self.debug(f">> SUB : Received message for event '{event_name}'.")
                     return
 
         poll_marker = asyncio.wait_for(_poll_marker(), timeout=timeout)
@@ -98,7 +97,7 @@ async def _listen_for_message():
                 task.cancel()
 
         except asyncio.TimeoutError:
-            self.log.error(f"Timeout waiting for event '{event_name}' after {timeout} seconds.")
+            self.error(f"Timeout waiting for event '{event_name}' after {timeout} seconds.")
             for task in tasks:
                 task.cancel()
             raise
@@ -112,9 +111,12 @@ async def _listen_for_message():
 
     async def set_metadata(self, metadata: str) -> None:
         """Store transfer metadata."""
-        if int (await self.redis.exists(self._k_meta)) > 0:
-            raise KeyError(f"Metadata for transfer '{self.transfer_id}' already exists.")
-        await self.redis.set(self._k_meta, metadata, nx=True)
+        challenge = random.randbytes(8)
+        await self.redis.set(self._k_meta, challenge, nx=True)
+        if await self.redis.get(self._k_meta) == challenge:
+            await self.redis.set(self._k_meta, metadata, ex=300)
+        else:
+            raise KeyError("Metadata already set for this transfer.")
 
     async def get_metadata(self) -> str | None:
         """Retrieve transfer metadata."""
@@ -179,6 +181,6 @@ async def cleanup(self) -> int:
                 break
 
         if keys_to_delete:
-            self.log.debug(f"- Cleaning up {len(keys_to_delete)} keys")
+            self.debug(f"- Cleaning up {len(keys_to_delete)} keys")
             return await self.redis.delete(*keys_to_delete)
         return 0

lib/transfer.py

@@ -1,14 +1,29 @@
 import asyncio
 from starlette.responses import ClientDisconnect
 from starlette.websockets import WebSocketDisconnect
-from typing import AsyncIterator, Callable, Awaitable
+from typing import AsyncIterator, Callable, Awaitable, Optional, Any
 
 from lib.store import Store
-from lib.logging import get_logger
 from lib.metadata import FileMetadata
+from lib.logging import HasLogging, get_logger
+logger = get_logger('transfer')
 
 
-class FileTransfer:
+class TransferError(Exception):
+    """Custom exception for transfer errors with optional propagation control."""
+    def __init__(self, *args, propagate: bool = False, **extra: Any) -> None:
+        super().__init__(*args)
+        self.propagate = propagate
+        self.extra = extra
+
+    @property
+    def shutdown(self) -> bool:
+        """Indicates if the transfer should be shut down (usually the opposite of `propagate`)."""
+        return self.extra.get('shutdown', not self.propagate)
+
+
+class FileTransfer(metaclass=HasLogging, name_from='uid'):
+    """Handles file transfers, including metadata queries and data streaming."""
 
     DONE_FLAG = b'\x00\xFF'
     DEAD_FLAG = b'\xDE\xAD'
@@ -20,9 +35,6 @@ def __init__(self, uid: str, file: FileMetadata):
         self.bytes_uploaded = 0
         self.bytes_downloaded = 0
 
-        log = get_logger(self.uid)
-        self.debug, self.info, self.warning, self.error = log.debug, log.info, log.warning, log.error
-
     @classmethod
     async def create(cls, uid: str, file: FileMetadata):
         transfer = cls(uid, file)
@@ -86,27 +98,33 @@ async def collect_upload(self, stream: AsyncIterator[bytes], on_error: Callable[
                     break
 
                 if await self.is_interrupted():
-                    raise ClientDisconnect("Transfer was interrupted by the receiver.")
+                    raise TransferError("Transfer was interrupted by the receiver.", propagate=False)
 
                 await self.store.put_in_queue(chunk)
                 self.bytes_uploaded += len(chunk)
 
             if self.bytes_uploaded < self.file.size:
-                raise ClientDisconnect("Received less data than expected.")
+                raise TransferError("Received less data than expected.", propagate=True)
 
             self.debug(f"△ End of upload, sending done marker.")
             await self.store.put_in_queue(self.DONE_FLAG)
 
         except (ClientDisconnect, WebSocketDisconnect) as e:
-            self.warning(f"△ Upload error: {str(e)}")
+            self.error(f"△ Unexpected upload error: {e}")
             await self.store.put_in_queue(self.DEAD_FLAG)
-            await on_error(e)
 
         except asyncio.TimeoutError as e:
             self.warning(f"△ Timeout during upload.")
             await on_error("Timeout during upload.")
 
-        else:
+        except TransferError as e:
+            self.warning(f"△ Upload error: {e}")
+            if e.propagate:
+                await self.store.put_in_queue(self.DEAD_FLAG)
+            else:
+                await on_error(e)
+
+        finally:
             await asyncio.sleep(1.0)
 
     async def supply_download(self, on_error: Callable[[Exception | str], Awaitable[None]]) -> AsyncIterator[bytes]:
@@ -117,10 +135,10 @@ async def supply_download(self, on_error: Callable[[Exception | str], Awaitable[
                 chunk = await self.store.get_from_queue()
 
                 if chunk == self.DEAD_FLAG:
-                    raise ClientDisconnect("Sender disconnected.")
+                    raise TransferError("Sender disconnected.")
 
                 if chunk == self.DONE_FLAG and self.bytes_downloaded < self.file.size:
-                    raise ClientDisconnect("Received less data than expected.")
+                    raise TransferError("Received less data than expected.")
 
                 elif chunk == self.DONE_FLAG:
                     self.debug(f"▼ Done marker received, ending download.")
@@ -129,16 +147,14 @@ async def supply_download(self, on_error: Callable[[Exception | str], Awaitable[
                 self.bytes_downloaded += len(chunk)
                 yield chunk
 
-        except (ClientDisconnect, WebSocketDisconnect) as e:
-            self.warning(f"▼ Download error: {e}")
-            await self.set_interrupted()
-
-        except asyncio.TimeoutError:
-            self.warning(f"▼ Timeout during download.")
-            await on_error("Timeout during download.")
+        except Exception as e:
+            self.error(f"▼ Unexpected download error!", exc_info=True)
+            self.debug("Debug info:", stack_info=True)
+            await on_error(e)
 
-        else:
-            await asyncio.sleep(1.0)
+        except TransferError as e:
+            self.warning(f"▼ Download error")
+            await on_error(e)
 
     async def cleanup(self):
         try:
@@ -148,10 +164,15 @@ async def cleanup(self):
             pass
 
     async def finalize_download(self):
-        self.debug("▼ Finalizing download...")
+        # self.debug("▼ Finalizing download...")
+        if self.bytes_downloaded < self.file.size and not await self.is_interrupted():
+            self.warning("▼ Client disconnected before download was complete.")
+            await self.set_interrupted()
+
+        await self.cleanup()
+        # self.debug("▼ Finalizing download...")
         if self.bytes_downloaded < self.file.size and not await self.is_interrupted():
             self.warning("▼ Client disconnected before download was complete.")
             await self.set_interrupted()
 
-        await asyncio.sleep(4.0)
         await self.cleanup()

static/index.html

@@ -69,7 +69,7 @@ <h3>Using cURL</h3>
                     <p><small>The <code>-JLO</code> flags downloads the file with its original name and follows redirects.</small></p>
                     <div class="code-block">
                         <code><span class="code-comment"># Example</span>
-curl -T <span class="code-variable">/music/song.mp3</span> https://transit.sh/<span class="code-string">music-for-dad</span>/
+curl -T <span class="code-variable">/music/song.mp3</span> https://transit.sh/<span class="code-string">music-for-dad</span>/ --expect100-timeout 300
 curl -JLO https://transit.sh/<span class="code-string">music-for-dad</span>/</code>
                     </div>
                 </div>