test: Enhance testing capabilities with new pytest configurations and fixtures

edenreich · edenreich · commit 298b6f88ade7 · 2025-02-03T02:31:35.000Z
Signed-off-by: Eden Reich &lt;eden.reich@gmail.com&gt;
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
@@ -6,7 +6,9 @@ ENV ZSH_CUSTOM=/home/vscode/.oh-my-zsh/custom \
     PYLINT_VERSION=3.3.3 \
     BUILD_VERSION=1.2.2.post1 \
     TWINE_VERSION=6.0.1 \
-    TASK_VERSION=v3.41.0
+    TASK_VERSION=v3.41.0 \
+    PYTEST_VERSION=8.3.4 \
+    PYTEST_WATCH_VERSION=4.2.0
 
 RUN apt-get update && \
     # Install nodejs and npm
@@ -16,12 +18,17 @@ RUN apt-get update && \
     curl -s https://taskfile.dev/install.sh | sh -s -- -b /usr/local/bin ${TASK_VERSION} && \
     # Install pipx
     python -m pip install --upgrade pip && \
-    # Install black, isort, pylint using pip
-    pip install black==${BLACK_VERSION} && \
-    pip install isort==${ISORT_VERSION} && \
-    pip install pylint==${PYLINT_VERSION} && \
-    pip install build==${BUILD_VERSION} && \
-    pip install twine==${TWINE_VERSION} && \
+    # Install development tools using pip
+    pip install black==${BLACK_VERSION} \
+                isort==${ISORT_VERSION} \
+                pylint==${PYLINT_VERSION} \
+                build==${BUILD_VERSION} \
+                twine==${TWINE_VERSION} \
+                pytest==${PYTEST_VERSION} \
+                pytest-watch==${PYTEST_WATCH_VERSION} \
+                pytest-cov \
+                pytest-xdist \
+                debugpy && \
     # Clean up
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -34,7 +34,13 @@
         "dev.containers.copyGitConfig": true,
         "githubPullRequests.experimental.chat": true,
         "githubPullRequests.experimental.notificationsView": true,
-        "files.insertFinalNewline": true
+        "files.insertFinalNewline": true,
+        "python.testing.pytestEnabled": true,
+        "python.testing.unittestEnabled": false,
+        "python.testing.nosetestsEnabled": false,
+        "python.testing.pytestArgs": [
+            "tests"
+        ]
       }
     }
   },
diff --git a/.devcontainer/launch.json b/.devcontainer/launch.json
@@ -2,11 +2,30 @@
   "version": "0.2.0",
   "configurations": [
     {
-      "name": "Python: Current File",
+      "name": "Python: Debug Tests",
       "type": "debugpy",
       "request": "launch",
-      "program": "${file}",
-      "console": "integratedTerminal"
+      "program": "/usr/local/bin/python",
+      "args": [
+        "-v",
+        "--no-cov",
+        "tests/"
+      ],
+      "console": "integratedTerminal",
+      "justMyCode": false
+    },
+    {
+      "name": "Python: Debug Current Test",
+      "type": "debugpy", 
+      "request": "launch",
+      "program": "/usr/local/bin/python",
+      "args": [
+        "-v",
+        "--no-cov",
+        "${file}"
+      ],
+      "console": "integratedTerminal",
+      "justMyCode": false
     }
   ]
 }
diff --git a/Taskfile.yml b/Taskfile.yml
@@ -15,7 +15,22 @@ tasks:
   test:
     desc: Run tests
     cmds:
-      - pytest tests/
+      - pytest tests/ -v
+
+  test:watch:
+    desc: Run tests in watch mode
+    cmds:
+      - ptw tests/ -- -v
+
+  test:coverage:
+    desc: Run tests with coverage report
+    cmds:
+      - pytest tests/ -v --cov=inference_gateway --cov-report=term-missing
+
+  test:debug:
+    desc: Run tests with debugger enabled
+    cmds:
+      - pytest tests/ -v --pdb
 
   clean:
     desc: Clean up
diff --git a/inference_gateway/client.py b/inference_gateway/client.py
@@ -1,8 +1,7 @@
-from typing import Generator, Optional
+from typing import Generator, Optional, Union, List, Dict, Optional
 import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Dict, Optional
 import requests
 
 
@@ -37,36 +36,67 @@ def to_dict(self) -> Dict[str, str]:
 @dataclass
 class Model:
     """Represents an LLM model"""
+
     name: str
 
 
 @dataclass
 class ProviderModels:
     """Groups models by provider"""
+
     provider: Provider
     models: List[Model]
 
 
 @dataclass
 class ResponseTokens:
     """Response tokens structure as defined in the API spec"""
+
     role: str
     model: str
     content: str
 
+    @classmethod
+    def from_dict(cls, data: dict) -> "ResponseTokens":
+        """Create ResponseTokens from dictionary data
+
+        Args:
+            data: Dictionary containing response data
+
+        Returns:
+            ResponseTokens instance
+
+        Raises:
+            TypeError: If data is not a dictionary
+            ValueError: If required fields are missing
+        """
+        if not isinstance(data, dict):
+            raise TypeError(f"Expected dict, got {type(data)}")
+
+        required = ["role", "model", "content"]
+        missing = [field for field in required if field not in data]
+
+        if missing:
+            raise ValueError(
+                f"Missing required arguments: {
+                    ', '.join(missing)}"
+            )
+
+        return cls(role=data["role"], model=data["model"], content=data["content"])
+
 
 @dataclass
 class GenerateResponse:
     """Response structure for token generation"""
+
     provider: str
     response: ResponseTokens
 
     @classmethod
-    def from_dict(cls, data: dict) -> 'GenerateResponse':
+    def from_dict(cls, data: dict) -> "GenerateResponse":
         """Create GenerateResponse from dictionary data"""
         return cls(
-            provider=data.get('provider', ''),
-            response=ResponseTokens(**data.get('response', {}))
+            provider=data.get("provider", ""), response=ResponseTokens(**data.get("response", {}))
         )
 
 
@@ -86,9 +116,79 @@ def list_models(self) -> List[ProviderModels]:
         response.raise_for_status()
         return response.json()
 
+    def _parse_sse_chunk(self, chunk: bytes) -> dict:
+        """Parse an SSE message chunk into structured event data
+
+        Args:
+            chunk: Raw SSE message chunk in bytes format
+
+        Returns:
+            dict: Parsed SSE message with event type and data fields
+
+        Raises:
+            json.JSONDecodeError: If chunk format or content is invalid
+        """
+        if not isinstance(chunk, bytes):
+            raise TypeError(f"Expected bytes, got {type(chunk)}")
+
+        try:
+            decoded = chunk.decode("utf-8")
+            message = {}
+
+            for line in (l.strip() for l in decoded.split("\n") if l.strip()):
+                if line.startswith("event: "):
+                    message["event"] = line.removeprefix("event: ")
+                elif line.startswith("data: "):
+                    try:
+                        json_str = line.removeprefix("data: ")
+                        data = json.loads(json_str)
+                        if not isinstance(data, dict):
+                            raise json.JSONDecodeError(
+                                f"Invalid SSE data format - expected object, got: {
+                                    json_str}",
+                                json_str,
+                                0,
+                            )
+                        message["data"] = data
+                    except json.JSONDecodeError as e:
+                        raise json.JSONDecodeError(f"Invalid SSE JSON: {json_str}", e.doc, e.pos)
+
+            if not message.get("data"):
+                raise json.JSONDecodeError(
+                    f"Missing or invalid data field in SSE message: {
+                        decoded}",
+                    decoded,
+                    0,
+                )
+
+            return message
+
+        except UnicodeDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Invalid UTF-8 encoding in SSE chunk: {
+                    chunk!r}",
+                str(chunk),
+                0,
+            )
+
+    def _parse_json_line(self, line: bytes) -> ResponseTokens:
+        """Parse a single JSON line into GenerateResponse"""
+        try:
+            decoded_line = line.decode("utf-8")
+            data = json.loads(decoded_line)
+            return ResponseTokens.from_dict(data)
+        except UnicodeDecodeError as e:
+            raise json.JSONDecodeError(f"Invalid UTF-8 encoding: {line}", str(line), 0)
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Invalid JSON response: {
+                    decoded_line}",
+                e.doc,
+                e.pos,
+            )
+
     def generate_content(self, provider: Provider, model: str, messages: List[Message]) -> Dict:
-        payload = {"model": model, "messages": [
-            msg.to_dict() for msg in messages]}
+        payload = {"model": model, "messages": [msg.to_dict() for msg in messages]}
 
         response = self.session.post(
             f"{self.base_url}/llms/{provider.value}/generate", json=payload
@@ -97,12 +197,8 @@ def generate_content(self, provider: Provider, model: str, messages: List[Messag
         return response.json()
 
     def generate_content_stream(
-        self,
-        provider: Provider,
-        model: str,
-        messages: List[Message],
-        use_sse: bool = False
-    ) -> Generator[Union[GenerateResponse, dict], None, None]:
+        self, provider: Provider, model: str, messages: List[Message], use_sse: bool = False
+    ) -> Generator[Union[ResponseTokens, dict], None, None]:
         """Stream content generation from the model
 
         Args:
@@ -112,33 +208,37 @@ def generate_content_stream(
             use_sse: Whether to use Server-Sent Events format
 
         Yields:
-            Either GenerateResponse objects (for raw JSON) or dicts (for SSE)
+            Either ResponseTokens objects (for raw JSON) or dicts (for SSE)
         """
         payload = {
             "model": model,
             "messages": [msg.to_dict() for msg in messages],
             "stream": True,
-            "ssevents": use_sse
+            "ssevents": use_sse,
         }
 
-        with self.session.post(
-            f"{self.base_url}/llms/{provider.value}/generate",
-            json=payload,
-            stream=True
-        ) as response:
-            response.raise_for_status()
+        response = self.session.post(
+            f"{self.base_url}/llms/{provider.value}/generate", json=payload, stream=True
+        )
+        response.raise_for_status()
+
+        if use_sse:
+            buffer = []
 
             for line in response.iter_lines():
-                if line:
-                    if use_sse and line.startswith(b'data: '):
-                        # Handle SSE format
-                        data = json.loads(line.decode(
-                            'utf-8').replace('data: ', ''))
-                        yield data
-                    else:
-                        # Handle raw JSON format
-                        data = json.loads(line)
-                        yield GenerateResponse.from_dict(data)
+                if not line:
+                    if buffer:
+                        chunk = b"\n".join(buffer)
+                        yield self._parse_sse_chunk(chunk)
+                        buffer = []
+                    continue
+
+                buffer.append(line)
+        else:
+            for line in response.iter_lines():
+                if not line:
+                    continue
+                yield self._parse_json_line(line)
 
     def health_check(self) -> bool:
         """Check if the API is healthy"""
diff --git a/tests/test_client.py b/tests/test_client.py
