import random import string from repository.brain.get_default_user_brain import get_user_default_brain 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" def test_set_as_default_brain_endpoint(client, api_key): 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, } # Making a POST request to the /brains/ endpoint response = client.post( "/brains/", json=payload, headers={"Authorization": "Bearer " + api_key}, ) response_data = response.json() brain_id = response_data["id"] # Make a POST request to set the brain as default for the user response = client.post( f"/brains/{brain_id}/default", headers={"Authorization": "Bearer " + api_key}, ) # Assert that the response status code is 200 (HTTP OK) assert response.status_code == 200 # Assert the response message assert response.json() == { "message": f"Brain {brain_id} has been set as default brain." } # Check if the brain is now the default for the user # Send a request to get user information response = client.get("/user", headers={"Authorization": "Bearer " + api_key}) # Assert that the response contains the expected fields user_info = response.json() user_id = user_info["id"] default_brain = get_user_default_brain(user_id) assert default_brain is not None assert default_brain.brain_id == brain_id