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 assert "status" 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"