quivr/backend/core/tests/test_brains.py

import random
import string


def test_retrieve_default_brain(client, api_key):
    # Making a GET request to the /brains/default/ endpoint
    response = client.get(
        "/brains/default/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    default_brain = response.json()
    assert "id" in default_brain
    assert "name" in default_brain


def test_create_brain(client, api_key):
    # Generate a random name for the brain
    random_brain_name = "".join(
        random.choices(string.ascii_letters + string.digits, k=10)
    )

    # Set up the request payload
    payload = {
        "name": random_brain_name,
        "status": "public",
        "model": "gpt-3.5-turbo-0613",
        "temperature": 0,
        "max_tokens": 256,
        "file_sha1": "",
    }

    # Making a POST request to the /brains/ endpoint
    response = client.post(
        "/brains/",
        json=payload,
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    # Optionally, assert on specific fields in the response
    response_data = response.json()
    # e.g., assert that the response contains a 'brain_id' field
    assert "id" in response_data
    assert "name" in response_data

    # Optionally, assert that the returned 'name' matches the one sent in the request
    assert response_data["name"] == payload["name"]


def test_retrieve_all_brains(client, api_key):
    # Making a GET request to the /brains/ endpoint to retrieve all brains for the current user
    response = client.get(
        "/brains/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    response_data = response.json()
    # Optionally, you can loop through the brains and assert on specific fields in each brain
    for brain in response_data["brains"]:
        assert "id" in brain
        assert "name" in brain


def test_retrieve_one_brain(client, api_key):
    # Making a GET request to the /brains/default/ endpoint
    response = client.get(
        "/brains/default/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    response_data = response.json()

    # Extract the brain_id from the response
    brain_id = response_data["id"]

    # Making a GET request to the /brains/{brain_id}/ endpoint
    response = client.get(
        f"/brains/{brain_id}/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    brain = response.json()
    assert "id" in brain
    assert "name" in brain


def test_delete_all_brains(client, api_key):
    # First, retrieve all brains for the current user
    response = client.get(
        "/brains/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200

    response_data = response.json()

    # Loop through each brain and send a DELETE request
    for brain in response_data["brains"]:
        brain_id = brain["id"]

        # Send a DELETE request to delete the specific brain
        delete_response = client.delete(
            f"/brains/{brain_id}/subscription",
            headers={"Authorization": "Bearer " + api_key},
        )

        # Assert that the DELETE response status code is 200 (HTTP OK)
        assert delete_response.status_code == 200

    # Finally, retrieve all brains for the current user
    response = client.get(
        "/brains/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200
    response_data = response.json()
    assert len(response_data["brains"]) == 0


def test_delete_all_brains_and_get_default_brain(client, api_key):
    # First create a new brain
    test_create_brain(client, api_key)

    # Now, retrieve all brains for the current user
    response = client.get(
        "/brains/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200
    assert len(response.json()["brains"]) > 0

    test_delete_all_brains(client, api_key)

    # Get the default brain, it should create one if it doesn't exist
    response = client.get(
        "/brains/default/",
        headers={"Authorization": "Bearer " + api_key},
    )

    # Assert that the response status code is 200 (HTTP OK)
    assert response.status_code == 200
    assert response.json()["name"] == "Default brain"