mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(agents): code quality fixes from self-review
- Fix test_generate_key_helper_fn_agent_id: remove bare except clause, use explicit assert mock_insert.called, use .kwargs for clean arg access - Remove no-op conditional in handleNext (both branches were identical) - Validate selectedExistingKey before calling keyUpdateCall - Validate selectedAgentId before setting on formValues in create_key_button
This commit is contained in:
parent
10a5417a21
commit
6599143729
3 changed files with 35 additions and 25 deletions
|
|
@ -6178,31 +6178,36 @@ async def test_generate_key_with_agent_id():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_helper_fn_agent_id():
|
||||
"""Test that generate_key_helper_fn correctly stores agent_id in key_data."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
"""Test that generate_key_helper_fn passes agent_id into the insert_data call."""
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
import litellm.proxy.management_endpoints.key_management_endpoints as km
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_insert = AsyncMock(return_value=MagicMock(token="sk-test", created_at=None, updated_at=None, litellm_budget_table=None))
|
||||
mock_insert = AsyncMock(
|
||||
return_value=MagicMock(
|
||||
token="sk-test",
|
||||
created_at=None,
|
||||
updated_at=None,
|
||||
litellm_budget_table=None,
|
||||
)
|
||||
)
|
||||
mock_prisma_client.insert_data = mock_insert
|
||||
|
||||
with patch.object(km, "prisma_client", mock_prisma_client):
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
||||
try:
|
||||
result = await generate_key_helper_fn(
|
||||
request_type="key",
|
||||
agent_id="test-agent-456",
|
||||
key_alias="test-agent-key",
|
||||
models=[],
|
||||
table_name="key",
|
||||
)
|
||||
except Exception:
|
||||
# We just verify the agent_id was passed to insert_data
|
||||
pass
|
||||
await generate_key_helper_fn(
|
||||
request_type="key",
|
||||
agent_id="test-agent-456",
|
||||
key_alias="test-agent-key",
|
||||
models=[],
|
||||
table_name="key",
|
||||
)
|
||||
|
||||
# Verify insert_data was called with agent_id in key_data
|
||||
if mock_insert.called:
|
||||
call_kwargs = mock_insert.call_args
|
||||
key_data = call_kwargs[1].get("data") or (call_kwargs[0][0] if call_kwargs[0] else {})
|
||||
assert key_data.get("agent_id") == "test-agent-456"
|
||||
assert mock_insert.called, "insert_data was never called"
|
||||
# insert_data is called as insert_data(data=key_data, ...)
|
||||
call_kwargs = mock_insert.call_args.kwargs
|
||||
key_data = call_kwargs.get("data", {})
|
||||
assert key_data.get("agent_id") == "test-agent-456", (
|
||||
f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -93,11 +93,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
|
|||
const handleNext = async () => {
|
||||
try {
|
||||
if (currentStep === 0) {
|
||||
const fieldsToValidate =
|
||||
agentType === CUSTOM_AGENT_TYPE
|
||||
? ["agent_name"]
|
||||
: ["agent_name"];
|
||||
await form.validateFields(fieldsToValidate);
|
||||
await form.validateFields(["agent_name"]);
|
||||
const agentName = form.getFieldValue("agent_name");
|
||||
if (agentName && !newKeyName) {
|
||||
setNewKeyName(`${agentName}-key`);
|
||||
|
|
@ -184,7 +180,12 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
|
|||
newKeyModels,
|
||||
);
|
||||
setCreatedKeyValue(keyResponse.key || null);
|
||||
} else if (keyAssignOption === "existing_key" && selectedExistingKey) {
|
||||
} else if (keyAssignOption === "existing_key") {
|
||||
if (!selectedExistingKey) {
|
||||
message.error("Please select an existing key to assign");
|
||||
setIsSubmitting(false);
|
||||
return;
|
||||
}
|
||||
await keyUpdateCall(accessToken, {
|
||||
key: selectedExistingKey,
|
||||
agent_id: agentId,
|
||||
|
|
|
|||
|
|
@ -298,6 +298,10 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
|||
if (keyOwner === "you") {
|
||||
formValues.user_id = userID;
|
||||
} else if (keyOwner === "agent") {
|
||||
if (!selectedAgentId) {
|
||||
message.error("Please select an agent");
|
||||
return;
|
||||
}
|
||||
formValues.agent_id = selectedAgentId;
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue