add api and processor tests

Author: divakaivan

tests/test_api.py

@@ -0,0 +1,49 @@
+import pytest
+from fastapi.testclient import TestClient
+
+from make_api.app.main import app
+
+client = TestClient(app)
+
+
+@pytest.fixture
+def sample_input():
+    return {
+        "passenger_count": 2,
+        "trip_type": 1,
+        "congestion_surcharge": 2.5,
+        "mean_distance": 3.2,
+        "mean_duration": 7.5,
+        "rush_hour": 1,
+        "vendor_id": "VTS",
+    }
+
+
+def test_predict_one(sample_input):
+    """Test the predict_one endpoint works"""
+
+    response = client.post("/predict", json=sample_input)
+
+    assert response.status_code == 200
+
+    json_response = response.json()
+
+    assert "prediction" in json_response
+    assert isinstance(json_response["prediction"], float)
+
+
+def test_predict_missing_input():
+    """Test the predict_one endpoint fails with missing input"""
+
+    incomplete_input = {
+        "passenger_count": 2,
+        "trip_type": 1,
+        "congestion_surcharge": 2.5,
+        "mean_distance": 3.2,
+        "rush_hour": 1,
+        "vendor_id": "VTS",
+    }  # missing mean_duration
+
+    response = client.post("/predict", json=incomplete_input)
+
+    assert response.status_code == 422

tests/test_data_processor.py

@@ -0,0 +1,54 @@
+from unittest.mock import MagicMock
+
+import pandas as pd
+import pytest
+
+from make_data.data_processor import DataProcessor
+
+
+@pytest.fixture
+def sample_dataframe():
+    data = {
+        "lpep_pickup_datetime": ["2024-02-08 10:00:00", "2024-02-08 15:00:00"],
+        "lpep_dropoff_datetime": ["2024-02-08 10:15:00", "2024-02-08 15:30:00"],
+        "fare_amount": [10, -5],
+        "PULocationID": [1, 2],
+        "DOLocationID": [3, 4],
+        "trip_distance": [2.5, 3.5],
+        "VendorID": [1, 2],
+    }
+    df = pd.DataFrame(data)
+    return df
+
+
+@pytest.fixture
+def mock_config():
+    config = MagicMock()
+    config.num_features = ["fare_amount", "trip_distance"]
+    config.cat_features = ["vendor_id"]
+    config.target = ["duration"]
+    return config
+
+
+def test_process_data(sample_dataframe, mock_config):
+    """Test the process_data method"""
+
+    processor = DataProcessor(sample_dataframe, mock_config)
+    processor.process_data()
+
+    assert "duration" in processor.df.columns
+    assert "vendor_id" in processor.df.columns
+    assert processor.df["fare_amount"].min() >= 0
+    assert processor.df["duration"].min() >= 0
+
+
+def test_split_data(sample_dataframe, mock_config):
+    """Test the split_data method"""
+
+    processor = DataProcessor(sample_dataframe, mock_config)
+    processor.process_data()
+    train, test = processor.split_data()
+
+    assert len(train) > 0
+    assert len(test) > 0
+    assert len(train) + len(test) == len(processor.df)