Add test files for files, predictions, models, and feedback (#38)

Added test files following the structure of test_hub.py:

- test_files.py
- test_predictions.py
- test_models.py
- test_feedback.py

All tests follow similar patterns and use the mock_client fixture.

Link to Devin run:
https://app.devin.ai/sessions/11db6d50f539482f9424308c5e7ee913

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Sudeep Pillai <[email protected]>

Author: devin-ai-integration[bot]

tests/conftest.py

@@ -10,6 +10,9 @@
     HubInfoResponse,
     HubDomainsResponse,
     HubSchemaQueryResponse,
+    FileResponse,
+    PredictionResponse,
+    FeedbackSubmitResponse,
 )
 
 
@@ -24,6 +27,20 @@ def mock_client(monkeypatch):
     """Mock the Client class."""
 
     class MockClient:
+        class AudioPredictions:
+            def __init__(self, client):
+                self._client = client
+
+            def generate(self, *args, **kwargs):
+                return PredictionResponse(
+                    id="prediction1",
+                    status="completed",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at="2024-01-01T00:00:01Z",
+                    response={"result": "test"},
+                    usage={"total_tokens": 100}
+                )
+
         def __init__(self, api_key=None, base_url=None):
             self.api_key = api_key or "test-key"
             self.base_url = base_url or "https://api.vlm.run"
@@ -33,9 +50,11 @@ def __init__(self, api_key=None, base_url=None):
             self.files = self.Files(self)
             self.models = self.Models(self)
             self.hub = self.Hub(self)
-            self.image = self.Image(self)
-            self.video = self.Video(self)
-            self.document = self.Document(self)
+            self.image = self.ImagePredictions(self)
+            self.video = self.VideoPredictions(self)
+            self.document = self.DocumentPredictions(self)
+            self.audio = self.AudioPredictions(self)
+            self.feedback = self.Feedback(self)
 
         class Dataset:
             def __init__(self, client):
@@ -99,44 +118,89 @@ def __init__(self, client):
                 self._client = client
 
             def create(self, model, prompt, **kwargs):
-                return {"id": "prediction1"}
+                return PredictionResponse(
+                    id="prediction1",
+                    status="running",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at=None,
+                    response=None,
+                    usage={"total_tokens": 0}
+                )
 
             def list(self):
-                return [{"id": "prediction1", "status": "running"}]
+                return [PredictionResponse(
+                    id="prediction1",
+                    status="running",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at=None,
+                    response=None,
+                    usage={"total_tokens": 0}
+                )]
 
             def get(self, prediction_id):
-                return {"id": prediction_id, "status": "running"}
+                return PredictionResponse(
+                    id=prediction_id,
+                    status="running",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at=None,
+                    response=None,
+                    usage={"total_tokens": 0}
+                )
+                
+            def wait(self, prediction_id, timeout=60, sleep=1):
+                return PredictionResponse(
+                    id=prediction_id,
+                    status="completed",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at="2024-01-01T00:00:01Z",
+                    response={"result": "test"},
+                    usage={"total_tokens": 100}
+                )
 
         class Files:
             def __init__(self, client):
                 self._client = client
 
             def list(self):
                 return [
-                    {
-                        "id": "file1",
-                        "filename": "test.txt",
-                        "size": 100,
-                        "created_at": "2024-01-01",
-                    }
+                    FileResponse(
+                        id="file1",
+                        filename="test.txt",
+                        bytes=b"test content",
+                        purpose="assistants",
+                        created_at="2024-01-01T00:00:00Z"
+                    )
                 ]
 
             def upload(self, file_path, purpose="fine-tune"):
-                return {"id": "file1", "filename": file_path}
+                return FileResponse(
+                    id="file1",
+                    filename=str(file_path),
+                    bytes=b"test content",
+                    purpose=purpose,
+                    created_at="2024-01-01T00:00:00Z"
+                )
 
             def get(self, file_id):
-                return {
-                    "id": file_id,
-                    "filename": "test.txt",
-                    "size": 100,
-                    "created_at": "2024-01-01",
-                }
+                return FileResponse(
+                    id=file_id,
+                    filename="test.txt",
+                    bytes=b"test content",
+                    purpose="assistants",
+                    created_at="2024-01-01T00:00:00Z"
+                )
 
             def get_content(self, file_id):
                 return b"test content"
 
             def delete(self, file_id):
-                return True
+                return FileResponse(
+                    id=file_id,
+                    filename="test.txt",
+                    bytes=b"test content",
+                    purpose="assistants",
+                    created_at="2024-01-01T00:00:00Z"
+                )
 
         class Models:
             def __init__(self, client):
@@ -175,26 +239,67 @@ def get_schema(self, domain):
                     schema_hash="abcd1234",
                 )
 
-        class Image:
+        class ImagePredictions:
             def __init__(self, client):
                 self._client = client
 
             def generate(self, *args, **kwargs):
-                return b"image data"
+                return PredictionResponse(
+                    id="prediction1",
+                    status="completed",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at="2024-01-01T00:00:01Z",
+                    response={"result": "test"},
+                    usage={"total_tokens": 100}
+                )
 
-        class Video:
+        class VideoPredictions:
             def __init__(self, client):
                 self._client = client
 
             def generate(self, *args, **kwargs):
-                return b"video data"
+                return PredictionResponse(
+                    id="prediction1",
+                    status="completed",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at="2024-01-01T00:00:01Z",
+                    response={"result": "test"},
+                    usage={"total_tokens": 100}
+                )
 
-        class Document:
+        class DocumentPredictions:
             def __init__(self, client):
                 self._client = client
 
             def generate(self, *args, **kwargs):
-                return b"document data"
+                return PredictionResponse(
+                    id="prediction1",
+                    status="completed",
+                    created_at="2024-01-01T00:00:00Z",
+                    completed_at="2024-01-01T00:00:01Z",
+                    response={"result": "test"},
+                    usage={"total_tokens": 100}
+                )
+
+        class Feedback:
+            def __init__(self, client):
+                self._client = client
+
+            def submit(self, id, label=None, notes=None, flag=None):
+                return FeedbackSubmitResponse(
+                    id="feedback1",
+                    created_at="2024-01-01T00:00:00Z",
+                    request_id=id,
+                    response=label
+                )
+
+            def get(self, id):
+                return FeedbackSubmitResponse(
+                    id="feedback1",
+                    created_at="2024-01-01T00:00:00Z",
+                    request_id=id,
+                    response=None
+                )
 
     monkeypatch.setattr("vlmrun.cli.cli.Client", MockClient)
     return MockClient()

tests/test_feedback.py

@@ -0,0 +1,30 @@
+"""Tests for feedback operations."""
+
+from pydantic import BaseModel
+from vlmrun.client.types import FeedbackSubmitResponse
+
+
+class TestLabel(BaseModel):
+    """Test label model."""
+    score: int
+    comment: str
+
+
+def test_submit_feedback(mock_client):
+    """Test submitting feedback for a prediction."""
+    label = TestLabel(score=5, comment="Great prediction!")
+    response = mock_client.feedback.submit(
+        id="prediction1",
+        label=label,
+        notes="Test feedback",
+        flag=False
+    )
+    assert isinstance(response, FeedbackSubmitResponse)
+    assert response.id == "feedback1"
+
+
+def test_get_feedback(mock_client):
+    """Test getting feedback by ID."""
+    response = mock_client.feedback.get("prediction1")
+    assert isinstance(response, FeedbackSubmitResponse)
+    assert response.id == "feedback1"

tests/test_files.py

@@ -0,0 +1,49 @@
+"""Tests for files operations."""
+
+from pathlib import Path
+from vlmrun.client.types import FileResponse
+
+
+def test_list_files(mock_client):
+    """Test listing files."""
+    response = mock_client.files.list()
+    assert isinstance(response, list)
+    assert all(isinstance(file, FileResponse) for file in response)
+    assert len(response) > 0
+    assert response[0].id == "file1"
+    assert response[0].filename == "test.txt"
+
+
+def test_upload_file(mock_client, tmp_path):
+    """Test uploading a file."""
+    # Create a temporary file
+    test_file = tmp_path / "test.txt"
+    test_file.write_text("test content")
+    
+    response = mock_client.files.upload(test_file)
+    assert isinstance(response, FileResponse)
+    assert response.id == "file1"
+    assert response.filename == str(test_file)
+
+
+def test_get_file(mock_client):
+    """Test getting file metadata."""
+    response = mock_client.files.get("file1")
+    assert isinstance(response, FileResponse)
+    assert response.id == "file1"
+    assert response.filename == "test.txt"
+    assert len(response.bytes) == len(b"test content")
+
+
+def test_get_content(mock_client):
+    """Test getting file content."""
+    response = mock_client.files.get_content("file1")
+    assert isinstance(response, bytes)
+    assert response == b"test content"
+
+
+def test_delete_file(mock_client):
+    """Test deleting a file."""
+    response = mock_client.files.delete("file1")
+    assert isinstance(response, FileResponse)
+    assert response.id == "file1"

tests/test_models.py

@@ -0,0 +1,16 @@
+"""Tests for models operations."""
+
+from vlmrun.client.types import ModelInfoResponse
+
+
+def test_list_models(mock_client):
+    """Test listing available models."""
+    response = mock_client.models.list()
+    assert isinstance(response, list)
+    assert all(isinstance(model, ModelInfoResponse) for model in response)
+    assert len(response) > 0
+    model = response[0]
+    assert isinstance(model.model, str)
+    assert isinstance(model.domain, str)
+    assert model.model == "model1"
+    assert model.domain == "test-domain"