Create test_main.py

.github/workflows/python-app.yml

@@ -28,6 +28,8 @@ jobs:
         python -m pip install --upgrade pip
         pip install flake8 pytest
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+        # create .env file with an API Key
+        echo "YT_API_KEY=${{ secrets.YT_API_KEY }}" > .env
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

main.py

@@ -8,8 +8,8 @@ from models import init_emotions_model
 
 class Settings(BaseSettings):
     YT_API_KEY: str
-    PRED_BATCH_SIZE: int
-    MAX_COMMENT_SIZE: int
+    PRED_BATCH_SIZE: int = 512
+    MAX_COMMENT_SIZE: int = 300
     model_config = SettingsConfigDict(env_file='.env')
 
 

requirements.txt

@@ -7,3 +7,5 @@ torchvision
 torchaudio
 transformers
 pandas
+pytest
+httpx

test_main.py

@@ -0,0 +1,27 @@
+from fastapi.testclient import TestClient
+from main import app
+import pandas as pd
+
+
+client = TestClient(app)
+
+
+def test_home():
+    """Test home page."""
+    response = client.get("/")
+    assert response.status_code == 200
+
+
+def test_predict():
+    """Test predict method on an example video."""
+    TEST_VIDEO_ID = "0peXnOnDgQ8"
+    response = client.get(
+        "/predict/",
+        params={"video_id": TEST_VIDEO_ID}
+    )
+    df = pd.read_json(response, orient='records')
+
+    # Ensure the DataFrame has the right amount of columns
+    assert df.shape[1] == 39
+    # Ensure there are no NaN values
+    assert df.isna().sum().sum() == 0