Repository: aipotheosis-labs/aci
File: aipolabs/server/tests/routes/projects/test_projects_agents.py (Python)

Code before:
        headers={"Authorization": f"Bearer {dummy_user_bearer_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    agent_public = AgentPublic.model_validate(response.json())
    assert agent_public.name == body.name
    assert agent_public.description == body.description
    assert agent_public.project_id == dummy_project_1.id

    # Verify the agent was actually created in the database and values match returned values
    agent = db_session.execute(
        select(Agent).filter(Agent.id == agent_public.id)
    ).scalar_one_or_none()

    assert agent is not None
    assert agent_public.model_dump() == AgentPublic.model_validate(agent).model_dump()

    # check api keys
    api_key = db_session.execute(
        select(APIKey).filter(APIKey.agent_id == agent.id)
    ).scalar_one_or_none()
    assert api_key is not None
    assert len(agent_public.api_keys) == 1
    assert agent_public.api_keys[0].key == api_key.key


def test_create_custom_instructions(
    test_client: TestClient,
    db_session: Session,
    dummy_user: User,
    dummy_agent_1: Agent,
    dummy_app_google: App,
    dummy_user_bearer_token: str,
) -> None:

    ENDPOINT = f"{config.ROUTER_PREFIX_PROJECTS}/agents/{dummy_agent_1.id}"
    response = test_client.patch(
        ENDPOINT,
        json={"custom_instructions": {str(dummy_app_google.id): "test instructions"}},
        headers={"Authorization": f"Bearer {dummy_user_bearer_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    agent_public = AgentPublic.model_validate(response.json())
    assert agent_public.custom_instructions == {str(dummy_app_google.id): "test instructions"}


def test_update_agent(
    test_client: TestClient,
    db_session: Session,
    dummy_agent_1: Agent,
    dummy_app_google: App,
    dummy_app_github: App,

Review comment (refactor):
this test case is duplicate with part of the testing in test_update_agent.
(additional the name of this test case "create_custom_instructions" is a bit confusing as you use "patch" in the actual implementation. maybe "test_update_agent_add_custom_instructions" would be better. But recommend just deleting it instead) 

Code after:
        headers={"Authorization": f"Bearer {dummy_user_bearer_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    agent_public = AgentPublic.model_validate(response.json())
    assert agent_public.name == body.name
    assert agent_public.description == body.description
    assert agent_public.project_id == dummy_project_1.id

    # Verify the agent was actually created in the database and values match returned values
    agent = db_session.execute(
        select(Agent).filter(Agent.id == agent_public.id)
    ).scalar_one_or_none()

    assert agent is not None
    assert agent_public.model_dump() == AgentPublic.model_validate(agent).model_dump()

    # check api keys
    api_key = db_session.execute(
        select(APIKey).filter(APIKey.agent_id == agent.id)
    ).scalar_one_or_none()
    assert api_key is not None
    assert len(agent_public.api_keys) == 1
    assert agent_public.api_keys[0].key == api_key.key


def test_update_agent(
    test_client: TestClient,
    db_session: Session,
    dummy_project_1: Project,
    dummy_agent_1: Agent,
    dummy_app_google: App,
    dummy_app_github: App,
    dummy_user_bearer_token: str,
) -> None:
    ENDPOINT = f"{config.ROUTER_PREFIX_PROJECTS}/{dummy_project_1.id}/agents/{dummy_agent_1.id}"

    # Test updating name and description
    response = test_client.patch(
        ENDPOINT,
        json={"name": "Updated Agent Name", "description": "Updated description"},
        headers={"Authorization": f"Bearer {dummy_user_bearer_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    agent_public = AgentPublic.model_validate(response.json())
    assert agent_public.name == "Updated Agent Name"
    assert agent_public.description == "Updated description"

    # Test updating excluded apps and functions
    response = test_client.patch(
        ENDPOINT,
        json={