Unit Tests for /columns API Endpoint (#1060)

* unit test for /columns django endpoint

* 🎨 Auto-generated directory tree for repository in Architecture.md

* 🎨 Format Python code with psf/black

* black formatting

* remove unneeded prints

---------

Co-authored-by: karkir0003 <[email protected]>

training/tests/test_dataset_endpoint.py

@@ -0,0 +1,29 @@
+import pytest
+from tests.utils.test_utils import mock_authenticate, get_test_bearer_token
+from django.test import Client
+from training.core.authenticator import FirebaseAuth
+
+
+@pytest.mark.parametrize(
+    "dataset_name", ["IRIS", "CALIFORNIA_HOUSING", "DIABETES", "WINE"]
+)
+def test_columns_endpoint(monkeypatch, dataset_name):
+    client = Client()
+    # Use monkeypatch to replace FirebaseAuth.authenticate with our mock function
+    monkeypatch.setattr(FirebaseAuth, "authenticate", mock_authenticate)
+    # Set the Authorization header with the fake token
+    headers = get_test_bearer_token()
+    response = client.get(f"/api/datasets/default/{dataset_name}/columns", **headers)
+    assert response.status_code == 200
+    assert isinstance(response.json(), dict)
+
+
+@pytest.mark.parametrize("dataset_name", ["TEST_DATASET", "HELLO", None])
+def test_columns_invalid_default(monkeypatch, dataset_name):
+    client = Client()
+    # Use monkeypatch to replace FirebaseAuth.authenticate with our mock function
+    monkeypatch.setattr(FirebaseAuth, "authenticate", mock_authenticate)
+    headers = get_test_bearer_token()
+    response = client.get(f"/api/datasets/default/{dataset_name}/columns", **headers)
+    assert response.status_code == 404
+    assert response.content.decode("utf-8") == '{"message": "Dataset not found"}'

training/tests/utils/test_utils.py

@@ -0,0 +1,32 @@
+"""
+File that houses helper functions for testing the training backend 
+"""
+import jwt
+import datetime
+
+
+def mock_authenticate(*args, **kwargs) -> str:
+    """
+    Function that gives a test JWT Token for testing (not necessarily real user data)
+    Django API Endpoints that require user authentication
+
+    Returns:
+        token: Bearer Token
+    """
+    payload = {
+        "sub": "1234567890",
+        "name": "John Doe",
+        "iat": datetime.datetime.utcnow(),
+        "exp": datetime.datetime.utcnow() + datetime.timedelta(days=1),
+    }
+    secret = "secret"
+    token = jwt.encode(payload, secret, algorithm="HS256")
+    return token
+
+
+def get_test_bearer_token() -> dict:
+    """
+    Wrapper that uses mock_authenticate function to build a bearer token
+    in a format that Django accepts
+    """
+    return {"HTTP_AUTHORIZATION": "Bearer " + mock_authenticate()}