diff --git a/training/tests/test_dataset_endpoint.py b/training/tests/test_dataset_endpoint.py new file mode 100644 index 000000000..48b353a7e --- /dev/null +++ b/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"}' diff --git a/training/tests/utils/__init__.py b/training/tests/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/training/tests/utils/test_utils.py b/training/tests/utils/test_utils.py new file mode 100644 index 000000000..c522a295c --- /dev/null +++ b/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()}