gradientai_test.go

package godo

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var listAgentResponse = `
{
	"agents": [
		{
			"uuid": "00000000-0000-0000-0000-000000000000",
			"name": "AMD ROCm Documentation Agent",
			"created_at": "2025-05-08T03:37:28Z",
			"updated_at": "2025-05-09T16:02:39Z",
			"instruction": "providing answers to common questions about the platform",
			"model": {
				"uuid": "00000000-0000-0000-0000-000000000000",
				"name": "Llama 3.3 Instruct (70B)",
				"inference_name": "llama3.3-70b-instruct",
				"version": {
					"major": 1
				},
				"inference_version": "4200000096",
				"is_foundational": true,
				"upload_complete": true,
				"created_at": "2025-01-13T20:56:20Z",
				"updated_at": "2025-05-13T15:16:21Z",
				"metadata": {
					"agreements": {
						"list": [
							{
								"title": "Licensing Agreement",
								"url": "https://llama.meta.com/llama3_1/license/"
							}
						],
						"title": "Llama 3.3"
					},
					"description": "An advanced language model with greater capabilities due to its larger size.",
					"max_tokens": {
						"default": 512,
						"max": 2048,
						"min": 256
					},
					"temperature": {
						"default": 0.7,
						"max": 1,
						"min": 0
					},
					"top_p": {
						"default": 0.9,
						"max": 1,
						"min": 0.1
					}
				},
				"parent_uuid": "00000000-0000-0000-0000-000000000000",
				"agreement": {
					"uuid": "00000000-0000-0000-0000-000000000000",
					"name": "Meta Llama 3.3 Community License",
					"description": "Meta Llama 3.3 is licensed under the Meta Llama 3.3 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved. By purchasing, deploying, accessing, or using this model, you agree to comply with the",
					"url": "https://www.llama.com/llama3_3/license/"
				},
				"usecases": [
					"MODEL_USECASE_AGENT",
					"MODEL_USECASE_SERVERLESS"
				]
			},
			"deployment": {
				"uuid": "00000000-0000-0000-0000-000000000000",
				"url": "https://bndu2aqk2tnxj6ldwsatdmyu.agents.do-ai.run",
				"status": "STATUS_RUNNING",
				"visibility": "VISIBILITY_PUBLIC",
				"created_at": "2025-05-08T03:37:28Z",
				"updated_at": "2025-05-12T21:07:15Z"
			},
			"k": 0,
			"temperature": 0.7,
			"top_p": 0.9,
			"max_tokens": 2048,
			"project_id": "00000000-0000-0000-0000-000000000000",
			"route_uuid": "00000000-0000-0000-0000-000000000000",
			"region": "tor1",
			"chatbot": {
				"name": "AMD ROCm Documentation Agent",
				"primary_color": "#031B4E",
				"secondary_color": "#E5E8ED",
				"starting_message": "Hello! How can assist you with some documentation today?",
				"button_background_color": "#0061EB",
				"logo": "http://159.203.26.153/AMD-Logo.svg"
			},
			"route_created_at": "0001-01-01T00:00:00Z",
			"user_id": "14567334",
			"chatbot_identifiers": [
				{
					"agent_chatbot_identifier": "dfsvgdvscfvgf"
				}
			],
			"retrieval_method": "RETRIEVAL_METHOD_NONE"
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/agents?page=1&per_page=1",
			"previous": "https://api.digitalocean.com/v2/gen-ai/agents?page=1&per_page=1",
			"next": "https://api.digitalocean.com/v2/gen-ai/agents?page=3&per_page=1",
			"last": "https://api.digitalocean.com/v2/gen-ai/agents?page=34&per_page=1"
		}
	},
	"meta": {
		"total": 34,
		"page": 2,
		"pages": 34
	}
}
`

var agentResponse = `
{
	"agent": {
		"uuid": "00000000-0000-0000-0000-000000000000",
		"name": "testing-godo",
		"created_at": "2025-05-14T13:18:05Z",
		"updated_at": "2025-05-14T13:18:05Z",
		"instruction": "You are an agent who thinks deeply about the world",
		"description": "My Agent Description",
		"model": {
			"uuid": "00000000-0000-0000-0000-000000000000",
			"name": "Llama 3.3 Instruct (70B)",
			"inference_name": "llama3.3-70b-instruct",
			"version": {
				"major": 1
			},
			"inference_version": "423242296",
			"is_foundational": true,
			"upload_complete": true,
			"created_at": "2025-01-13T20:56:20Z",
			"updated_at": "2025-05-13T15:16:21Z",
			"metadata": {
				"agreements": {
					"list": [
						{
							"title": "Licensing Agreement",
							"url": "https://llama.meta.com/llama3_1/license/"
						}
					],
					"title": "Llama 3.3"
				},
				"description": "An advanced language model with greater capabilities due to its larger size.",
				"max_tokens": {
					"default": 512,
					"max": 2048,
					"min": 256
				},
				"temperature": {
					"default": 0.7,
					"max": 1,
					"min": 0
				},
				"top_p": {
					"default": 0.9,
					"max": 1,
					"min": 0.1
				}
			},
			"parent_uuid": "00000000-0000-0000-0000-000000000000",
			"agreement": {
				"uuid": "00000000-0000-0000-0000-000000000000",
				"name": "Meta Llama 3.3 Community License",
				"description": "Meta Llama 3.3 is licensed under the Meta Llama 3.3 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved. By purchasing, deploying, accessing, or using this model, you agree to comply with the",
				"url": "https://www.llama.com/llama3_3/license/"
			},
			"usecases": [
				"MODEL_USECASE_AGENT",
				"MODEL_USECASE_SERVERLESS"
			]
		},
		"deployment": {
			"uuid": "00000000-0000-0000-0000-000000000000",
			"status": "STATUS_WAITING_FOR_DEPLOYMENT",
			"visibility": "VISIBILITY_PRIVATE",
			"created_at": "2025-05-14T13:18:05Z",
			"updated_at": "2025-05-14T13:18:05Z"
		},
		"knowledge_bases": [
			{
				"uuid": "00000000-0000-0000-0000-000000000000",
				"name": "genai-api-testing",
				"created_at": "2025-05-14T07:52:27Z",
				"updated_at": "2025-05-14T07:57:27Z",
				"region": "nyc3",
				"embedding_model_uuid": "00000000-0000-0000-0000-000000000000",
				"project_id": "00000000-0000-0000-0000-000000000000",
				"database_id": "00000000-0000-0000-0000-000000000000",
				"last_indexing_job": {
					"uuid": "00000000-0000-0000-0000-000000000000",
					"knowledge_base_uuid": "00000000-0000-0000-0000-000000000000",
					"created_at": "2025-05-14T07:57:27Z",
					"updated_at": "2025-05-14T08:07:14Z",
					"started_at": "2025-05-14T07:57:27Z",
					"finished_at": "2025-05-14T07:58:05Z",
					"phase": "BATCH_JOB_PHASE_SUCCEEDED",
					"total_datasources": 1,
					"completed_datasources": 1,
					"tokens": 10060
				}
			}
		],
		"api_keys": [
			{
				"api_key": "dassfbgnhgbdsfvbgdvscfvg"
			}
		],
		"k": 10,
		"temperature": 0.7,
		"top_p": 0.9,
		"max_tokens": 512,
		"tags": [
			"string"
		],
		"project_id": "00000000-0000-0000-0000-000000000000",
		"route_uuid": "00000000-0000-0000-0000-000000000000",
		"region": "TOR1",
		"route_created_at": "0001-01-01T00:00:00Z",
		"user_id": "18919793",
		"chatbot_identifiers": [
			{
				"agent_chatbot_identifier": "dsfgbfsdsadsfbgfdfvsc"
			}
		],
		"retrieval_method": "RETRIEVAL_METHOD_NONE",
		"functions": [
        {
            "name": "godo-test-function",
            "description": "Creating Function Route",
            "faas_name": "godo-test-faasname",
            "faas_namespace": "fn-00000000-0000-0000-0000-000000000000",
            "input_schema": {
                "parameters": [
                    {
                        "name": "zipCode",
                        "in": "query",
                        "schema": {
                            "type": "string"
                        },
                        "required": false,
                        "description": "The ZIP code for which to fetch the weather"
                    },
                    {
                        "name": "measurement",
                        "in": "query",
                        "schema": {
                            "type": "string",
                            "enum": [
                                "F",
                                "C"
                            ]
                        },
                        "required": false,
                        "description": "The measurement unit for temperature (F or C)"
                    }
                ]
            },
            "output_schema": {
                "properties": [
                    {
                        "name": "temperature",
                        "type": "number",
                        "description": "The temperature for the specified location"
                    },
                    {
                        "name": "measurement",
                        "type": "string",
                        "description": "The measurement unit used for the temperature (F or C)"
                    },
                    {
                        "name": "conditions",
                        "type": "string",
                        "description": "A description of the current weather conditions (Sunny, Cloudy, etc)"
                    }
                ]
            }
        }
    ]
		
	}
}
`

var agentUpdateResponse = `
{
	"agent": {
		"uuid": "00000000-0000-0000-0000-000000000000",
		"name": "My Agent",
		"created_at": "2025-05-14T13:22:05Z",
		"updated_at": "2025-05-14T13:46:48Z",
		"instruction": "You are an agent who thinks deeply about the world",
		"description": "My Agent Description",
		"deployment": {
			"uuid": "00000000-0000-0000-0000-000000000000",
			"url": "https://hv5kn7pdaawp5772uk7nmyzd.agents.do-ai.run",
			"status": "STATUS_RUNNING",
			"visibility": "VISIBILITY_PUBLIC",
			"created_at": "2025-05-14T13:22:05Z",
			"updated_at": "2025-05-14T13:59:39Z"
		},
		"k": 10,
		"temperature": 0.7,
		"top_p": 0.9,
		"max_tokens": 512,
		"tags": [
			"updated",
			"example"
		],
		"project_id": "00000000-0000-0000-0000-000000000000",
		"route_uuid": "00000000-0000-0000-0000-000000000000",
		"region": "tor1",
		"chatbot": {
			"name": "My Agent Chatbot",
			"primary_color": "#031B4E",
			"secondary_color": "#E5E8ED",
			"starting_message": "Hello! How can I help you today?",
			"button_background_color": "#0061EB"
		},
		"route_created_at": "0001-01-01T00:00:00Z",
		"user_id": "18793",
		"retrieval_method": "RETRIEVAL_METHOD_NONE",
		"functions": [
        {
            "name": "godo-test-function",
            "description": "Updating Function Route",
            "faas_name": "godo-test-faasname",
            "faas_namespace": "fn-00000000-0000-0000-0000-000000000000",
            "input_schema": {
                "parameters": [
                    {
                        "name": "zipCode",
                        "in": "query",
                        "schema": {
                            "type": "string"
                        },
                        "required": false,
                        "description": "The ZIP code for which to fetch the weather"
                    },
                    {
                        "name": "measurement",
                        "in": "query",
                        "schema": {
                            "type": "string",
                            "enum": [
                                "F",
                                "C"
                            ]
                        },
                        "required": false,
                        "description": "The measurement unit for temperature (F or C)"
                    }
                ]
            },
            "output_schema": {
                "properties": [
                    {
                        "name": "temperature",
                        "type": "number",
                        "description": "The temperature for the specified location"
                    },
                    {
                        "name": "measurement",
                        "type": "string",
                        "description": "The measurement unit used for the temperature (F or C)"
                    },
                    {
                        "name": "conditions",
                        "type": "string",
                        "description": "A description of the current weather conditions (Sunny, Cloudy, etc)"
                    }
                ]
            }
        }
    ]
	}
}
`

var listAvailableModelsResponse = `
{
	"models": [
		{
			"uuid": "00000000-0000-0000-0000-000000000000",
			"name": "Llama 3.3 Instruct (70B)",
			"description": "An advanced language model with greater capabilities due to its larger size.",
			"version": {
				"major": 1
			},
			"is_foundational": true,
			"upload_complete": true,
			"created_at": "2025-01-13T20:56:20Z",
			"updated_at": "2025-05-13T15:16:21Z",
			"parent_uuid": "00000000-0000-0000-0000-000000000000",
			"agreement": {
				"uuid": "00000000-0000-0000-0000-000000000000",
				"name": "Meta Llama 3.3 Community License",
				"description": "Meta Llama 3.3 is licensed under the Meta Llama 3.3 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved. By purchasing, deploying, accessing, or using this model, you agree to comply with the",
				"url": "https://www.llama.com/llama3_3/license/"
			},
			"model_availability": "SERVERLESS",
			"context_window": "131072",
			"capabilities": ["inference", "chat"],
			"modalities": {
				"input": ["text"],
				"output": ["text"]
			},
			"parameter_count": 70.0,
			"type": "chat"
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/models?page=1&per_page=1",
			"next": "https://api.digitalocean.com/v2/gen-ai/models?page=2&per_page=1",
			"last": "https://api.digitalocean.com/v2/gen-ai/models?page=15&per_page=1"
		}
	},
	"meta": {
		"total": 15,
		"page": 1,
		"pages": 15
	}
}
`

var listInferenceRoutersResponse = `
{
	"model_routers": [
		{
			"uuid": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
			"name": "router-one",
			"description": "Test router for routing prompts."
		}
	],
	"meta": {
		"total": 1,
		"page": 1,
		"pages": 1
	}
}
`

var listInferenceRouterTaskPresetsResponse = `
{
	"tasks": [
		{
			"task_slug": "translation",
			"name": "Translation",
			"category": "general",
			"description": "Translate text between languages.",
			"models": ["openai-gpt-oss-120b"],
			"tags": ["general"]
		},
		{
			"task_slug": "code-generation",
			"name": "Code Generation",
			"category": "software-engineering",
			"description": "Generate code from natural language.",
			"models": ["openai-gpt-oss-120b", "anthropic-claude-3-5-sonnet"],
			"tags": ["code"]
		}
	],
	"meta": {
		"total": 2,
		"page": 1,
		"pages": 1
	}
}
`

var getInferenceRouterResponse = `
{
	"model_router": {
		"uuid": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
		"name": "router-one",
		"description": "Router used for unit tests.",
		"created_at": "2025-06-01T12:00:00Z",
		"updated_at": "2025-06-02T15:30:00Z",
		"config": {
			"fallback_models": ["openai-gpt-4"],
			"policies": [
				{"task_slug": "code-generation", "models": ["llama3"], "selection_policy": {"prefer": "fastest"}}
			]
		}
	}
}
`

var listAPIKeysResponse = `
{
    "api_key_infos": [
        {
            "uuid": "00000000-0000-0000-0000-000000000000",
            "name": "Key One",
            "secret_key": "1000000",
            "created_at": "2025-05-14T13:18:05Z",
            "created_by": "12345678"
        },
        {
            "uuid": "00000000-0000-0000-0000-000000000000",
            "name": "Key Two",
            "secret_key": "1000000",
            "created_at": "2025-05-15T13:18:05Z",
            "created_by": "12345678"
        }
    ]
}
`

var apiKeyInfoResponse = `
{
    "api_key_info": {
        "uuid": "00000000-0000-0000-0000-000000000000",
        "name": "Key One",
        "secret_key": "1000000",
        "created_at": "2025-05-14T13:18:05Z",
        "created_by": "12345678"
    }
}
`

var listKnowledgeBaseResponse = `
{
	"knowledge_bases": [
		{
			"created_at": "2025-05-08T03:37:28Z",
			"database_id": "11111111-1111-1111-1111-111111111111",
			"embedding_model_uuid": "11111111-1111-1111-1111-111111111111",
			"last_indexing_job": {
				"created_at": "2025-05-08T03:37:28Z",
				"finished_at": "2025-05-08T03:38:28Z",
				"knowledge_base_uuid": "11111111-1111-1111-1111-111111111111",
				"phase": "BATCH_JOB_PHASE_SUCCEEDED",
				"started_at": "2025-05-08T03:37:28Z",
				"updated_at": "2025-05-08T03:38:28Z",
				"total_datasources": 1,
				"completed_datasources": 1,
				"tokens": 5000,
				"uuid": "11111111-1111-1111-1111-111111111111"
			},
			"name": "Test Knowledge Base",
			"project_id": "11111111-1111-1111-1111-111111111111",
			"region": "tor1",
			"tags": ["test"],
			"uuid": "11111111-1111-1111-1111-111111111111",
			"updated_at": "2025-05-09T16:02:39Z",
			"is_public": false,
			"user_id": "14567334"
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases?page=1&per_page=1",
			"previous": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases?page=1&per_page=1",
			"next": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases?page=3&per_page=1",
			"last": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases?page=10&per_page=1"
		}
	},
	"meta": {
		"total": 10,
		"page": 1,
		"pages": 10
	}
}
`

var listIndexingJobsResponse = `
{
	"jobs": [
		{
			"uuid": "22222222-2222-2222-2222-222222222222",
			"knowledge_base_uuid": "11111111-1111-1111-1111-111111111111",
			"created_at": "2025-05-08T03:37:28Z",
			"started_at": "2025-05-08T03:37:30Z",
			"finished_at": "2025-05-08T03:38:28Z",
			"updated_at": "2025-05-08T03:38:28Z",
			"phase": "BATCH_JOB_PHASE_SUCCEEDED",
			"total_datasources": 2,
			"completed_datasources": 2,
			"tokens": 10500,
			"data_source_uuids": ["33333333-3333-3333-3333-333333333333", "44444444-4444-4444-4444-444444444444"]
		},
		{
			"uuid": "55555555-5555-5555-5555-555555555555",
			"knowledge_base_uuid": "66666666-6666-6666-6666-666666666666",
			"created_at": "2025-05-07T02:30:15Z",
			"started_at": "2025-05-07T02:30:20Z",
			"finished_at": "2025-05-07T02:32:45Z",
			"updated_at": "2025-05-07T02:32:45Z",
			"phase": "BATCH_JOB_PHASE_SUCCEEDED",
			"total_datasources": 1,
			"completed_datasources": 1,
			"tokens": 7500,
			"data_source_uuids": ["77777777-7777-7777-7777-777777777777"]
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/indexing_jobs?page=1&per_page=2",
			"next": "https://api.digitalocean.com/v2/gen-ai/indexing_jobs?page=2&per_page=2",
			"last": "https://api.digitalocean.com/v2/gen-ai/indexing_jobs?page=5&per_page=2"
		}
	},
	"meta": {
		"total": 9,
		"page": 1,
		"pages": 5
	}
}
`

var indexingJobDataSourcesResponse = `
{
	"indexed_data_sources": [
		{
			"completed_at": "2025-05-08T03:38:28Z",
			"data_source_uuid": "33333333-3333-3333-3333-333333333333",
			"error_details": "",
			"error_msg": "",
			"failed_item_count": "0",
			"indexed_file_count": "5",
			"indexed_item_count": "150",
			"removed_item_count": "0",
			"skipped_item_count": "2",
			"started_at": "2025-05-08T03:37:30Z",
			"status": "DATA_SOURCE_STATUS_COMPLETED",
			"total_bytes": "2048000",
			"total_bytes_indexed": "1900000",
			"total_file_count": "7"
		},
		{
			"completed_at": "2025-05-08T03:38:25Z",
			"data_source_uuid": "44444444-4444-4444-4444-444444444444",
			"error_details": "File format not supported",
			"error_msg": "Some files could not be processed",
			"failed_item_count": "3",
			"indexed_file_count": "8",
			"indexed_item_count": "200",
			"removed_item_count": "1",
			"skipped_item_count": "5",
			"started_at": "2025-05-08T03:37:32Z",
			"status": "DATA_SOURCE_STATUS_COMPLETED_WITH_ERRORS",
			"total_bytes": "3072000",
			"total_bytes_indexed": "2800000",
		"total_file_count": "12"
	}
]
}
`

var indexingJobResponse = `
{
	"job": {
		"uuid": "22222222-2222-2222-2222-222222222222",
		"knowledge_base_uuid": "11111111-1111-1111-1111-111111111111",
		"status": "INDEX_JOB_STATUS_COMPLETED",
		"phase": "BATCH_JOB_PHASE_COMPLETED",
		"created_at": "2025-05-08T03:30:00Z",
		"started_at": "2025-05-08T03:30:05Z",
		"finished_at": "2025-05-08T03:37:45Z",
		"updated_at": "2025-05-08T03:37:45Z",
		"data_source_uuids": [
			"33333333-3333-3333-3333-333333333333",
			"44444444-4444-4444-4444-444444444444"
		],
		"total_datasources": 2,
		"completed_datasources": 2,
		"tokens": 1250,
		"total_items_indexed": "350",
		"total_items_failed": "3",
		"total_items_skipped": "7"
	}
}
`

var cancelIndexingJobResponse = `
{
	"job": {
		"uuid": "22222222-2222-2222-2222-222222222222",
		"knowledge_base_uuid": "11111111-1111-1111-1111-111111111111",
		"status": "INDEX_JOB_STATUS_CANCELLED",
		"phase": "BATCH_JOB_PHASE_CANCELLED",
		"created_at": "2025-05-08T03:30:00Z",
		"started_at": "2025-05-08T03:30:05Z",
		"finished_at": "2025-05-08T03:35:00Z",
		"updated_at": "2025-05-08T03:35:00Z",
		"data_source_uuids": [
			"33333333-3333-3333-3333-333333333333"
		],
		"total_datasources": 1,
		"completed_datasources": 0,
		"tokens": 0
	}
}
`

var knowledgeBaseResponse = `
{
	"knowledge_base": {
		"uuid": "11111111-1111-1111-1111-111111111111",
		"name": "testing-kb",
		"created_at": "2025-05-14T13:18:05Z",
		"updated_at": "2025-05-14T13:18:05Z",
		"region": "tor1",
		"project_id": "11111111-1111-1111-1111-111111111111",
		"embedding_model_uuid": "11111111-1111-1111-1111-111111111111",
		"database_id": "11111111-1111-1111-1111-111111111111",
		"is_public": false,
		"tags": ["string"],
		"user_id": "18919793"
	}
}
`

var knowledgeBaseGetResponse = `
{
	"database_status" : "ONLINE",
	"knowledge_base": {
		"uuid": "11111111-1111-1111-1111-111111111111",
		"name": "testing-kb",
		"created_at": "2025-05-14T13:18:05Z",
		"updated_at": "2025-05-14T13:18:05Z",
		"region": "tor1",
		"project_id": "11111111-1111-1111-1111-111111111111",
		"embedding_model_uuid": "11111111-1111-1111-1111-111111111111",
		"database_id": "11111111-1111-1111-1111-111111111111",
		"is_public": false,
		"tags": ["string"],
		"user_id": "18919793"
	}
}
`

var knowledgeBaseUpdateResponse = `
{
	"knowledge_base": {
		"uuid": "11111111-1111-1111-1111-111111111111",
		"name": "Updated-Knowledge-Base",
		"created_at": "2025-05-14T13:18:05Z",
		"updated_at": "2025-05-14T13:46:48Z",
		"region": "tor1",
		"project_id": "11111111-1111-1111-1111-111111111111",
		"embedding_model_uuid": "11111111-1111-1111-1111-111111111111",
		"database_id": "11111111-1111-1111-1111-111111111111",
		"is_public": true,
		"tags": ["updated", "example"],
		"user_id": "18919793"
	}
}
`

var listDataSourcesResponse = `
{
	"knowledge_base_data_sources": [
		{
			"uuid": "22222222-2222-2222-2222-222222222222",
			"created_at": "2025-05-14T13:18:05Z",
			"updated_at": "2025-05-14T13:18:05Z",
			"spaces_data_source": {
				"bucket_name": "test-bucket",
				"item_path": "/docs/test.pdf",
				"region": "tor1"
			}
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources?page=1&per_page=1",
			"previous": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources?page=1&per_page=1",
			"next": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources?page=2&per_page=1",
			"last": "https://api.digitalocean.com/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources?page=3&per_page=1"
		}
	},
	"meta": { 
		"total": 3,
		"page": 1,
		"pages": 3
	}
}
`

var addDataSourceResponse = `
{
	"knowledge_base_data_source": {
		"uuid": "22222222-2222-2222-2222-222222222222",
		"created_at": "2025-05-14T13:18:05Z",
		"updated_at": "2025-05-14T13:18:05Z",
		"spaces_data_source": {
			"bucket_name": "test-bucket",
			"item_path": "/docs/test.pdf",
			"region": "tor1"
		}
	}
}
`

var deleteDataSourceResponse = `
{
	"knowledge_base_uuid": "11111111-1111-1111-1111-111111111111",
	"data_source_uuid": "22222222-2222-2222-2222-222222222222"
}
`

var deleteKnowledgeBaseResponse = `
{
	"uuid": "11111111-1111-1111-1111-111111111111"
}`

var agentRouteResponse = `
{
    "child_agent_uuid": "00000000-0000-0000-0000-000000000001",
    "parent_agent_uuid": "00000000-0000-0000-0000-000000000000",
    "rollback": false,
    "uuid": "00000000-0000-0000-0000-000000000003"
}
`

var listAgentVersionsResponse = `
{
	"agent_versions": [
		{
			"agent_uuid": "00000000-0000-0000-0000-000000000000",
			"currently_applied": true,
			"version_hash": "00000000000000000000000000000000000000000000000000000000000000",
			"created_at": "2025-06-25T09:32:26Z",
			"created_by_email": "example@gmail.com",
			"trigger_action": "Rolled back from version ABCDEDF",
			"model_name": "Llama 3.3 Instruct (70B)",
			"name": "normal-agent",
			"instruction": "You are an agent who thinks deeply about the world",
			"description": "Think about the world deeply",
			"tags": [
				"example string"
			],
			"k": 10,
			"temperature": 0.7,
			"top_p": 0.9,
			"max_tokens": 512,
			"retrieval_method": "RETRIEVAL_METHOD_NONE",
			"attached_functions": [
				{
					"name": "terraform-tf",
					"faas_name": "default/testing",
					"faas_namespace": "fn-b90f0000-0000-0000-0000-75edfbb6f397",
					"is_deleted": true
				}
			],
			"can_rollback": true
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/agents/01efde4/versions?page=1&per_page=5",
			"next": "https://api.digitalocean.com/v2/gen-ai/agents/01efde4/versions?page=2&per_page=5",
			"last": "https://api.digitalocean.com/v2/gen-ai/agents/01efde4/versions?page=13&per_page=5"
		}
	},
	"meta": {
		"total": 65,
		"page": 1,
		"pages": 13
	}
}
`
var listAnthropicAPIKeysResponse = `
{
    "api_key_infos": [
        {
            "uuid": "11111111-1111-1111-1111-111111111111",
            "name": "Anthropic Key One",
             "api_key": "sk-ant-1",
            "created_at": "2025-05-14T13:18:05Z",
            "created_by": "user-1"
        },
        {
            "uuid": "22222222-2222-2222-2222-222222222222",
            "name": "Anthropic Key Two",
             "api_key": "sk-ant-2",
            "created_at": "2025-05-15T13:18:05Z",
            "created_by": "user-2"
        }
    ],
    "links": {
        "pages": {
            "first": "https://api.digitalocean.com/v2/gen-ai/anthropic/keys?page=1&per_page=1",
            "next": "https://api.digitalocean.com/v2/gen-ai/anthropic/keys?page=2&per_page=1",
            "last": "https://api.digitalocean.com/v2/gen-ai/anthropic/keys?page=10&per_page=1"
        }
    },
    "meta": {
        "total": 2,
        "page": 1,
        "pages": 10
    }
}
`

var anthropicAPIKeyInfoResponse = `
{
    "api_key_info": {
        "uuid": "11111111-1111-1111-1111-111111111111",
        "name": "Anthropic Key One",
        "api_key": "sk-ant-1",
        "created_at": "2025-05-14T13:18:05Z",
        "created_by": "user-1"
    }
}
`
var listAgentsByAnthropicAPIKeyResponse = `
{
  "agents": [
    {
      "uuid": "00000000-0000-0000-0000-000000000000",
      "name": "Anthropic Agent 1"
    },
    {
      "uuid": "00000000-0000-0000-0000-000000000001",
      "name": "Anthropic Agent 2"
    }
  ],
  "links": {
    "pages": {
      "first": "https://api.digitalocean.com/v2/gen-ai/anthropic/keys?page=1&per_page=1"
    }
  },
  "meta": {
    "total": 2,
    "page": 1,
    "pages": 1
  }
}
`

var rollbackResponse = `
{
	"version_hash": "00000000000000000000000000000000000000000000000000000000000001"
}
`

var listOpenAIAPIKeysResponse = `
{
    "api_key_infos": [
        {
            "uuid": "11111111-1111-1111-1111-111111111111",
            "created_at": "2025-05-14T13:18:05Z",
            "created_by": "user-1",
			"name": "openai-key1",
			"models": [{
				"agreement": {
					"description": "openai example",
					"name": "example name",
					"url": "example string",
					"uuid": "123e4567-e89b-12d3-a456-426614174000"
				},
				"created_at": "2023-01-01T00:00:00Z",
				"inference_name": "example name",
				"inference_version": "example string",
				"is_foundational": true,
				"metadata": { },
				"name": "example name",
				"parent_uuid": "123e4567-e89b-12d3-a456-426614174000",
				"provider": "MODEL_PROVIDER_DIGITALOCEAN",
				"updated_at": "2023-01-01T00:00:00Z",
				"upload_complete": true,
				"url": "example string",
				"usecases": ["MODEL_USECASE_AGENT","MODEL_USECASE_GUARDRAIL"],
				"uuid": "123e4567-e89b-12d3-a456-426614174000",
				"version": {
					"major": 123,
					"minor": 123,
					"patch": 123
				}
			}]
        },
		 {
            "uuid": "22222222-2222-2222-2222-222222222222",
            "created_at": "2025-05-14T13:18:05Z",
            "created_by": "user-2",
			"name": "openai-key2",
			"models": [{
				"agreement": {
					"description": "openai example",
					"name": "example name",
					"url": "example string",
					"uuid": "123e4567-e89b-12d3-a456-426614174000"
				},
				"created_at": "2023-01-01T00:00:00Z",
				"inference_name": "example name",
				"inference_version": "example string",
				"is_foundational": true,
				"metadata": { },
				"name": "example name",
				"parent_uuid": "123e4567-e89b-12d3-a456-426614174000",
				"provider": "MODEL_PROVIDER_DIGITALOCEAN",
				"updated_at": "2023-01-01T00:00:00Z",
				"upload_complete": true,
				"url": "example string",
				"usecases": ["MODEL_USECASE_AGENT","MODEL_USECASE_GUARDRAIL"],
				"uuid": "123e4567-e89b-12d3-a456-426614174000",
				"version": {
					"major": 123,
					"minor": 123,
					"patch": 123
				}
			}]
        }
    ],
    "links": {
        "pages": {
            "first": "https://api.digitalocean.com/v2/gen-ai/openai/keys?page=1&per_page=1",
            "next": "https://api.digitalocean.com/v2/gen-ai/openai/keys?page=2&per_page=1",
            "last": "https://api.digitalocean.com/v2/gen-ai/openai/keys?page=10&per_page=1"
        }
    },
    "meta": {
        "total": 2,
        "page": 1,
        "pages": 10
    }
}
`

var openaiAPIKeyInfoResponse = `
{
    "api_key_info": {
            "uuid": "11111111-1111-1111-1111-111111111111",
            "created_at": "2025-05-14T13:18:05Z",
            "created_by": "user-1",
			"name": "OpenAI One",
			"models": [{
				"agreement": {
					"description": "openai example",
					"name": "example name",
					"url": "example string",
					"uuid": "123e4567-e89b-12d3-a456-426614174000"
				},
				"created_at": "2023-01-01T00:00:00Z",
				"inference_name": "example name",
				"inference_version": "example string",
				"is_foundational": true,
				"metadata": { },
				"name": "example name",
				"parent_uuid": "123e4567-e89b-12d3-a456-426614174000",
				"provider": "MODEL_PROVIDER_DIGITALOCEAN",
				"updated_at": "2023-01-01T00:00:00Z",
				"upload_complete": true,
				"url": "example string",
				"usecases": ["MODEL_USECASE_AGENT","MODEL_USECASE_GUARDRAIL"],
				"uuid": "123e4567-e89b-12d3-a456-426614174000",
				"version": {
					"major": 123,
					"minor": 123,
					"patch": 123
				}
			}]
        }
}
`

var listAgentsByOpenAIAPIKeyResponse = `
{
  "agents": [
    {
      "uuid": "00000000-0000-0000-0000-000000000000",
      "name": "OpenAI Agent 1"
    },
    {
      "uuid": "00000000-0000-0000-0000-000000000001",
      "name": "OpenAI Agent 2"
    }
  ],
  "links": {
    "pages": {
      "first": "https://api.digitalocean.com/v2/gen-ai/openai/keys?page=1&per_page=1"
    }
  },
  "meta": {
    "total": 2,
    "page": 1,
    "pages": 1
  }
}
`

var listDatacenterRegionsResponse = `
{
	"regions": [
		{
			"region": "tor1",
			"inference_url": "https://tor1.gen-ai.digitalocean.com",
			"serves_batch": true,
			"serves_inference": true,
			"stream_inference_url": "https://tor1.gen-ai.digitalocean.com/stream"
		},
		{
			"region": "nyc3",
			"inference_url": "https://nyc3.gen-ai.digitalocean.com",
			"serves_batch": false,
			"serves_inference": true,
			"stream_inference_url": "https://nyc3.gen-ai.digitalocean.com/stream"
		}
	]
}
`

func TestListAgents(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "1",
		})

		fmt.Fprint(w, listAgentResponse)
	})

	req := &ListOptions{
		Page:    1,
		PerPage: 1,
	}

	agents, resp, err := client.GradientAI.ListAgents(ctx, req)
	if err != nil {
		t.Errorf("GenAI.ListAgents returned error: %v", err)
	}

	var agent Agent
	err = json.Unmarshal([]byte(listAgentResponse), &agent)
	if err != nil {
		t.Errorf("Error unmarshalling response: %v", err)
	}

	assert.Equal(t, 34, resp.Meta.Total)
	expectedString := fmt.Sprintf("%v", agents[0])
	assert.Equal(t, expectedString, agents[0].String())
	assert.Equal(t, agent.K, agents[0].K)
}

func TestCreateAgent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, agentResponse)
	})

	req := &AgentCreateRequest{
		Description:       "My Agent Description",
		Instruction:       "You are an agent who thinks deeply about the world",
		KnowledgeBaseUuid: []string{"00000000-0000-0000-0000-000000000000"},
		ModelUuid:         "00000000-0000-0000-0000-000000000000",
		Name:              "testing-godo",
		ProjectId:         "00000000-0000-0000-0000-000000000000",
		Region:            "tor1",
		Tags:              []string{"string"},
		WorkspaceUuid:     "00000000-0000-0000-0000-000000000000",
	}

	res, _, err := client.GradientAI.CreateAgent(ctx, req)
	if err != nil {
		t.Errorf("GenAI.Create returned error: %v", err)
	}

	assert.Equal(t, res.Name, req.Name)
	assert.Equal(t, res.Description, req.Description)
}

func TestListAPIKeys(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/api_keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAPIKeysResponse)
	})

	keys, resp, err := client.GradientAI.ListAgentAPIKeys(ctx, "00000000-0000-0000-0000-000000000000", nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, 2, len(keys))
	assert.Equal(t, "Key One", keys[0].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", keys[0].Uuid)
	assert.Equal(t, "Key Two", keys[1].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", keys[1].Uuid)
}

func TestCreateAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/api_keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, apiKeyInfoResponse)
	})

	req := &AgentAPIKeyCreateRequest{
		AgentUuid: "00000000-0000-0000-0000-000000000000",
		Name:      "Key One",
	}

	key, resp, err := client.GradientAI.CreateAgentAPIKey(ctx, "00000000-0000-0000-0000-000000000000", req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "Key One", key.Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", key.Uuid)
}

func TestUpdateAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/api_keys/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, apiKeyInfoResponse)
	})

	req := &AgentAPIKeyUpdateRequest{
		AgentUuid:  "00000000-0000-0000-0000-000000000000",
		APIKeyUuid: "00000000-0000-0000-0000-000000000000",
		Name:       "Key One",
	}

	key, resp, err := client.GradientAI.UpdateAgentAPIKey(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000000", req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "Key One", key.Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", key.Uuid)
}

func TestDeleteAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/api_keys/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, apiKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.DeleteAgentAPIKey(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000000")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "Key One", key.Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", key.Uuid)
}

func TestRegenerateAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/api_keys/00000000-0000-0000-0000-000000000000/regenerate", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, apiKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.RegenerateAgentAPIKey(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000000")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "Key One", key.Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", key.Uuid)
}

func TestGetAgent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, agentResponse)
	})

	res, resp, err := client.GradientAI.GetAgent(ctx, "00000000-0000-0000-0000-000000000000")
	if err != nil {
		t.Errorf("GenAI.Get returned error: %v", err)
	}

	assert.Equal(t, res.Name, "testing-godo")
	assert.Equal(t, resp.Response.StatusCode, 200)
	expectedString := fmt.Sprintf("%v", res)
	assert.Equal(t, expectedString, res.String())
}

func TestDeleteAgent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/def5d52c-30c5-11f0-bf8f-4e013e2ddde4", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, agentResponse)
	})

	res, resp, err := client.GradientAI.DeleteAgent(ctx, "def5d52c-30c5-11f0-bf8f-4e013e2ddde4")
	if err != nil {
		t.Errorf("GenAI.Delete returned error: %v", err)
	}

	assert.Equal(t, resp.Response.StatusCode, 200)
	assert.Equal(t, "testing-godo", res.Name)
}

func TestUpdateAgent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, agentUpdateResponse)
	})

	req := &AgentUpdateRequest{
		Tags: []string{"updated", "example"},
	}

	res, resp, err := client.GradientAI.UpdateAgent(ctx, "00000000-0000-0000-0000-000000000000", req)
	if err != nil {
		t.Errorf("GenAI.Update returned error: %v", err)
	}

	assert.Equal(t, res.Tags[0], req.Tags[0])
	assert.Equal(t, resp.Response.StatusCode, 200)
}

func TestUpdateAgentVisibility(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/deployment_visibility", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, agentUpdateResponse)
	})

	req := &AgentVisibilityUpdateRequest{
		Uuid:       "00000000-0000-0000-0000-000000000000",
		Visibility: "VISIBILITY_PRIVATE",
	}

	res, resp, err := client.GradientAI.UpdateAgentVisibility(ctx, "00000000-0000-0000-0000-000000000000", req)
	if err != nil {
		t.Errorf("GenAI.UpdateVisibility returned error: %v", err)
	}

	assert.Equal(t, res.Uuid, req.Uuid)
	assert.Equal(t, res.Name, "My Agent")
	assert.Equal(t, resp.Response.StatusCode, 200)
}

func TestListKnowledgeBases(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "1",
		})

		fmt.Fprint(w, listKnowledgeBaseResponse)
	})

	req := &ListOptions{
		Page:    1,
		PerPage: 1,
	}

	knowledgeBases, resp, err := client.GradientAI.ListKnowledgeBases(ctx, req)
	if err != nil {
		t.Errorf("GenAI.ListKnowledgeBases returned error: %v", err)
	}

	assert.Equal(t, 10, resp.Meta.Total)
	assert.Equal(t, "Test Knowledge Base", knowledgeBases[0].Name)
}

func TestListIndexingJobs(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/indexing_jobs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "2",
		})

		fmt.Fprint(w, listIndexingJobsResponse)
	})

	req := &ListOptions{
		Page:    1,
		PerPage: 2,
	}

	result, resp, err := client.GradientAI.ListIndexingJobs(ctx, req)
	if err != nil {
		t.Errorf("GenAI.ListIndexingJobs returned error: %v", err)
	}

	assert.Equal(t, 9, resp.Meta.Total)
	assert.Equal(t, 2, len(result.Jobs))
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", result.Jobs[0].Uuid)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", result.Jobs[0].KnowledgeBaseUuid)
	assert.Equal(t, "BATCH_JOB_PHASE_SUCCEEDED", result.Jobs[0].Phase)
	assert.Equal(t, 2, result.Jobs[0].TotalDatasources)
	assert.Equal(t, 2, result.Jobs[0].CompletedDatasources)
	assert.Equal(t, 10500, result.Jobs[0].Tokens)
	assert.Equal(t, 2, len(result.Jobs[0].DataSourceUuids))
	assert.Equal(t, "55555555-5555-5555-5555-555555555555", result.Jobs[1].Uuid)
}

func TestListIndexingJobDataSources(t *testing.T) {
	setup()
	defer teardown()

	indexingJobUUID := "22222222-2222-2222-2222-222222222222"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/indexing_jobs/%s/data_sources", indexingJobUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, indexingJobDataSourcesResponse)
	})

	result, resp, err := client.GradientAI.ListIndexingJobDataSources(ctx, indexingJobUUID)
	if err != nil {
		t.Errorf("GenAI.ListIndexingJobDataSources returned error: %v", err)
	}

	assert.Equal(t, 2, len(result.IndexedDataSources))

	// Test first data source
	ds1 := result.IndexedDataSources[0]
	assert.Equal(t, "33333333-3333-3333-3333-333333333333", ds1.DataSourceUuid)
	assert.Equal(t, "DATA_SOURCE_STATUS_COMPLETED", ds1.Status)
	assert.Equal(t, "5", ds1.IndexedFileCount)
	assert.Equal(t, "150", ds1.IndexedItemCount)
	assert.Equal(t, "0", ds1.FailedItemCount)
	assert.Equal(t, "2", ds1.SkippedItemCount)
	assert.Equal(t, "2048000", ds1.TotalBytes)
	assert.Equal(t, "1900000", ds1.TotalBytesIndexed)
	assert.Equal(t, "7", ds1.TotalFileCount)

	// Test second data source with errors
	ds2 := result.IndexedDataSources[1]
	assert.Equal(t, "44444444-4444-4444-4444-444444444444", ds2.DataSourceUuid)
	assert.Equal(t, "DATA_SOURCE_STATUS_COMPLETED_WITH_ERRORS", ds2.Status)
	assert.Equal(t, "8", ds2.IndexedFileCount)
	assert.Equal(t, "200", ds2.IndexedItemCount)
	assert.Equal(t, "3", ds2.FailedItemCount)
	assert.Equal(t, "5", ds2.SkippedItemCount)
	assert.Equal(t, "File format not supported", ds2.ErrorDetails)
	assert.Equal(t, "Some files could not be processed", ds2.ErrorMsg)

	_ = resp // Mark as used to avoid unused variable warning
}

func TestGetIndexingJob(t *testing.T) {
	setup()
	defer teardown()

	indexingJobUUID := "22222222-2222-2222-2222-222222222222"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/indexing_jobs/%s", indexingJobUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, indexingJobResponse)
	})

	result, resp, err := client.GradientAI.GetIndexingJob(ctx, indexingJobUUID)
	if err != nil {
		t.Errorf("GenAI.GetIndexingJob returned error: %v", err)
	}

	// Test the job details
	job := result.Job
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", job.Uuid)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", job.KnowledgeBaseUuid)
	assert.Equal(t, "INDEX_JOB_STATUS_COMPLETED", job.Status)
	assert.Equal(t, "BATCH_JOB_PHASE_COMPLETED", job.Phase)
	assert.Equal(t, 2, job.TotalDatasources)
	assert.Equal(t, 2, job.CompletedDatasources)
	assert.Equal(t, 1250, job.Tokens)
	assert.Equal(t, "350", job.TotalItemsIndexed)
	assert.Equal(t, "3", job.TotalItemsFailed)
	assert.Equal(t, "7", job.TotalItemsSkipped)

	// Test data source UUIDs array
	assert.Equal(t, 2, len(job.DataSourceUuids))
	assert.Equal(t, "33333333-3333-3333-3333-333333333333", job.DataSourceUuids[0])
	assert.Equal(t, "44444444-4444-4444-4444-444444444444", job.DataSourceUuids[1])

	// Test timestamps
	assert.NotNil(t, job.CreatedAt)
	assert.NotNil(t, job.StartedAt)
	assert.NotNil(t, job.FinishedAt)
	assert.NotNil(t, job.UpdatedAt)

	_ = resp // Mark as used to avoid unused variable warning
}

func TestCancelIndexingJob(t *testing.T) {
	setup()
	defer teardown()

	indexingJobUUID := "22222222-2222-2222-2222-222222222222"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/indexing_jobs/%s/cancel", indexingJobUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)

		// Verify the request payload
		v := new(CancelIndexingJobRequest)
		err := json.NewDecoder(r.Body).Decode(v)
		if err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}
		assert.Equal(t, indexingJobUUID, v.UUID)

		fmt.Fprint(w, cancelIndexingJobResponse)
	})

	result, resp, err := client.GradientAI.CancelIndexingJob(ctx, indexingJobUUID)
	if err != nil {
		t.Errorf("GenAI.CancelIndexingJob returned error: %v", err)
	}

	// Test the cancelled job details
	job := result.Job
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", job.Uuid)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", job.KnowledgeBaseUuid)
	assert.Equal(t, "INDEX_JOB_STATUS_CANCELLED", job.Status)
	assert.Equal(t, "BATCH_JOB_PHASE_CANCELLED", job.Phase)
	assert.Equal(t, 1, job.TotalDatasources)
	assert.Equal(t, 0, job.CompletedDatasources)
	assert.Equal(t, 0, job.Tokens)

	// Test data source UUIDs array
	assert.Equal(t, 1, len(job.DataSourceUuids))
	assert.Equal(t, "33333333-3333-3333-3333-333333333333", job.DataSourceUuids[0])

	// Test timestamps
	assert.NotNil(t, job.CreatedAt)
	assert.NotNil(t, job.StartedAt)
	assert.NotNil(t, job.FinishedAt)
	assert.NotNil(t, job.UpdatedAt)

	_ = resp // Mark as used to avoid unused variable warning
}

func TestCreateKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, knowledgeBaseResponse)
	})

	req := &KnowledgeBaseCreateRequest{
		Name:               "testing-kb",
		ProjectID:          "11111111-1111-1111-1111-111111111111",
		Region:             "tor1",
		EmbeddingModelUuid: "11111111-1111-1111-1111-111111111111",
		Tags:               []string{"string"},
		DataSources: []KnowledgeBaseDataSource{
			{
				WebCrawlerDataSource: &WebCrawlerDataSource{
					BaseUrl: "https://www.example.com",
				},
			},
		},
	}

	res, _, err := client.GradientAI.CreateKnowledgeBase(ctx, req)
	if err != nil {
		t.Errorf("GenAI.CreateKnowledgeBase returned error: %v", err)
	}

	assert.Equal(t, res.Name, req.Name)
	assert.Equal(t, res.ProjectId, req.ProjectID)
	assert.Equal(t, req.EmbeddingModelUuid, res.EmbeddingModelUuid)
	assert.Equal(t, req.Region, res.Region)
}

func TestListDataSources(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "1",
		})

		fmt.Fprint(w, listDataSourcesResponse)
	})

	req := &ListOptions{
		Page:    1,
		PerPage: 1,
	}

	dataSources, resp, err := client.GradientAI.ListKnowledgeBaseDataSources(ctx, "11111111-1111-1111-1111-111111111111", req)
	if err != nil {
		t.Errorf("GenAI.ListDataSources returned error: %v", err)
	}

	assert.Equal(t, 3, resp.Meta.Total)
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", dataSources[0].Uuid)
}

func TestAddDataSource(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, addDataSourceResponse)
	})

	req := &AddKnowledgeBaseDataSourceRequest{
		KnowledgeBaseUuid: "11111111-1111-1111-1111-111111111111",
		SpacesDataSource: &SpacesDataSource{
			BucketName: "test-bucket",
			ItemPath:   "/docs/test.pdf",
			Region:     "tor1",
		},
	}

	res, _, err := client.GradientAI.AddKnowledgeBaseDataSource(ctx, "11111111-1111-1111-1111-111111111111", req)
	if err != nil {
		t.Errorf("GenAI.AddDataSource returned error: %v", err)
	}

	assert.Equal(t, "test-bucket", res.SpacesDataSource.BucketName)
	assert.Equal(t, "/docs/test.pdf", res.SpacesDataSource.ItemPath)
}

func TestDeleteDataSource(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111/data_sources/22222222-2222-2222-2222-222222222222", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, deleteDataSourceResponse)
	})

	kbUUID, dsUUID, resp, err := client.GradientAI.DeleteKnowledgeBaseDataSource(ctx, "11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222")
	if err != nil {
		t.Errorf("GenAI.DeleteDataSource returned error: %v", err)
	}

	assert.Equal(t, "11111111-1111-1111-1111-111111111111", kbUUID)
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", dsUUID)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestGetKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, knowledgeBaseGetResponse)
	})

	res, dbStatus, resp, err := client.GradientAI.GetKnowledgeBase(ctx, "11111111-1111-1111-1111-111111111111")
	if err != nil {
		t.Errorf("GenAI.GetKnowledgeBase returned error: %v", err)
	}

	assert.Equal(t, "testing-kb", res.Name)
	assert.Equal(t, "ONLINE", dbStatus)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", res.Uuid)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestUpdateKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, knowledgeBaseUpdateResponse)
	})

	req := &UpdateKnowledgeBaseRequest{
		Name: "Updated-Knowledge-Base",
		Tags: []string{"updated", "example"},
	}

	res, resp, err := client.GradientAI.UpdateKnowledgeBase(ctx, "11111111-1111-1111-1111-111111111111", req)
	if err != nil {
		t.Errorf("GenAI.UpdateKnowledgeBase returned error: %v", err)
	}

	assert.Equal(t, req.Name, res.Name)
	assert.Equal(t, req.Tags, res.Tags)
	assert.Equal(t, res.Tags[0], req.Tags[0])
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestDeleteKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/knowledge_bases/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, deleteKnowledgeBaseResponse)
	})

	kbUUID, resp, err := client.GradientAI.DeleteKnowledgeBase(ctx, "11111111-1111-1111-1111-111111111111")
	if err != nil {
		t.Errorf("GenAI.DeleteKnowledgeBase returned error: %v", err)
	}

	assert.Equal(t, "11111111-1111-1111-1111-111111111111", kbUUID)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestAttachKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/knowledge_bases/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, agentResponse)
	})

	res, resp, err := client.GradientAI.AttachKnowledgeBaseToAgent(ctx, "00000000-0000-0000-0000-000000000000", "11111111-1111-1111-1111-111111111111")
	if err != nil {
		t.Errorf("GenAI.AttachKnowledgBase returned error: %v", err)
	}

	assert.Equal(t, "testing-godo", res.Name)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestDetachKnowledgeBase(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/knowledge_bases/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, agentResponse)
	})

	res, resp, err := client.GradientAI.DetachKnowledgeBaseToAgent(ctx, "00000000-0000-0000-0000-000000000000", "11111111-1111-1111-1111-111111111111")
	fmt.Print(res)
	fmt.Print(resp)

	if err != nil {
		t.Fatalf("GenAI.DetachKnowledgBase returned error: %v", err)
	}

	if res == nil {
		t.Fatalf("GenAI.DetachKnowledgBase returned nil response or result")
	}

	assert.Equal(t, "testing-godo", res.Name)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestAddAgentRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/child_agents/00000000-0000-0000-0000-000000000001", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, agentRouteResponse)
	})

	req := &AgentRouteCreateRequest{
		ChildAgentUuid:  "00000000-0000-0000-0000-000000000001",
		IfCase:          "use this to get weather information",
		ParentAgentUuid: "00000000-0000-0000-0000-000000000000",
		RouteName:       "weather route app",
	}

	res, resp, err := client.GradientAI.AddAgentRoute(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000001", req)
	if err != nil {
		t.Errorf("GenAI.AddAgentRoute returned error: %v", err)
	}
	fmt.Println(res)
	assert.Equal(t, res.ChildAgentUuid, "00000000-0000-0000-0000-000000000001")
	assert.Equal(t, resp.Response.StatusCode, 200)
}

func TestDeleteAgentRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/child_agents/00000000-0000-0000-0000-000000000001", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, agentRouteResponse)
	})

	res, resp, err := client.GradientAI.DeleteAgentRoute(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000001")
	if err != nil {
		t.Errorf("GenAI.DeleteAgentRoute returned error: %v", err)
	}

	assert.Equal(t, res.ChildAgentUuid, "00000000-0000-0000-0000-000000000001")
	assert.Equal(t, resp.Response.StatusCode, 200)
}

func TestUpdateAgentRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/child_agents/00000000-0000-0000-0000-000000000001", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, agentRouteResponse)
	})

	req := &AgentRouteUpdateRequest{
		ChildAgentUuid:  "00000000-0000-0000-0000-000000000001",
		IfCase:          "use this to get weather information",
		ParentAgentUuid: "00000000-0000-0000-0000-000000000000",
		RouteName:       "weather route app",
	}

	res, resp, err := client.GradientAI.UpdateAgentRoute(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000001", req)
	if err != nil {
		t.Errorf("GenAI.UpdateAgentRoute returned error: %v", err)
	}

	assert.Equal(t, res.ChildAgentUuid, "00000000-0000-0000-0000-000000000001")
	assert.Equal(t, resp.Response.StatusCode, 200)
}

func TestListVersions(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/versions", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAgentVersionsResponse)
	})

	versions, resp, err := client.GradientAI.ListAgentVersions(ctx, "00000000-0000-0000-0000-000000000000", nil)
	if err != nil {
		t.Errorf("GenAI.ListAgentVersions returned error: %v", err)
	}
	fmt.Println(versions)
	assert.Equal(t, 1, len(versions))
	assert.Equal(t, "00000000000000000000000000000000000000000000000000000000000000", versions[0].VersionHash)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestRollbackVersion(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/versions", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, rollbackResponse)
	})

	versions, resp, err := client.GradientAI.RollbackAgentVersion(ctx, "00000000-0000-0000-0000-000000000000", "00000000000000000000000000000000000000000000000000000000000000")
	if err != nil {
		t.Errorf("GenAI.RollbackVersion returned error: %v", err)
	}

	assert.Equal(t, "00000000000000000000000000000000000000000000000000000000000001", versions)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestListAnthropicAPIKeys(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAnthropicAPIKeysResponse)
	})

	keys, resp, err := client.GradientAI.ListAnthropicAPIKeys(ctx, nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, 2, len(keys))
	assert.Equal(t, "Anthropic Key One", keys[0].Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", keys[0].Uuid)
	assert.Equal(t, "Anthropic Key Two", keys[1].Name)
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", keys[1].Uuid)
}

func TestCreateAnthropicAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, anthropicAPIKeyInfoResponse)
	})

	req := &AnthropicAPIKeyCreateRequest{
		Name:   "Anthropic Key One",
		ApiKey: "11111111-1111-1111-1111-111111111111",
	}

	key, resp, err := client.GradientAI.CreateAnthropicAPIKey(ctx, req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "Anthropic Key One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestGetAnthropicAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, anthropicAPIKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.GetAnthropicAPIKey(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "Anthropic Key One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestUpdateAnthropicAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, anthropicAPIKeyInfoResponse)
	})

	req := &AnthropicAPIKeyUpdateRequest{
		Name:       "Anthropic Key One",
		ApiKey:     "11111111-1111-1111-1111-111111111111",
		ApiKeyUuid: "11111111-1111-1111-1111-111111111111",
	}

	key, resp, err := client.GradientAI.UpdateAnthropicAPIKey(ctx, "11111111-1111-1111-1111-111111111111", req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "Anthropic Key One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestDeleteAnthropicAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, anthropicAPIKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.DeleteAnthropicAPIKey(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "Anthropic Key One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestListAgentsByAnthropicAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/anthropic/keys/11111111-1111-1111-1111-111111111111/agents", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAgentsByAnthropicAPIKeyResponse)
	})

	agents, resp, err := client.GradientAI.ListAgentsByAnthropicAPIKey(ctx, "11111111-1111-1111-1111-111111111111", nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, 2, len(agents))
	assert.Equal(t, "Anthropic Agent 1", agents[0].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", agents[0].Uuid)
	assert.Equal(t, "Anthropic Agent 2", agents[1].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000001", agents[1].Uuid)
	assert.NotNil(t, resp.Meta)
	assert.Equal(t, 2, resp.Meta.Total)
	assert.NotNil(t, resp.Links)
	assert.Equal(t, "https://api.digitalocean.com/v2/gen-ai/anthropic/keys?page=1&per_page=1", resp.Links.Pages.First)
}

func TestListOpenAIAPIKeys(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listOpenAIAPIKeysResponse)
	})

	keys, resp, err := client.GradientAI.ListOpenAIAPIKeys(ctx, nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, 2, len(keys))
	assert.Equal(t, "openai-key1", keys[0].Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", keys[0].Uuid)
	assert.Equal(t, "openai-key2", keys[1].Name)
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", keys[1].Uuid)
}

func TestCreateOpenAIAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, openaiAPIKeyInfoResponse)
	})

	req := &OpenAIAPIKeyCreateRequest{
		Name:   "OpenAI One",
		ApiKey: "11111111-1111-1111-1111-111111111111",
	}

	key, resp, err := client.GradientAI.CreateOpenAIAPIKey(ctx, req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "OpenAI One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestGetOpenAIAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, openaiAPIKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.GetOpenAIAPIKey(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "OpenAI One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestUpdateOpenAIAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, openaiAPIKeyInfoResponse)
	})

	req := &OpenAIAPIKeyUpdateRequest{
		Name:       "OpenAI One",
		ApiKey:     "11111111-1111-1111-1111-111111111111",
		ApiKeyUuid: "11111111-1111-1111-1111-111111111111",
	}

	key, resp, err := client.GradientAI.UpdateOpenAIAPIKey(ctx, "11111111-1111-1111-1111-111111111111", req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "OpenAI One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestDeleteOpenAIAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, openaiAPIKeyInfoResponse)
	})

	key, resp, err := client.GradientAI.DeleteOpenAIAPIKey(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, "OpenAI One", key.Name)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", key.Uuid)
}

func TestListAgentsByOpenAIAPIKey(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/openai/keys/11111111-1111-1111-1111-111111111111/agents", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAgentsByOpenAIAPIKeyResponse)
	})

	agents, resp, err := client.GradientAI.ListAgentsByOpenAIAPIKey(ctx, "11111111-1111-1111-1111-111111111111", nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
	assert.Equal(t, 2, len(agents))
	assert.Equal(t, "OpenAI Agent 1", agents[0].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", agents[0].Uuid)
	assert.Equal(t, "OpenAI Agent 2", agents[1].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000001", agents[1].Uuid)
	assert.NotNil(t, resp.Meta)
	assert.Equal(t, 2, resp.Meta.Total)
	assert.NotNil(t, resp.Links)
	assert.Equal(t, "https://api.digitalocean.com/v2/gen-ai/openai/keys?page=1&per_page=1", resp.Links.Pages.First)
}

func TestCreateFunctionRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/functions", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		// Respond with an agent object that includes a non-empty functions array
		fmt.Fprint(w, agentResponse)
	})

	req := &FunctionRouteCreateRequest{
		Description:   "Creating Function Route",
		AgentUuid:     "00000000-0000-0000-0000-000000000000",
		FaasName:      "godo-test-faasname",
		FaasNamespace: "fn-00000000-0000-0000-0000-000000000000",
		InputSchema: FunctionInputSchema{
			Parameters: []OpenAPIParameterSchema{
				{
					Name: "zipCode",
					In:   "query",
					Schema: NestedSchema{
						Type: "string",
					},
					Required:    false,
					Description: "The ZIP code for which to fetch the weather",
				},
				{
					Name: "measurement",
					In:   "query",
					Schema: NestedSchema{
						Type: "string",
						Enum: []string{"F", "C"},
					},
					Required:    false,
					Description: "The measurement unit for temperature (F or C)",
				},
			},
		},
		FunctionName: "godo-test-function",
		OutputSchema: json.RawMessage(`{
  "properties": [
    {
      "name": "temperature",
      "type": "number",
      "description": "The temperature for the specified location"
    },
    {
      "name": "measurement",
      "type": "string",
      "description": "The measurement unit used for the temperature (F or C)"
    },
    {
      "name": "conditions",
      "type": "string",
      "description": "A description of the current weather conditions (Sunny, Cloudy, etc)"
    }
  ]
}`),
	}

	agent, res, err := client.GradientAI.CreateFunctionRoute(ctx, "00000000-0000-0000-0000-000000000000", req)
	if err != nil {
		t.Errorf("GenAI.Create returned error: %v", err)
	}

	assert.NoError(t, err)
	assert.Equal(t, 200, res.Response.StatusCode)
	t.Log(agent)
	assert.Equal(t, req.FunctionName, agent.Functions[0].Name)
}

func TestUpdateFunctionRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/functions/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, agentUpdateResponse)
	})

	req := &FunctionRouteUpdateRequest{
		Description: "Updating Function Route",
		InputSchema: FunctionInputSchema{
			Parameters: []OpenAPIParameterSchema{
				{
					Name: "zipCode",
					In:   "query",
					Schema: NestedSchema{
						Type: "string",
					},
					Required:    false,
					Description: "The ZIP code for which to fetch the weather",
				},
				{
					Name: "measurement",
					In:   "query",
					Schema: NestedSchema{
						Type: "string",
						Enum: []string{"F", "C"},
					},
					Required:    false,
					Description: "The measurement unit for temperature (F or C)",
				},
			},
		},
		OutputSchema: json.RawMessage(`{
            "properties": [
                {
                    "name": "temperature",
                    "type": "number",
                    "description": "The temperature for the specified location"
                }
            ]
        }`),
	}

	agent, resp, err := client.GradientAI.UpdateFunctionRoute(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000000", req)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, req.Description, agent.Functions[0].Description)
}

func TestDeleteFunctionRoute(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/agents/00000000-0000-0000-0000-000000000000/functions/00000000-0000-0000-0000-000000000000", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `{}`)
	})

	_, resp, err := client.GradientAI.DeleteFunctionRoute(ctx, "00000000-0000-0000-0000-000000000000", "00000000-0000-0000-0000-000000000000")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestListInferenceRouters(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "10",
		})
		fmt.Fprint(w, listInferenceRoutersResponse)
	})

	routers, resp, err := client.GradientAI.ListInferenceRouters(ctx, &ListOptions{Page: 1, PerPage: 10})
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.Len(t, routers, 1)
	assert.Equal(t, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", routers[0].UUID)
	assert.Equal(t, "router-one", routers[0].Name)
	assert.Equal(t, "Test router for routing prompts.", routers[0].Description)
	assert.NotNil(t, resp.Meta)
	assert.Equal(t, 1, resp.Meta.Total)
}

func TestListInferenceRouterTaskPresets(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers/tasks/presets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "10",
		})
		fmt.Fprint(w, listInferenceRouterTaskPresetsResponse)
	})

	tasks, resp, err := client.GradientAI.ListInferenceRouterTaskPresets(ctx, &ListOptions{Page: 1, PerPage: 10})
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.Len(t, tasks, 2)
	assert.Equal(t, "translation", tasks[0].TaskSlug)
	assert.Equal(t, "Translation", tasks[0].Name)
	assert.Equal(t, "general", tasks[0].Category)
	assert.Equal(t, "code-generation", tasks[1].TaskSlug)
	assert.Equal(t, "software-engineering", tasks[1].Category)
	require.NotNil(t, resp.Meta)
	assert.Equal(t, 2, resp.Meta.Total)
}

func TestGetInferenceRouter(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers/aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, getInferenceRouterResponse)
	})

	router, resp, err := client.GradientAI.GetInferenceRouter(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.NotNil(t, router)
	assert.Equal(t, "router-one", router.Name)
	assert.Equal(t, "Router used for unit tests.", router.Description)
	require.NotNil(t, router.Config)
	assert.Equal(t, []string{"openai-gpt-4"}, router.Config.FallbackModels)
	assert.Contains(t, string(router.Config.Policies), "code-generation")
}

func TestCreateInferenceRouter(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, getInferenceRouterResponse)
	})

	create := &InferenceRouterCreateRequest{
		Name:           "router-one",
		Description:    "Router used for unit tests.",
		FallbackModels: []string{"openai-gpt-4"},
	}
	router, resp, err := client.GradientAI.CreateInferenceRouter(ctx, create)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.NotNil(t, router)
	assert.Equal(t, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", router.UUID)
}

func TestUpdateInferenceRouter(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers/aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprint(w, getInferenceRouterResponse)
	})

	name := "router-renamed"
	update := &InferenceRouterUpdateRequest{Name: name}
	router, resp, err := client.GradientAI.UpdateInferenceRouter(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", update)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.NotNil(t, router)
	assert.Equal(t, "router-one", router.Name)
}

func TestDeleteInferenceRouter(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models/routers/aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, `{"uuid":"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"}`)
	})

	out, resp, err := client.GradientAI.DeleteInferenceRouter(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	require.NotNil(t, out)
	assert.Equal(t, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", out.UUID)
}

func TestCreateInferenceRouterValidation(t *testing.T) {
	setup()
	defer teardown()

	_, _, err := client.GradientAI.CreateInferenceRouter(ctx, nil)
	assert.Error(t, err)

	_, _, err = client.GradientAI.CreateInferenceRouter(ctx, &InferenceRouterCreateRequest{Name: "n"})
	assert.Error(t, err)

	_, _, err = client.GradientAI.CreateInferenceRouter(ctx, &InferenceRouterCreateRequest{
		Name:           "router",
		FallbackModels: []string{"a", ""},
	})
	assert.Error(t, err)
}

func TestUpdateInferenceRouterValidation(t *testing.T) {
	setup()
	defer teardown()

	_, _, err := client.GradientAI.UpdateInferenceRouter(ctx, "", &InferenceRouterUpdateRequest{Name: "x"})
	assert.Error(t, err)

	_, _, err = client.GradientAI.UpdateInferenceRouter(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", nil)
	assert.Error(t, err)

	_, _, err = client.GradientAI.UpdateInferenceRouter(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", &InferenceRouterUpdateRequest{})
	assert.Error(t, err)
}

func TestListAvailableModels(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "1",
		})

		fmt.Fprint(w, listAvailableModelsResponse)
	})

	req := &ListOptions{
		Page:    1,
		PerPage: 1,
	}

	models, resp, err := client.GradientAI.ListAvailableModels(ctx, req)
	if err != nil {
		t.Fatalf("GenAI ListAvailableModels returned error: %v", err)
	}

	assert.Equal(t, "Llama 3.3 Instruct (70B)", models[0].Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", models[0].Uuid)
	assert.Equal(t, "SERVERLESS", models[0].ModelAvailability)
	assert.Equal(t, "An advanced language model with greater capabilities due to its larger size.", models[0].Description)
	assert.Equal(t, "131072", models[0].ContextWindow)
	assert.Equal(t, []string{"inference", "chat"}, models[0].Capabilities)
	assert.Equal(t, 70.0, models[0].ParameterCount)
	assert.Equal(t, "chat", models[0].Type)
	assert.NotNil(t, models[0].Modalities)
	assert.Equal(t, []string{"text"}, models[0].Modalities.Input)
	assert.Equal(t, []string{"text"}, models[0].Modalities.Output)
	expectedString := fmt.Sprintf("%v", models[0])
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, expectedString, models[0].String())
}

func TestSearchModels(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAvailableModelsResponse)
	})

	// Test matching query
	uuids, resp, err := client.GradientAI.SearchModels(ctx, "llama")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, 1, len(uuids))
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", uuids[0])
}

func TestSearchModelsNoMatch(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAvailableModelsResponse)
	})

	// Test non-matching query
	uuids, resp, err := client.GradientAI.SearchModels(ctx, "nonexistent")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, 0, len(uuids))
}

func TestGetModelByUUID(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAvailableModelsResponse)
	})

	// Test existing UUID
	model, resp, err := client.GradientAI.GetModelByUUID(ctx, "00000000-0000-0000-0000-000000000000")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.NotNil(t, model)
	assert.Equal(t, "Llama 3.3 Instruct (70B)", model.Name)
	assert.Equal(t, "00000000-0000-0000-0000-000000000000", model.Uuid)
}

func TestGetModelByUUIDNotFound(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listAvailableModelsResponse)
	})

	// Test non-existing UUID
	model, resp, err := client.GradientAI.GetModelByUUID(ctx, "99999999-9999-9999-9999-999999999999")
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Nil(t, model)
}

func TestListDatacenterRegions(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/regions", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listDatacenterRegionsResponse)
	})

	regions, resp, err := client.GradientAI.ListDatacenterRegions(ctx, nil, nil)
	if err != nil {
		t.Fatalf("GenAI ListDatacenterRegions returned error: %v", err)
	}

	assert.Equal(t, 2, len(regions))
	assert.Equal(t, "tor1", regions[0].Region)
	assert.Equal(t, "https://tor1.gen-ai.digitalocean.com", regions[0].InferenceUrl)
	assert.Equal(t, true, regions[0].ServesBatch)
	assert.Equal(t, true, regions[0].ServesInference)
	assert.Equal(t, "https://tor1.gen-ai.digitalocean.com/stream", regions[0].StreamInferenceUrl)

	assert.Equal(t, "nyc3", regions[1].Region)
	assert.Equal(t, "https://nyc3.gen-ai.digitalocean.com", regions[1].InferenceUrl)
	assert.Equal(t, false, regions[1].ServesBatch)
	assert.Equal(t, true, regions[1].ServesInference)
	assert.Equal(t, "https://nyc3.gen-ai.digitalocean.com/stream", regions[1].StreamInferenceUrl)

	assert.Equal(t, 200, resp.Response.StatusCode)
}

var listCustomModelsResponse = `
{
	"models": [
		{
			"uuid": "11111111-1111-1111-1111-111111111111",
			"name": "team/my-model",
			"description": "An imported HuggingFace model",
			"status": "STATUS_READY",
			"architecture": "LlamaForCausalLM",
			"source_type": "SOURCE_TYPE_HUGGINGFACE",
			"source_ref": {
				"repo_id": "team/my-model",
				"commit_sha": "abc123",
				"access_type": "ACCESS_TYPE_PUBLIC"
			},
			"total_size_bytes": "1073741824",
			"file_count": 7,
			"license": "apache-2.0",
			"tags": {"tags": ["llama", "chat"]},
			"created_at": "2026-01-02T03:04:05Z",
			"updated_at": "2026-01-02T04:05:06Z",
			"active_deployments": [
				{
					"id": "dep-1",
					"name": "primary",
					"region_slug": "tor1",
					"state": "ACTIVE",
					"endpoints": {
						"public_endpoint_fqdn": "public.example.com",
						"private_endpoint_fqdn": "private.example.com"
					},
					"created_at": "2026-01-02T03:04:05Z",
					"updated_at": "2026-01-02T04:05:06Z"
				}
			],
			"context_length": 131072,
			"cost_estimate_per_month": 123,
			"input_modalities": ["text"],
			"output_modalities": ["text"],
			"parameters": "70000000000",
			"team_id": "42",
			"config_json": {"architectures": ["LlamaForCausalLM"]},
			"storage_region": "nyc3"
		},
		{
			"uuid": "22222222-2222-2222-2222-222222222222",
			"name": "team/other-model",
			"status": "STATUS_IMPORTING",
			"source_type": "SOURCE_TYPE_SPACES_BUCKET",
			"source_ref": {
				"bucket": "my-bucket",
				"region": "nyc3",
				"prefix": "models/other"
			}
		}
	],
	"links": {
		"pages": {
			"first": "https://api.digitalocean.com/v2/gen-ai/custom_models?page=1&per_page=2",
			"next": "https://api.digitalocean.com/v2/gen-ai/custom_models?page=2&per_page=2",
			"last": "https://api.digitalocean.com/v2/gen-ai/custom_models?page=5&per_page=2"
		}
	},
	"meta": {
		"total": 10,
		"page": 1,
		"pages": 5
	},
	"max_threshold": 25
}
`

var getCustomModelResponse = `
{
	"model": {
		"uuid": "11111111-1111-1111-1111-111111111111",
		"name": "team/my-model",
		"description": "An imported HuggingFace model",
		"status": "STATUS_READY",
		"architecture": "LlamaForCausalLM",
		"source_type": "SOURCE_TYPE_HUGGINGFACE",
		"source_ref": {
			"repo_id": "team/my-model",
			"commit_sha": "abc123",
			"access_type": "ACCESS_TYPE_PUBLIC"
		},
		"total_size_bytes": "1073741824",
		"file_count": 7,
		"license": "apache-2.0",
		"tags": {"tags": ["llama", "chat"]},
		"created_at": "2026-01-02T03:04:05Z",
		"updated_at": "2026-01-02T04:05:06Z",
		"active_deployments": [
			{
				"id": "dep-1",
				"name": "primary",
				"region_slug": "tor1",
				"state": "ACTIVE",
				"endpoints": {
					"public_endpoint_fqdn": "public.example.com",
					"private_endpoint_fqdn": "private.example.com"
				},
				"created_at": "2026-01-02T03:04:05Z",
				"updated_at": "2026-01-02T04:05:06Z"
			}
		],
		"context_length": 131072,
		"cost_estimate_per_month": 123,
		"input_modalities": ["text"],
		"output_modalities": ["text"],
		"parameters": "70000000000",
		"team_id": "42",
		"config_json": {"architectures": ["LlamaForCausalLM"]},
		"storage_region": "nyc3"
	}
}
`

var getFailedCustomModelResponse = `
{
	"model": {
		"uuid": "55555555-5555-5555-5555-555555555555",
		"name": "team/failed-model",
		"status": "STATUS_FAILED",
		"error_message": "Repository not found or access denied",
		"source_type": "SOURCE_TYPE_HUGGINGFACE",
		"source_ref": {
			"repo_id": "team/failed-model",
			"access_type": "ACCESS_TYPE_PRIVATE"
		}
	}
}
`

var importCustomModelResponse = `
{
	"model": {
		"uuid": "33333333-3333-3333-3333-333333333333",
		"name": "team/new-model",
		"status": "STATUS_IMPORTING",
		"source_type": "SOURCE_TYPE_HUGGINGFACE",
		"source_ref": {
			"repo_id": "team/new-model",
			"access_type": "ACCESS_TYPE_PUBLIC"
		},
		"tags": {"tags": ["fast", "small"]}
	},
	"import_job": {
		"uuid": "44444444-4444-4444-4444-444444444444",
		"status": "RUNNING",
		"files_total": 7,
		"files_done": 0,
		"bytes_total": "1073741824",
		"bytes_done": "0",
		"started_at": "2026-01-02T03:04:05Z",
		"created_at": "2026-01-02T03:04:05Z"
	},
	"validation_steps": [
		{"name": "config-check", "passed": true},
		{"name": "license-check", "passed": false, "error": "no license file"}
	],
	"error": ""
}
`

var deleteCustomModelResponse = `
{
	"status": "DELETE_CUSTOM_MODEL_STATUS_SUCCESS",
	"error": ""
}
`

var updateCustomModelMetadataResponse = `
{
	"model": {
		"uuid": "11111111-1111-1111-1111-111111111111",
		"name": "team/my-model",
		"description": "Updated description",
		"status": "STATUS_READY",
		"source_type": "SOURCE_TYPE_HUGGINGFACE",
		"tags": {"tags": ["updated", "v2"]}
	}
}
`

var updateCustomModelMetadataSpacesResponse = `
{
	"model": {
		"uuid": "22222222-2222-2222-2222-222222222222",
		"name": "team/spaces-model",
		"description": "Spaces import",
		"status": "STATUS_READY",
		"source_type": "SOURCE_TYPE_SPACES_BUCKET",
		"input_modalities": ["text", "image"],
		"output_modalities": ["text"],
		"parameters": "7000000000",
		"license": "apache-2.0"
	}
}
`

var customEvaluationMetricResponse = `
{
	"metric": {
		"metric_uuid": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
		"metric_name": "My domain tone metric",
		"description": "Scores adherence to our support macros",
		"source": "EVALUATION_METRIC_SOURCE_CUSTOM",
		"custom_eval_config": {
			"requires_ground_truth": true,
			"scoring_prompt": "Grade for brand tone and factuality."
		},
		"associated_presets": [
			{
				"eval_preset_uuid": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
				"name": "Support preset"
			}
		]
	}
}
`

func TestListCustomModels(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testFormValues(t, r, values{
			"page":     "1",
			"per_page": "2",
			"status":   "STATUS_READY",
		})
		fmt.Fprint(w, listCustomModelsResponse)
	})

	opt := &CustomModelListOptions{
		Status: CustomModelStatusReady,
		ListOptions: ListOptions{
			Page:    1,
			PerPage: 2,
		},
	}

	out, resp, err := client.GradientAI.ListCustomModels(ctx, opt)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)

	assert.Equal(t, 2, len(out.Models))
	assert.Equal(t, 25, out.MaxThreshold)
	assert.Equal(t, 10, resp.Meta.Total)
	assert.Equal(t, 1, resp.Meta.Page)
	assert.Equal(t, 5, resp.Meta.Pages)
	assert.NotNil(t, resp.Links)

	first := out.Models[0]
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", first.Uuid)
	assert.Equal(t, "team/my-model", first.Name)
	assert.Equal(t, CustomModelStatusReady, first.Status)
	assert.Equal(t, CustomModelSourceTypeHuggingFace, first.SourceType)
	assert.NotNil(t, first.SourceRef)
	assert.Equal(t, "team/my-model", first.SourceRef.RepoId)
	assert.Equal(t, CustomModelSourceRefAccessTypePublic, first.SourceRef.AccessType)
	assert.Equal(t, "1073741824", first.TotalSizeBytes)
	assert.Equal(t, 7, first.FileCount)
	assert.NotNil(t, first.Tags)
	assert.Equal(t, []string{"llama", "chat"}, first.Tags.Tags)
	assert.Equal(t, 131072, first.ContextLength)
	assert.Equal(t, 123, first.CostEstimatePerMonth)
	assert.Equal(t, "70000000000", first.Parameters)
	assert.Equal(t, "42", first.TeamId)
	assert.Equal(t, "nyc3", first.StorageRegion)
	assert.NotNil(t, first.ConfigJson)
	assert.Equal(t, []string{"text"}, first.InputModalities)

	assert.Equal(t, 1, len(first.ActiveDeployments))
	dep := first.ActiveDeployments[0]
	assert.Equal(t, "dep-1", dep.Id)
	assert.Equal(t, "primary", dep.Name)
	assert.Equal(t, "tor1", dep.RegionSlug)
	assert.Equal(t, "ACTIVE", dep.State)
	assert.NotNil(t, dep.Endpoints)
	assert.Equal(t, "public.example.com", dep.Endpoints.PublicEndpointFqdn)
	assert.Equal(t, "private.example.com", dep.Endpoints.PrivateEndpointFqdn)

	second := out.Models[1]
	assert.Equal(t, CustomModelStatusImporting, second.Status)
	assert.Equal(t, CustomModelSourceTypeSpacesBucket, second.SourceType)
	assert.NotNil(t, second.SourceRef)
	assert.Equal(t, "my-bucket", second.SourceRef.Bucket)
	assert.Equal(t, "nyc3", second.SourceRef.Region)
	assert.Equal(t, "models/other", second.SourceRef.Prefix)
}

func TestListCustomModelsNoOptions(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		assert.Empty(t, r.URL.Query().Get("status"))
		assert.Empty(t, r.URL.Query().Get("page"))
		fmt.Fprint(w, listCustomModelsResponse)
	})

	out, resp, err := client.GradientAI.ListCustomModels(ctx, nil)
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, 2, len(out.Models))
}

func TestGetCustomModel(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, getCustomModelResponse)
	})

	model, resp, err := client.GradientAI.GetCustomModel(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.NotNil(t, model)
	assert.Equal(t, 200, resp.Response.StatusCode)

	assert.Equal(t, "11111111-1111-1111-1111-111111111111", model.Uuid)
	assert.Equal(t, "team/my-model", model.Name)
	assert.Equal(t, "An imported HuggingFace model", model.Description)
	assert.Equal(t, CustomModelStatusReady, model.Status)
	assert.Equal(t, "LlamaForCausalLM", model.Architecture)
	assert.Equal(t, CustomModelSourceTypeHuggingFace, model.SourceType)

	assert.NotNil(t, model.SourceRef)
	assert.Equal(t, "team/my-model", model.SourceRef.RepoId)
	assert.Equal(t, "abc123", model.SourceRef.CommitSha)
	assert.Equal(t, CustomModelSourceRefAccessTypePublic, model.SourceRef.AccessType)

	assert.Equal(t, "1073741824", model.TotalSizeBytes)
	assert.Equal(t, 7, model.FileCount)
	assert.Equal(t, "apache-2.0", model.License)
	assert.NotNil(t, model.Tags)
	assert.Equal(t, []string{"llama", "chat"}, model.Tags.Tags)
	assert.Equal(t, 131072, model.ContextLength)
	assert.Equal(t, 123, model.CostEstimatePerMonth)
	assert.Equal(t, "70000000000", model.Parameters)
	assert.Equal(t, "42", model.TeamId)
	assert.Equal(t, "nyc3", model.StorageRegion)
	assert.NotNil(t, model.ConfigJson)
	assert.Equal(t, []string{"text"}, model.InputModalities)
	assert.Equal(t, []string{"text"}, model.OutputModalities)

	assert.Equal(t, 1, len(model.ActiveDeployments))
	dep := model.ActiveDeployments[0]
	assert.Equal(t, "dep-1", dep.Id)
	assert.Equal(t, "primary", dep.Name)
	assert.Equal(t, "tor1", dep.RegionSlug)
	assert.Equal(t, "ACTIVE", dep.State)
	assert.NotNil(t, dep.Endpoints)
	assert.Equal(t, "public.example.com", dep.Endpoints.PublicEndpointFqdn)
	assert.Equal(t, "private.example.com", dep.Endpoints.PrivateEndpointFqdn)
}

func TestGetCustomModelFailed(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/55555555-5555-5555-5555-555555555555", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, getFailedCustomModelResponse)
	})

	model, resp, err := client.GradientAI.GetCustomModel(ctx, "55555555-5555-5555-5555-555555555555")
	assert.NoError(t, err)
	assert.NotNil(t, model)
	assert.Equal(t, 200, resp.Response.StatusCode)

	assert.Equal(t, "55555555-5555-5555-5555-555555555555", model.Uuid)
	assert.Equal(t, "team/failed-model", model.Name)
	assert.Equal(t, CustomModelStatusFailed, model.Status)
	assert.Equal(t, "Repository not found or access denied", model.ErrorMessage)
	assert.Equal(t, CustomModelSourceTypeHuggingFace, model.SourceType)
	assert.NotNil(t, model.SourceRef)
	assert.Equal(t, "team/failed-model", model.SourceRef.RepoId)
	assert.Equal(t, CustomModelSourceRefAccessTypePrivate, model.SourceRef.AccessType)
}

func TestGetCustomModelMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	model, resp, err := client.GradientAI.GetCustomModel(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.Nil(t, resp)
}

func TestGetCustomModelInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	model, resp, err := client.GradientAI.GetCustomModel(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.Nil(t, resp)
}

func TestGetCustomModelServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"not_found","message":"missing"}`, http.StatusNotFound)
	})

	model, resp, err := client.GradientAI.GetCustomModel(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestImportCustomModel(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/import", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)
		var got CustomModelImportRequest
		assert.NoError(t, json.Unmarshal(body, &got))
		assert.Equal(t, "team/new-model", got.Name)
		assert.Equal(t, CustomModelSourceTypeHuggingFace, got.SourceType)
		assert.NotNil(t, got.SourceRef)
		assert.Equal(t, "team/new-model", got.SourceRef.RepoId)
		assert.Equal(t, "tor1", got.PreferredGpuRegion)
		assert.True(t, got.AcceptTermsAndConditions)
		assert.NotNil(t, got.Tags)
		assert.Equal(t, []string{"fast", "small"}, got.Tags.Tags)

		fmt.Fprint(w, importCustomModelResponse)
	})

	req := &CustomModelImportRequest{
		Name:       "team/new-model",
		SourceType: CustomModelSourceTypeHuggingFace,
		SourceRef: &CustomModelSourceRef{
			RepoId:     "team/new-model",
			AccessType: CustomModelSourceRefAccessTypePublic,
		},
		Description:              "an imported model",
		PreferredGpuRegion:       "tor1",
		AcceptTermsAndConditions: true,
		Tags:                     &CustomModelTags{Tags: []string{"fast", "small"}},
	}

	out, resp, err := client.GradientAI.ImportCustomModel(ctx, req)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)

	assert.NotNil(t, out.Model)
	assert.Equal(t, "33333333-3333-3333-3333-333333333333", out.Model.Uuid)
	assert.Equal(t, CustomModelStatusImporting, out.Model.Status)

	assert.NotNil(t, out.ImportJob)
	assert.Equal(t, "44444444-4444-4444-4444-444444444444", out.ImportJob.Uuid)
	assert.Equal(t, "RUNNING", out.ImportJob.Status)
	assert.Equal(t, 7, out.ImportJob.FilesTotal)
	assert.Equal(t, 0, out.ImportJob.FilesDone)
	assert.Equal(t, "1073741824", out.ImportJob.BytesTotal)
	assert.Equal(t, "0", out.ImportJob.BytesDone)

	assert.Equal(t, 2, len(out.ValidationSteps))
	assert.Equal(t, "config-check", out.ValidationSteps[0].Name)
	assert.True(t, out.ValidationSteps[0].Passed)
	assert.Equal(t, "license-check", out.ValidationSteps[1].Name)
	assert.False(t, out.ValidationSteps[1].Passed)
	assert.Equal(t, "no license file", out.ValidationSteps[1].Error)
}

func TestImportCustomModelNilRequest(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.ImportCustomModel(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestImportCustomModelMissingName(t *testing.T) {
	setup()
	defer teardown()

	req := &CustomModelImportRequest{
		SourceType: CustomModelSourceTypeHuggingFace,
	}

	out, resp, err := client.GradientAI.ImportCustomModel(ctx, req)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestImportCustomModelMissingSourceType(t *testing.T) {
	setup()
	defer teardown()

	req := &CustomModelImportRequest{
		Name: "team/new-model",
	}

	out, resp, err := client.GradientAI.ImportCustomModel(ctx, req)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteCustomModel(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/11111111-1111-1111-1111-111111111111", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, deleteCustomModelResponse)
	})

	out, resp, err := client.GradientAI.DeleteCustomModel(ctx, "11111111-1111-1111-1111-111111111111")
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, DeleteCustomModelStatusSuccess, out.Status)
	assert.Empty(t, out.Error)
}

func TestDeleteCustomModelMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteCustomModel(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteCustomModelServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		http.Error(w, `{"id":"not_found","message":"model not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.DeleteCustomModel(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestUpdateCustomModelMetadata(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/11111111-1111-1111-1111-111111111111/metadata", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)
		var got CustomModelMetadataUpdateRequest
		assert.NoError(t, json.Unmarshal(body, &got))
		assert.Equal(t, "Updated description", got.Description)
		assert.NotNil(t, got.Tags)
		assert.Equal(t, []string{"updated", "v2"}, got.Tags.Tags)

		fmt.Fprint(w, updateCustomModelMetadataResponse)
	})

	req := &CustomModelMetadataUpdateRequest{
		Description: "Updated description",
		Tags:        &CustomModelTags{Tags: []string{"updated", "v2"}},
	}

	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "11111111-1111-1111-1111-111111111111", req)
	assert.NoError(t, err)
	assert.NotNil(t, model)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", model.Uuid)
	assert.Equal(t, "Updated description", model.Description)
	assert.NotNil(t, model.Tags)
	assert.Equal(t, []string{"updated", "v2"}, model.Tags.Tags)
}

func TestUpdateCustomModelMetadataSpacesFields(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/22222222-2222-2222-2222-222222222222/metadata", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)
		var got CustomModelMetadataUpdateRequest
		assert.NoError(t, json.Unmarshal(body, &got))
		assert.Equal(t, []string{"text", "image"}, got.InputModalities)
		assert.Equal(t, []string{"text"}, got.OutputModalities)
		assert.Equal(t, "7000000000", got.Parameters)
		assert.Equal(t, "apache-2.0", got.License)

		fmt.Fprint(w, updateCustomModelMetadataSpacesResponse)
	})

	req := &CustomModelMetadataUpdateRequest{
		InputModalities:  []string{"text", "image"},
		OutputModalities: []string{"text"},
		Parameters:       "7000000000",
		License:          "apache-2.0",
	}

	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "22222222-2222-2222-2222-222222222222", req)
	assert.NoError(t, err)
	assert.NotNil(t, model)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "22222222-2222-2222-2222-222222222222", model.Uuid)
	assert.Equal(t, []string{"text", "image"}, model.InputModalities)
	assert.Equal(t, []string{"text"}, model.OutputModalities)
	assert.Equal(t, "7000000000", model.Parameters)
	assert.Equal(t, "apache-2.0", model.License)
}

func TestUpdateCustomModelMetadataOmitsUnsetSpacesFields(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/11111111-1111-1111-1111-111111111111/metadata", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)

		var raw map[string]interface{}
		assert.NoError(t, json.Unmarshal(body, &raw))
		_, hasInput := raw["input_modalities"]
		_, hasOutput := raw["output_modalities"]
		_, hasParams := raw["parameters"]
		_, hasLicense := raw["license"]
		assert.False(t, hasInput, "input_modalities should be omitted when unset")
		assert.False(t, hasOutput, "output_modalities should be omitted when unset")
		assert.False(t, hasParams, "parameters should be omitted when unset")
		assert.False(t, hasLicense, "license should be omitted when unset")

		fmt.Fprint(w, updateCustomModelMetadataResponse)
	})

	req := &CustomModelMetadataUpdateRequest{Description: "Updated description"}
	_, _, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "11111111-1111-1111-1111-111111111111", req)
	assert.NoError(t, err)
}

func TestUpdateCustomModelMetadataMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	req := &CustomModelMetadataUpdateRequest{Description: "x"}

	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "", req)
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.Nil(t, resp)
}

func TestUpdateCustomModelMetadataNilRequest(t *testing.T) {
	setup()
	defer teardown()

	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "11111111-1111-1111-1111-111111111111", nil)
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.Nil(t, resp)
}

func TestListCustomModelsServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"server_error","message":"boom"}`, http.StatusInternalServerError)
	})

	out, resp, err := client.GradientAI.ListCustomModels(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusInternalServerError, resp.Response.StatusCode)
}

func TestImportCustomModelServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/import", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		http.Error(w, `{"id":"invalid","message":"bad source"}`, http.StatusBadRequest)
	})

	req := &CustomModelImportRequest{
		Name:       "team/new-model",
		SourceType: CustomModelSourceTypeHuggingFace,
	}
	out, resp, err := client.GradientAI.ImportCustomModel(ctx, req)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusBadRequest, resp.Response.StatusCode)
}

func TestUpdateCustomModelMetadataServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_models/11111111-1111-1111-1111-111111111111/metadata", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)
		http.Error(w, `{"id":"not_found","message":"missing"}`, http.StatusNotFound)
	})

	req := &CustomModelMetadataUpdateRequest{Description: "x"}
	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "11111111-1111-1111-1111-111111111111", req)
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestDeleteCustomModelInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteCustomModel(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestUpdateCustomModelMetadataInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	model, resp, err := client.GradientAI.UpdateCustomModelMetadata(ctx, "bad\nuuid", &CustomModelMetadataUpdateRequest{Description: "x"})
	assert.Error(t, err)
	assert.Nil(t, model)
	assert.Nil(t, resp)
}

func TestCreateCustomEvaluationMetric(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_evaluation_metrics", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)
		var got CreateCustomEvaluationMetricRequest
		assert.NoError(t, json.Unmarshal(body, &got))
		assert.Equal(t, "My domain tone metric", got.MetricName)
		assert.Equal(t, "Scores adherence to our support macros", got.Description)
		assert.NotNil(t, got.Config)
		assert.True(t, got.Config.RequiresGroundTruth)
		assert.Equal(t, "Grade for brand tone and factuality.", got.Config.ScoringPrompt)

		fmt.Fprint(w, customEvaluationMetricResponse)
	})

	req := &CreateCustomEvaluationMetricRequest{
		MetricName:  "My domain tone metric",
		Description: "Scores adherence to our support macros",
		Config: &CustomEvaluationMetricConfig{
			RequiresGroundTruth: true,
			ScoringPrompt:       "Grade for brand tone and factuality.",
		},
	}

	metric, resp, err := client.GradientAI.CreateCustomEvaluationMetric(ctx, req)
	assert.NoError(t, err)
	assert.NotNil(t, metric)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", metric.MetricUUID)
	assert.Equal(t, "My domain tone metric", metric.MetricName)
	assert.Equal(t, "Scores adherence to our support macros", metric.Description)
	assert.Equal(t, EvaluationMetricSourceCustom, metric.Source)
	assert.NotNil(t, metric.CustomEvalConfig)
	assert.Equal(t, "Grade for brand tone and factuality.", metric.CustomEvalConfig.ScoringPrompt)
	assert.Equal(t, 1, len(metric.AssociatedPresets))
	assert.Equal(t, "Support preset", metric.AssociatedPresets[0].Name)
}

func TestCreateCustomEvaluationMetricNilRequest(t *testing.T) {
	setup()
	defer teardown()

	metric, resp, err := client.GradientAI.CreateCustomEvaluationMetric(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, metric)
	assert.Nil(t, resp)
}

func TestUpdateCustomEvaluationMetric(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_evaluation_metrics/aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)

		body, err := io.ReadAll(r.Body)
		assert.NoError(t, err)
		var got UpdateCustomEvaluationMetricRequest
		assert.NoError(t, json.Unmarshal(body, &got))
		assert.Equal(t, "My domain tone metric", got.MetricName)
		assert.Equal(t, "Scores adherence to our support macros", got.Description)
		assert.NotNil(t, got.Config)
		assert.True(t, got.Config.RequiresGroundTruth)

		fmt.Fprint(w, customEvaluationMetricResponse)
	})

	req := &UpdateCustomEvaluationMetricRequest{
		MetricUUID:  "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
		MetricName:  "My domain tone metric",
		Description: "Scores adherence to our support macros",
		Config: &CustomEvaluationMetricConfig{
			RequiresGroundTruth: true,
			ScoringPrompt:       "Grade for brand tone and factuality.",
		},
	}

	metric, resp, err := client.GradientAI.UpdateCustomEvaluationMetric(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", req)
	assert.NoError(t, err)
	assert.NotNil(t, metric)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", metric.MetricUUID)
	assert.NotNil(t, metric.CustomEvalConfig)
	assert.True(t, metric.CustomEvalConfig.RequiresGroundTruth)
}

func TestDeleteCustomEvaluationMetric(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/custom_evaluation_metrics/aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `{}`)
	})

	resp, err := client.GradientAI.DeleteCustomEvaluationMetric(ctx, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
	assert.NoError(t, err)
	assert.NotNil(t, resp)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestDeleteCustomEvaluationMetricMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	resp, err := client.GradientAI.DeleteCustomEvaluationMetric(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, resp)
}

var deleteModelEvaluationRunResponse = `
{
	"status": "DELETE_MODEL_EVALUATION_RUN_STATUS_SUCCESS",
	"error": ""
}
`

func TestDeleteModelEvaluationRun(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/12345678-1234-1234-1234-123456789012", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, deleteModelEvaluationRunResponse)
	})

	out, resp, err := client.GradientAI.DeleteModelEvaluationRun(ctx, "12345678-1234-1234-1234-123456789012")
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, DeleteModelEvaluationRunStatusSuccess, out.Status)
	assert.Empty(t, out.Error)
}

func TestDeleteModelEvaluationRunMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteModelEvaluationRun(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteModelEvaluationRunServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		http.Error(w, `{"id":"not_found","message":"evaluation run not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.DeleteModelEvaluationRun(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestDeleteModelEvaluationRunInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteModelEvaluationRun(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteModelEvaluationPreset(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_presets/12345678-1234-1234-1234-123456789012", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, `{}`)
	})

	out, resp, err := client.GradientAI.DeleteModelEvaluationPreset(ctx, "12345678-1234-1234-1234-123456789012")
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestDeleteModelEvaluationPresetMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteModelEvaluationPreset(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteModelEvaluationPresetServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_presets/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		http.Error(w, `{"id":"not_found","message":"evaluation preset not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.DeleteModelEvaluationPreset(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestDeleteModelEvaluationPresetInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteModelEvaluationPreset(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

var cancelModelEvaluationRunResponse = `
{
	"run": {
		"eval_run_uuid": "12345678-1234-1234-1234-123456789012",
		"name": "my-eval-run",
		"status": "MODEL_EVALUATION_RUN_CANCELLING",
		"candidate_model_uuid": "00000000-0000-0000-0000-000000000001",
		"candidate_model_name": "candidate-model",
		"candidate_model_source": "CANDIDATE_MODEL_SOURCE_SERVERLESS",
		"judge_model_uuid": "00000000-0000-0000-0000-000000000002",
		"judge_model_name": "judge-model",
		"dataset_uuid": "00000000-0000-0000-0000-000000000003",
		"dataset_name": "eval-dataset",
		"created_at": "2025-05-08T03:37:28Z"
	}
}
`

func TestCancelModelEvaluationRun(t *testing.T) {
	setup()
	defer teardown()

	evalRunUUID := "12345678-1234-1234-1234-123456789012"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_runs/%s/cancel", evalRunUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)

		v := new(CancelModelEvaluationRunRequest)
		if err := json.NewDecoder(r.Body).Decode(v); err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}
		assert.Equal(t, evalRunUUID, v.EvalRunUUID)

		fmt.Fprint(w, cancelModelEvaluationRunResponse)
	})

	out, resp, err := client.GradientAI.CancelModelEvaluationRun(ctx, evalRunUUID)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.NotNil(t, out.Run)
	assert.Equal(t, 200, resp.Response.StatusCode)

	run := out.Run
	assert.Equal(t, evalRunUUID, run.EvalRunUuid)
	assert.Equal(t, "my-eval-run", run.Name)
	assert.Equal(t, ModelEvaluationRunCancelling, run.Status)
	assert.Equal(t, "candidate-model", run.CandidateModelName)
	assert.Equal(t, CandidateModelSourceServerless, run.CandidateModelSource)
	assert.Equal(t, "judge-model", run.JudgeModelName)
	assert.Equal(t, "eval-dataset", run.DatasetName)
	assert.NotNil(t, run.CreatedAt)
}

func TestCancelModelEvaluationRunMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.CancelModelEvaluationRun(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestCancelModelEvaluationRunServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/99999999-9999-9999-9999-999999999999/cancel", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		http.Error(w, `{"id":"not_found","message":"evaluation run not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.CancelModelEvaluationRun(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestCancelModelEvaluationRunInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.CancelModelEvaluationRun(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

var updateModelEvaluationRunResponse = `
{
	"run": {
		"eval_run_uuid": "33c32a9b-aa48-4577-b513-0a0257dd00f2",
		"name": "Updated run",
		"candidate_model_name": "OpenAI GPT-oss-20b",
		"dataset_name": "political_reasoning_safety_NGT.jsonl",
		"status": "MODEL_EVALUATION_RUN_PARTIALLY_SUCCESSFUL",
		"created_at": "2026-06-05T12:17:08Z",
		"candidate_model_uuid": "c60da0be-73c4-11f0-b074-4e013e2ddde4",
		"judge_model_uuid": "3568b895-f2ca-11ef-bf8f-4e013e2ddde4",
		"judge_model_name": "OpenAI GPT-4o",
		"dataset_uuid": "11f15429-ff86-d769-9439-ca68c578b04b"
	}
}
`

func TestUpdateModelEvaluationRun(t *testing.T) {
	setup()
	defer teardown()

	evalRunUUID := "33c32a9b-aa48-4577-b513-0a0257dd00f2"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_runs/%s", evalRunUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)

		v := new(UpdateModelEvaluationRunRequest)
		if err := json.NewDecoder(r.Body).Decode(v); err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}
		assert.Equal(t, "Updated run", v.Name)

		fmt.Fprint(w, updateModelEvaluationRunResponse)
	})

	updateReq := &UpdateModelEvaluationRunRequest{
		Name: "Updated run",
	}

	out, resp, err := client.GradientAI.UpdateModelEvaluationRun(ctx, evalRunUUID, updateReq)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.NotNil(t, out.Run)
	assert.Equal(t, 200, resp.Response.StatusCode)

	run := out.Run
	assert.Equal(t, evalRunUUID, run.EvalRunUuid)
	assert.Equal(t, "Updated run", run.Name)
	assert.Equal(t, ModelEvaluationRunPartiallySuccessful, run.Status)
	assert.Equal(t, "OpenAI GPT-oss-20b", run.CandidateModelName)
	assert.Equal(t, "c60da0be-73c4-11f0-b074-4e013e2ddde4", run.CandidateModelUuid)
	assert.Equal(t, "OpenAI GPT-4o", run.JudgeModelName)
	assert.Equal(t, "3568b895-f2ca-11ef-bf8f-4e013e2ddde4", run.JudgeModelUuid)
	assert.Equal(t, "political_reasoning_safety_NGT.jsonl", run.DatasetName)
	assert.Equal(t, "11f15429-ff86-d769-9439-ca68c578b04b", run.DatasetUuid)
	assert.NotNil(t, run.CreatedAt)
}

func TestUpdateModelEvaluationRunMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.UpdateModelEvaluationRun(ctx, "", &UpdateModelEvaluationRunRequest{Name: "x"})
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestUpdateModelEvaluationRunMissingRequest(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.UpdateModelEvaluationRun(ctx, "12345678-1234-1234-1234-123456789012", nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestUpdateModelEvaluationRunServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPatch)
		http.Error(w, `{"id":"not_found","message":"evaluation run not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.UpdateModelEvaluationRun(ctx, "99999999-9999-9999-9999-999999999999", &UpdateModelEvaluationRunRequest{Name: "x"})
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestUpdateModelEvaluationRunInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.UpdateModelEvaluationRun(ctx, "bad\nuuid", &UpdateModelEvaluationRunRequest{Name: "x"})
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

var createModelEvaluationRunResponse = `
{
	"eval_run_uuid": "12345678-1234-1234-1234-123456789012"
}
`

func TestCreateModelEvaluationRun(t *testing.T) {
	setup()
	defer teardown()

	createReq := &CreateModelEvaluationRunRequest{
		Name:               "my-evaluation-run",
		CandidateModelUUID: "11111111-1111-1111-1111-111111111111",
		JudgeModelUUID:     "22222222-2222-2222-2222-222222222222",
		DatasetUUID:        "33333333-3333-3333-3333-333333333333",
		MetricUUIDs:        []string{"44444444-4444-4444-4444-444444444444"},
	}

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)

		v := new(CreateModelEvaluationRunRequest)
		if err := json.NewDecoder(r.Body).Decode(v); err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}
		assert.Equal(t, createReq.Name, v.Name)
		assert.Equal(t, createReq.CandidateModelUUID, v.CandidateModelUUID)
		assert.Equal(t, createReq.JudgeModelUUID, v.JudgeModelUUID)
		assert.Equal(t, createReq.DatasetUUID, v.DatasetUUID)
		assert.Equal(t, createReq.MetricUUIDs, v.MetricUUIDs)

		fmt.Fprint(w, createModelEvaluationRunResponse)
	})

	out, resp, err := client.GradientAI.CreateModelEvaluationRun(ctx, createReq)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "12345678-1234-1234-1234-123456789012", out.EvalRunUuid)
}

func TestCreateModelEvaluationRunNilRequest(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.CreateModelEvaluationRun(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestCreateModelEvaluationRunServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		http.Error(w, `{"id":"server_error","message":"boom"}`, http.StatusInternalServerError)
	})

	out, resp, err := client.GradientAI.CreateModelEvaluationRun(ctx, &CreateModelEvaluationRunRequest{Name: "x"})
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusInternalServerError, resp.Response.StatusCode)
}

var createModelEvalDatasetUploadPresignedURLsResponse = `
{
	"request_id": "abcdef01-2345-6789-abcd-ef0123456789",
	"uploads": [
		{
			"object_key": "obj-1",
			"original_file_name": "dataset.jsonl",
			"presigned_url": "https://example.com/upload/dataset.jsonl",
			"expires_at": "2025-05-08T03:37:28Z"
		}
	]
}
`

func TestCreateModelEvalDatasetUploadPresignedURLs(t *testing.T) {
	setup()
	defer teardown()

	createReq := &CreateModelEvalDatasetUploadPresignedURLsRequest{
		Files: []*PresignedUrlFile{
			{FileName: "dataset.jsonl", FileSize: "2048"},
		},
	}

	mux.HandleFunc("/v2/gen-ai/model_evaluation/datasets/file_upload_presigned_urls", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)

		v := new(CreateModelEvalDatasetUploadPresignedURLsRequest)
		if err := json.NewDecoder(r.Body).Decode(v); err != nil {
			t.Errorf("Error decoding request body: %v", err)
		}
		assert.Len(t, v.Files, 1)
		assert.Equal(t, "dataset.jsonl", v.Files[0].FileName)
		assert.Equal(t, "2048", v.Files[0].FileSize)

		fmt.Fprint(w, createModelEvalDatasetUploadPresignedURLsResponse)
	})

	out, resp, err := client.GradientAI.CreateModelEvalDatasetUploadPresignedURLs(ctx, createReq)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "abcdef01-2345-6789-abcd-ef0123456789", out.RequestID)
	assert.Len(t, out.Uploads, 1)
	assert.Equal(t, "dataset.jsonl", out.Uploads[0].OriginalFileName)
	assert.Equal(t, "https://example.com/upload/dataset.jsonl", out.Uploads[0].PresignedURL)
	assert.NotNil(t, out.Uploads[0].ExpiresAt)
}

func TestCreateModelEvalDatasetUploadPresignedURLsNilRequest(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.CreateModelEvalDatasetUploadPresignedURLs(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

var getModelEvaluationRunResponse = `
{
	"run": {
		"eval_run_uuid": "12345678-1234-1234-1234-123456789012",
		"name": "my-eval-run",
		"status": "MODEL_EVALUATION_RUN_SUCCESSFUL",
		"candidate_model_uuid": "00000000-0000-0000-0000-000000000001",
		"candidate_model_name": "candidate-model",
		"candidate_model_source": "CANDIDATE_MODEL_SOURCE_SERVERLESS",
		"judge_model_uuid": "00000000-0000-0000-0000-000000000002",
		"judge_model_name": "judge-model",
		"dataset_uuid": "00000000-0000-0000-0000-000000000003",
		"dataset_name": "eval-dataset",
		"created_at": "2025-05-08T03:37:28Z",
		"started_at": "2025-05-08T03:38:00Z",
		"completed_at": "2025-05-08T03:45:00Z",
		"progress": {
			"candidate_rows_evaluated": 80,
			"judge_rows_evaluated": 75,
			"total_rows": 100
		},
		"result_summary": {
			"overall_score_percent": 87.5,
			"total_duration_seconds": 420,
			"metric_summaries": [
				{
					"metric_uuid": "44444444-4444-4444-4444-444444444444",
					"metric_name": "accuracy",
					"pass_percent": 90.0,
					"fail_percent": 10.0
				}
			],
			"performance_metrics": {
				"candidate_latency": {
					"avg_e2e_latency_ms": 250.5,
					"p95_latency_ms": 450.0
				},
				"token_usage": {
					"total_candidate_tokens": "12345"
				}
			},
			"pricing": {
				"currency": "USD",
				"total_cost": 1.5
			}
		}
	},
	"results": [
		{
			"input": "What is 1 + 1?",
			"output": "2",
			"ground_truth": "2",
			"candidate_model_name": "candidate-model",
			"candidate_model_uuid": "00000000-0000-0000-0000-000000000001",
			"metric_results": [
				{
					"metric_name": "accuracy",
					"metric_value_type": "METRIC_VALUE_TYPE_PERCENTAGE",
					"number_value": 100.0,
					"reasoning": "exact match"
				}
			]
		}
	],
	"links": {},
	"meta": {
		"total": 1
	}
}
`

func TestGetModelEvaluationRun(t *testing.T) {
	setup()
	defer teardown()

	evalRunUUID := "12345678-1234-1234-1234-123456789012"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_runs/%s", evalRunUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		assert.Equal(t, "2", r.URL.Query().Get("page"))
		assert.Equal(t, "25", r.URL.Query().Get("per_page"))
		fmt.Fprint(w, getModelEvaluationRunResponse)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationRun(ctx, evalRunUUID, &ModelEvaluationRunGetOptions{Page: 2, PerPage: 25})
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)

	assert.NotNil(t, out.Run)
	assert.Equal(t, evalRunUUID, out.Run.EvalRunUuid)
	assert.Equal(t, "my-eval-run", out.Run.Name)
	assert.Equal(t, ModelEvaluationRunSuccessful, out.Run.Status)
	assert.Equal(t, CandidateModelSourceServerless, out.Run.CandidateModelSource)
	assert.NotNil(t, out.Run.Progress)
	assert.Equal(t, int64(80), *out.Run.Progress.CandidateRowsEvaluated)
	assert.Equal(t, int64(75), *out.Run.Progress.JudgeRowsEvaluated)
	assert.Equal(t, int64(100), *out.Run.Progress.TotalRows)
	assert.NotNil(t, out.Run.ResultSummary)
	assert.Equal(t, 87.5, out.Run.ResultSummary.OverallScorePercent)
	assert.Equal(t, int64(420), out.Run.ResultSummary.TotalDurationSeconds)
	assert.Len(t, out.Run.ResultSummary.MetricSummaries, 1)
	assert.Equal(t, "accuracy", out.Run.ResultSummary.MetricSummaries[0].MetricName)
	assert.Equal(t, 90.0, out.Run.ResultSummary.MetricSummaries[0].PassPercent)
	assert.NotNil(t, out.Run.ResultSummary.PerformanceMetrics)
	assert.NotNil(t, out.Run.ResultSummary.PerformanceMetrics.CandidateLatency)
	assert.Equal(t, 250.5, out.Run.ResultSummary.PerformanceMetrics.CandidateLatency.AvgE2ELatencyMs)
	assert.Equal(t, "USD", out.Run.ResultSummary.Pricing.Currency)

	assert.Len(t, out.Results, 1)
	assert.Equal(t, "What is 1 + 1?", out.Results[0].Input)
	assert.Equal(t, "2", out.Results[0].Output)
	assert.Len(t, out.Results[0].MetricResults, 1)
	assert.Equal(t, MetricValueTypePercentage, out.Results[0].MetricResults[0].MetricValueType)
	assert.Equal(t, 100.0, out.Results[0].MetricResults[0].NumberValue)
}

func TestGetModelEvaluationRunNoOptions(t *testing.T) {
	setup()
	defer teardown()

	evalRunUUID := "12345678-1234-1234-1234-123456789012"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_runs/%s", evalRunUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		assert.Empty(t, r.URL.Query().Get("page"))
		assert.Empty(t, r.URL.Query().Get("per_page"))
		fmt.Fprint(w, getModelEvaluationRunResponse)
	})

	out, _, err := client.GradientAI.GetModelEvaluationRun(ctx, evalRunUUID, nil)
	assert.NoError(t, err)
	assert.NotNil(t, out)
}

func TestGetModelEvaluationRunMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.GetModelEvaluationRun(ctx, "", nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestGetModelEvaluationRunServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"not_found","message":"evaluation run not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationRun(ctx, "99999999-9999-9999-9999-999999999999", nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

var getModelEvaluationPresetResponse = `
{
	"preset": {
		"eval_preset_uuid": "12345678-1234-1234-1234-123456789012",
		"name": "my-preset",
		"dataset_uuid": "00000000-0000-0000-0000-000000000003",
		"dataset_name": "eval-dataset",
		"judge_model_uuid": "00000000-0000-0000-0000-000000000002",
		"judge_model_name": "judge-model",
		"created_at": "2025-05-08T03:37:28Z",
		"metrics": [
			{
				"metric_uuid": "44444444-4444-4444-4444-444444444444",
				"metric_name": "accuracy",
				"category": "METRIC_CATEGORY_CORRECTNESS"
			}
		],
		"star_metric": {
			"metric_uuid": "44444444-4444-4444-4444-444444444444",
			"name": "accuracy",
			"success_threshold": 0.8
		}
	}
}
`

func TestGetModelEvaluationPreset(t *testing.T) {
	setup()
	defer teardown()

	evalPresetUUID := "12345678-1234-1234-1234-123456789012"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_presets/%s", evalPresetUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, getModelEvaluationPresetResponse)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationPreset(ctx, evalPresetUUID)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.NotNil(t, out.Preset)
	assert.Equal(t, 200, resp.Response.StatusCode)

	preset := out.Preset
	assert.Equal(t, evalPresetUUID, preset.EvalPresetUuid)
	assert.Equal(t, "my-preset", preset.Name)
	assert.Equal(t, "eval-dataset", preset.DatasetName)
	assert.Len(t, preset.Metrics, 1)
	assert.Equal(t, "accuracy", preset.Metrics[0].MetricName)
	assert.Equal(t, MetricCategoryCorrectness, preset.Metrics[0].Category)
	assert.NotNil(t, preset.StarMetric)
	assert.NotNil(t, preset.StarMetric.SuccessThreshold)
	assert.InDelta(t, 0.8, *preset.StarMetric.SuccessThreshold, 0.0001)
}

func TestGetModelEvaluationPresetMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.GetModelEvaluationPreset(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestGetModelEvaluationPresetServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_presets/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"not_found","message":"preset not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationPreset(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

var getModelEvaluationRunResultsDownloadURLResponse = `
{
	"download_url": "https://example.com/results/run.json.gz?token=abc",
	"expires_at": "2025-05-08T04:37:28Z"
}
`

func TestGetModelEvaluationRunResultsDownloadURL(t *testing.T) {
	setup()
	defer teardown()

	evalRunUUID := "12345678-1234-1234-1234-123456789012"

	mux.HandleFunc(fmt.Sprintf("/v2/gen-ai/model_evaluation_runs/%s/results/download_url", evalRunUUID), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, getModelEvaluationRunResultsDownloadURLResponse)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationRunResultsDownloadURL(ctx, evalRunUUID)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Equal(t, "https://example.com/results/run.json.gz?token=abc", out.DownloadURL)
	assert.NotNil(t, out.ExpiresAt)
}

func TestGetModelEvaluationRunResultsDownloadURLMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.GetModelEvaluationRunResultsDownloadURL(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestGetModelEvaluationRunResultsDownloadURLServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs/99999999-9999-9999-9999-999999999999/results/download_url", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"not_found","message":"run not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.GetModelEvaluationRunResultsDownloadURL(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

var listModelEvaluationRunsResponse = `
{
	"runs": [
		{
			"eval_run_uuid": "11111111-1111-1111-1111-111111111111",
			"name": "run-1",
			"status": "MODEL_EVALUATION_RUN_SUCCESSFUL",
			"candidate_model_name": "candidate-1",
			"candidate_model_source": "CANDIDATE_MODEL_SOURCE_SERVERLESS",
			"created_at": "2025-05-08T03:37:28Z",
			"progress": {
				"candidate_rows_evaluated": 40,
				"judge_rows_evaluated": 35,
				"total_rows": 100
			}
		},
		{
			"eval_run_uuid": "22222222-2222-2222-2222-222222222222",
			"name": "run-2",
			"status": "MODEL_EVALUATION_RUN_QUEUED",
			"candidate_model_name": "candidate-2",
			"candidate_model_source": "CANDIDATE_MODEL_SOURCE_DEDICATED",
			"created_at": "2025-05-08T04:37:28Z"
		}
	],
	"links": {},
	"meta": {
		"total": 2
	}
}
`

func TestListModelEvaluationRuns(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		q := r.URL.Query()
		assert.Equal(t, "preset-uuid", q.Get("eval_preset_uuid"))
		assert.Equal(t, string(ModelEvaluationRunSuccessful), q.Get("status"))
		assert.Equal(t, "1", q.Get("page"))
		assert.Equal(t, "10", q.Get("per_page"))
		assert.Equal(t, []string{
			string(ModelEvaluationRunSuccessful),
			string(ModelEvaluationRunPartiallySuccessful),
		}, q["statuses"])
		assert.Equal(t, []string{
			string(CandidateModelSourceServerless),
			string(CandidateModelSourceDedicated),
		}, q["candidate_types"])
		assert.Equal(t, "needle", q.Get("search"))
		assert.Equal(t, string(ModelEvaluationRunSortFieldCreatedAt), q.Get("sort_by"))
		assert.Equal(t, string(ModelEvaluationRunSortDirectionDesc), q.Get("sort_direction"))
		fmt.Fprint(w, listModelEvaluationRunsResponse)
	})

	out, resp, err := client.GradientAI.ListModelEvaluationRuns(ctx, &ModelEvaluationRunListOptions{
		EvalPresetUUID: "preset-uuid",
		Status:         ModelEvaluationRunSuccessful,
		Statuses: []ModelEvaluationRunStatus{
			ModelEvaluationRunSuccessful,
			ModelEvaluationRunPartiallySuccessful,
		},
		CandidateTypes: []CandidateModelSource{
			CandidateModelSourceServerless,
			CandidateModelSourceDedicated,
		},
		Search:        "needle",
		SortBy:        ModelEvaluationRunSortFieldCreatedAt,
		SortDirection: ModelEvaluationRunSortDirectionDesc,
		ListOptions: ListOptions{
			Page:    1,
			PerPage: 10,
		},
	})
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Len(t, out.Runs, 2)
	assert.Equal(t, "run-1", out.Runs[0].Name)
	assert.Equal(t, ModelEvaluationRunSuccessful, out.Runs[0].Status)
	assert.Equal(t, CandidateModelSourceServerless, out.Runs[0].CandidateModelSource)
	assert.NotNil(t, out.Runs[0].Progress)
	assert.Equal(t, int64(40), *out.Runs[0].Progress.CandidateRowsEvaluated)
	assert.Equal(t, int64(35), *out.Runs[0].Progress.JudgeRowsEvaluated)
	assert.Equal(t, int64(100), *out.Runs[0].Progress.TotalRows)
	assert.Nil(t, out.Runs[1].Progress)
	assert.Equal(t, ModelEvaluationRunQueued, out.Runs[1].Status)
	assert.Equal(t, CandidateModelSourceDedicated, out.Runs[1].CandidateModelSource)
}

func TestListModelEvaluationRunsNoOptions(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_runs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		q := r.URL.Query()
		assert.Empty(t, q.Get("eval_preset_uuid"))
		assert.Empty(t, q.Get("status"))
		assert.Empty(t, q["statuses"])
		assert.Empty(t, q["candidate_types"])
		assert.Empty(t, q.Get("search"))
		assert.Empty(t, q.Get("sort_by"))
		assert.Empty(t, q.Get("sort_direction"))
		fmt.Fprint(w, listModelEvaluationRunsResponse)
	})

	out, _, err := client.GradientAI.ListModelEvaluationRuns(ctx, nil)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Len(t, out.Runs, 2)
}

var listModelEvaluationPresetsResponse = `
{
	"presets": [
		{
			"eval_preset_uuid": "11111111-1111-1111-1111-111111111111",
			"name": "preset-1",
			"dataset_uuid": "00000000-0000-0000-0000-000000000003",
			"dataset_name": "eval-dataset",
			"created_at": "2025-05-08T03:37:28Z"
		},
		{
			"eval_preset_uuid": "22222222-2222-2222-2222-222222222222",
			"name": "preset-2",
			"dataset_uuid": "00000000-0000-0000-0000-000000000004",
			"dataset_name": "another-dataset",
			"created_at": "2025-05-08T04:37:28Z"
		}
	]
}
`

func TestListModelEvaluationPresets(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_presets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listModelEvaluationPresetsResponse)
	})

	out, resp, err := client.GradientAI.ListModelEvaluationPresets(ctx)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Len(t, out.Presets, 2)
	assert.Equal(t, "preset-1", out.Presets[0].Name)
	assert.Equal(t, "eval-dataset", out.Presets[0].DatasetName)
	assert.Equal(t, "preset-2", out.Presets[1].Name)
}

func TestListModelEvaluationPresetsServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_presets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"server_error","message":"boom"}`, http.StatusInternalServerError)
	})

	out, resp, err := client.GradientAI.ListModelEvaluationPresets(ctx)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusInternalServerError, resp.Response.StatusCode)
}

var listModelEvaluationMetricsResponse = `
{
	"metrics": [
		{
			"metric_uuid": "11111111-1111-1111-1111-111111111111",
			"metric_name": "accuracy",
			"description": "Measures the correctness of the response.",
			"metric_type": "METRIC_TYPE_GENERAL_QUALITY",
			"metric_value_type": "METRIC_VALUE_TYPE_PERCENTAGE",
			"category": "METRIC_CATEGORY_CORRECTNESS",
			"range_min": 0,
			"range_max": 1,
			"inverted": false
		},
		{
			"metric_uuid": "22222222-2222-2222-2222-222222222222",
			"metric_name": "toxicity",
			"description": "Measures whether the response contains toxic content.",
			"metric_type": "METRIC_TYPE_GENERAL_QUALITY",
			"metric_value_type": "METRIC_VALUE_TYPE_NUMBER",
			"category": "METRIC_CATEGORY_SAFETY_AND_SECURITY",
			"inverted": true
		}
	]
}
`

func TestListModelEvaluationMetrics(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_metrics", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listModelEvaluationMetricsResponse)
	})

	out, resp, err := client.GradientAI.ListModelEvaluationMetrics(ctx)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Len(t, out.Metrics, 2)
	assert.Equal(t, "accuracy", out.Metrics[0].MetricName)
	assert.Equal(t, MetricTypeGeneralQuality, out.Metrics[0].MetricType)
	assert.Equal(t, MetricValueTypePercentage, out.Metrics[0].MetricValueType)
	assert.Equal(t, MetricCategoryCorrectness, out.Metrics[0].Category)
	assert.Equal(t, "toxicity", out.Metrics[1].MetricName)
	assert.True(t, out.Metrics[1].Inverted)
}

func TestListModelEvaluationMetricsServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/model_evaluation_metrics", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"server_error","message":"boom"}`, http.StatusInternalServerError)
	})

	out, resp, err := client.GradientAI.ListModelEvaluationMetrics(ctx)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusInternalServerError, resp.Response.StatusCode)
}

var listEvaluationDatasetsResponse = `
{
	"evaluation_datasets": [
		{
			"created_at": "2023-01-01T00:00:00Z",
			"dataset_name": "dataset-1",
			"dataset_type": "EVALUATION_DATASET_TYPE_MODEL",
			"dataset_uuid": "11111111-1111-1111-1111-111111111111",
			"file_size": "12345",
			"has_ground_truth": true,
			"row_count": 100
		},
		{
			"created_at": "2023-02-01T00:00:00Z",
			"dataset_name": "dataset-2",
			"dataset_type": "EVALUATION_DATASET_TYPE_ADK",
			"dataset_uuid": "22222222-2222-2222-2222-222222222222",
			"file_size": "67890",
			"has_ground_truth": false,
			"row_count": 50
		}
	]
}
`

func TestListEvaluationDatasets(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/evaluation_datasets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, listEvaluationDatasetsResponse)
	})

	out, resp, err := client.GradientAI.ListEvaluationDatasets(ctx, nil)
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Len(t, out.EvaluationDatasets, 2)
	assert.Equal(t, "dataset-1", out.EvaluationDatasets[0].DatasetName)
	assert.Equal(t, EvaluationDatasetTypeModel, out.EvaluationDatasets[0].DatasetType)
	assert.Equal(t, "11111111-1111-1111-1111-111111111111", out.EvaluationDatasets[0].DatasetUUID)
	assert.Equal(t, int64(100), out.EvaluationDatasets[0].RowCount)
	assert.True(t, out.EvaluationDatasets[0].HasGroundTruth)
	assert.Equal(t, "dataset-2", out.EvaluationDatasets[1].DatasetName)
	assert.Equal(t, EvaluationDatasetTypeADK, out.EvaluationDatasets[1].DatasetType)
}

func TestListEvaluationDatasetsWithOptions(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/evaluation_datasets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		assert.Equal(t, string(EvaluationDatasetTypeADK), r.URL.Query().Get("dataset_type"))
		fmt.Fprint(w, listEvaluationDatasetsResponse)
	})

	out, resp, err := client.GradientAI.ListEvaluationDatasets(ctx, &EvaluationDatasetListOptions{
		DatasetType: EvaluationDatasetTypeADK,
	})
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
	assert.Len(t, out.EvaluationDatasets, 2)
}

func TestListEvaluationDatasetsServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/evaluation_datasets", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		http.Error(w, `{"id":"server_error","message":"boom"}`, http.StatusInternalServerError)
	})

	out, resp, err := client.GradientAI.ListEvaluationDatasets(ctx, nil)
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusInternalServerError, resp.Response.StatusCode)
}

func TestDeleteEvaluationDataset(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/evaluation_datasets/12345678-1234-1234-1234-123456789012", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		fmt.Fprint(w, `{}`)
	})

	out, resp, err := client.GradientAI.DeleteEvaluationDataset(ctx, "12345678-1234-1234-1234-123456789012")
	assert.NoError(t, err)
	assert.NotNil(t, out)
	assert.Equal(t, 200, resp.Response.StatusCode)
}

func TestDeleteEvaluationDatasetMissingUUID(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteEvaluationDataset(ctx, "")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}

func TestDeleteEvaluationDatasetServerError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/v2/gen-ai/evaluation_datasets/99999999-9999-9999-9999-999999999999", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
		http.Error(w, `{"id":"not_found","message":"evaluation dataset not found"}`, http.StatusNotFound)
	})

	out, resp, err := client.GradientAI.DeleteEvaluationDataset(ctx, "99999999-9999-9999-9999-999999999999")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.NotNil(t, resp)
	assert.Equal(t, http.StatusNotFound, resp.Response.StatusCode)
}

func TestDeleteEvaluationDatasetInvalidURL(t *testing.T) {
	setup()
	defer teardown()

	out, resp, err := client.GradientAI.DeleteEvaluationDataset(ctx, "bad\nuuid")
	assert.Error(t, err)
	assert.Nil(t, out)
	assert.Nil(t, resp)
}