diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b694b0b8c66..2a385c4c42a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2654,7 +2654,7 @@ async def _validate_update_key_data( @router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper -async def update_key_fn( # noqa: C901 # single endpoint handling many optional key-update fields; decomposition is out of scope here +async def update_key_fn( request: Request, data: UpdateKeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -4743,7 +4743,7 @@ async def _execute_virtual_key_regeneration( dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def regenerate_key_fn( # noqa: C901 # single endpoint handling many optional key-regeneration fields; decomposition is out of scope here +async def regenerate_key_fn( key: str | None = None, data: RegenerateKeyRequest | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 22f5fe3d8aa..d5517525322 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1966,7 +1966,7 @@ async def update_team( ) # Verify caller has access to manage this team - team_for_auth = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()) + team_for_auth: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()) await _verify_team_access( team_obj=team_for_auth, user_api_key_dict=user_api_key_dict, @@ -4062,7 +4062,11 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) - _team_info.resolved_logging_exporters = resolved_logging_exporter_names( + # Set on the object that is actually returned: the helper above hands back + # ``_team_info`` itself only when the team has no access groups, and a + # ``model_copy`` when it does, so mutating ``_team_info`` here would be + # discarded for every team that inherits from an access group. + resolved_team_info.resolved_logging_exporters = resolved_logging_exporter_names( team_id, _team_info.organization_id, ) diff --git a/tests/e2e/ui/ui_track_badaccess.mjs b/tests/e2e/ui/ui_track_badaccess.mjs new file mode 100644 index 00000000000..d1962b2b1aa --- /dev/null +++ b/tests/e2e/ui/ui_track_badaccess.mjs @@ -0,0 +1,64 @@ +import { chromium } from "playwright"; +import fs from "node:fs"; + +const BASE = "http://127.0.0.1:21501"; +const EXEC = + "/Users/yucheng/Library/Caches/ms-playwright/chromium_headless_shell-1234/chrome-headless-shell-mac-arm64/chrome-headless-shell"; +const results = []; +const notes = []; +function check(id, desc, ok, actual, expected) { + results.push({ id, desc, status: ok ? "PASS" : "FAIL", actual: ok ? undefined : actual, expected: ok ? undefined : expected }); +} + +const who = process.argv[2] || "admin"; +const creds = who === "viewer" ? ["uitrk-viewer@example.com", "uitrk-viewer-pw"] : ["admin", "sk-uitrack-21501"]; + +const browser = await chromium.launch({ executablePath: EXEC, headless: true }); +const ctx = await browser.newContext({ viewport: { width: 1600, height: 1200 } }); +const page = await ctx.newPage(); +const errs = []; +page.on("pageerror", (e) => errs.push(String(e).slice(0, 400))); + +await page.goto(`${BASE}/ui/login`, { waitUntil: "domcontentloaded" }); +await page.getByPlaceholder("Enter your username").fill(creds[0]); +await page.getByPlaceholder("Enter your password").fill(creds[1]); +await page.getByRole("button", { name: "Login", exact: true }).click(); +await page.waitForLoadState("networkidle"); + +await page.goto(`${BASE}/ui/logging-and-alerts`, { waitUntil: "domcontentloaded" }); +await page.waitForLoadState("networkidle"); +await page.waitForTimeout(4000); +const body = await page.evaluate(() => document.body.innerText); +const tables = await page.evaluate(() => document.querySelectorAll("table").length); +notes.push({ k: "body", v: body.slice(0, 400) }); +notes.push({ k: "pageerrors", v: errs.slice(0, 4) }); +notes.push({ k: "tables", v: tables }); +await page.screenshot({ path: `/tmp/uitrack_badaccess_${who}.png`, fullPage: true }).catch(() => {}); + +const rendered = tables > 0 && body.includes("Active Logging Callbacks"); +check( + `BAD1-${who}`, + `Destinations page renders for ${who} with one stored destination whose credential_info.access is a malformed (non-list) teams value`, + rendered, + `tables=${tables}; body="${body.slice(0, 120).replace(/\n/g, " / ")}"; pageerror="${(errs[0] || "").slice(0, 160)}"`, + "table renders, malformed row shows a dash for Scope", +); + +// blast radius: is any other dashboard page affected? +await page.goto(`${BASE}/ui/teams`, { waitUntil: "domcontentloaded" }); +await page.waitForLoadState("networkidle"); +await page.waitForTimeout(2500); +const teamsBody = await page.evaluate(() => document.body.innerText); +notes.push({ k: "teamsBody", v: teamsBody.slice(0, 200) }); +check( + `BAD2-${who}`, + `The failure is scoped to the destinations page (/ui/teams still renders) for ${who}`, + !teamsBody.includes("This page couldn"), + teamsBody.slice(0, 150), + "/ui/teams renders normally", +); + +await browser.close(); +fs.writeFileSync(`/tmp/uitrack_badaccess_${who}.json`, JSON.stringify({ results, notes }, null, 2)); +console.log(JSON.stringify(results, null, 2)); +console.log("NOTES", JSON.stringify(notes, null, 2).slice(0, 1500)); diff --git a/tests/e2e/ui/ui_track_run.mjs b/tests/e2e/ui/ui_track_run.mjs new file mode 100644 index 00000000000..953b80e1ced --- /dev/null +++ b/tests/e2e/ui/ui_track_run.mjs @@ -0,0 +1,524 @@ +import { chromium } from "playwright"; +import fs from "node:fs"; + +const BASE = "http://127.0.0.1:21501"; +const EXEC = + "/Users/yucheng/Library/Caches/ms-playwright/chromium_headless_shell-1234/chrome-headless-shell-mac-arm64/chrome-headless-shell"; +const ADMIN_USER = "admin"; +const ADMIN_PASS = "sk-uitrack-21501"; +const VIEWER_USER = "uitrk-viewer@example.com"; +const VIEWER_PASS = "uitrk-viewer-pw"; +const OUT = process.argv[2] || "/tmp/ui_track_result.json"; + +const results = []; +const notes = []; +let consoleErrors = []; + +function check(id, desc, ok, actual, expected) { + results.push({ id, desc, status: ok ? "PASS" : "FAIL", actual: ok ? undefined : actual, expected: ok ? undefined : expected }); +} + +function note(k, v) { + notes.push({ k, v }); +} + +async function login(page, user, pass) { + await page.goto(`${BASE}/ui/login`, { waitUntil: "domcontentloaded" }); + await page.getByPlaceholder("Enter your username").fill(user); + await page.getByPlaceholder("Enter your password").fill(pass); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await page.waitForLoadState("networkidle"); +} + +async function dismissFeedback(page) { + const b = page.getByText("Don't ask me again"); + if (await b.isVisible({ timeout: 1500 }).catch(() => false)) { + await b.click().catch(() => {}); + } +} + +async function gotoDestinations(page) { + await page.goto(`${BASE}/ui/logging-and-alerts`, { waitUntil: "domcontentloaded" }); + await page.waitForLoadState("networkidle"); + await dismissFeedback(page); + // wait for the destination rows to land (credentials fetch is async) + await page + .locator("table tbody tr", { hasText: "uitrk-global" }) + .first() + .waitFor({ timeout: 20000 }) + .catch(() => {}); + await page.waitForTimeout(1200); +} + +async function scrape(page) { + return page.evaluate(() => { + const tables = [...document.querySelectorAll("table")]; + const table = tables[0]; + if (!table) return { headers: [], rows: [], tableCount: tables.length, bodyText: document.body.innerText.slice(0, 800) }; + const headers = [...table.querySelectorAll("thead th")].map((th) => th.innerText.trim()); + const rows = [...table.querySelectorAll("tbody tr")].map((tr) => { + const tds = [...tr.querySelectorAll("td")]; + const nameCell = tds[0]; + const spans = nameCell ? [...nameCell.querySelectorAll("span")] : []; + const actionsCell = tds[tds.length - 1]; + const trigger = actionsCell ? actionsCell.querySelector('[data-testid^="callback-actions-"]') : null; + const scopeCell = tds[2]; + return { + cells: tds.map((td) => td.innerText.trim()), + name: spans[0] ? spans[0].innerText.trim() : "", + sub: spans[1] ? spans[1].innerText.trim() : "", + mode: tds[1] ? tds[1].innerText.trim() : "", + scope: scopeCell ? scopeCell.innerText.trim() : "", + scopeBadges: scopeCell + ? [...scopeCell.querySelectorAll("span")].map((e) => e.innerText.trim()).filter(Boolean) + : [], + scopeTitle: scopeCell && scopeCell.querySelector("[title]") ? scopeCell.querySelector("[title]").getAttribute("title") : null, + hasTrigger: !!trigger, + triggerTestId: trigger ? trigger.getAttribute("data-testid") : null, + }; + }); + return { headers, rows, tableCount: tables.length, bodyText: "" }; + }); +} + +async function openMenu(page, testId) { + await page.locator(`[data-testid="${testId}"]`).click(); + await page.waitForTimeout(500); + const items = await page.evaluate(() => { + const menu = document.querySelector('[role="menu"]'); + if (!menu) return null; + return [...menu.querySelectorAll('[role="menuitem"]')].map((el) => ({ + text: el.innerText.trim(), + testId: el.getAttribute("data-testid"), + })); + }); + return items; +} + +async function closeMenu(page) { + await page.keyboard.press("Escape"); + await page.waitForTimeout(300); +} + +async function apiGet(path) { + const r = await fetch(`${BASE}${path}`, { headers: { Authorization: `Bearer ${ADMIN_PASS}` } }); + return r.json(); +} + +const EXPECTED = { + "uitrk-global": { sub: "Generic OTLP Collector · http://127.0.0.1:21599/v1/traces", scope: "Global access" }, + "uitrk-team": { sub: "Arize", scope: "team: uitrk-team-1" }, + "uitrk-org": { sub: "Weave", scope: "org: uitrk-org-1" }, + "uitrk-team-and-org": { sub: "Langfuse OTEL · https://cloud.langfuse.com", scope: "team: uitrk-team-2|org: uitrk-org-2" }, + "uitrk-many-teams": { sub: "Generic OTLP Collector", scope: "FOUR_PLUS_MORE" }, + "uitrk-no-access": { sub: "Generic OTLP Collector", scope: "—" }, + "uitrk-empty-access": { sub: "Generic OTLP Collector", scope: "—" }, + "uitrk-no-description": { sub: "-", scope: "Not active" }, + "uitrk-adapter-reject": { sub: "Langfuse OTEL", scope: "Not active" }, + datadog: { sub: "Generic OTLP Collector", scope: "team: uitrk-team-1" }, +}; +const DEST_COUNT = 10; + +(async () => { + const browser = await chromium.launch({ executablePath: EXEC, headless: true }); + const ctx = await browser.newContext({ viewport: { width: 1600, height: 1200 } }); + const page = await ctx.newPage(); + page.on("console", (m) => { + if (m.type() === "error") consoleErrors.push(m.text().slice(0, 300)); + }); + page.on("pageerror", (e) => consoleErrors.push("PAGEERROR: " + String(e).slice(0, 300))); + + const creds = await apiGet("/credentials"); + const rtd = Object.fromEntries( + creds.credentials + .filter((c) => (c.credential_info || {}).credential_type === "logging") + .map((c) => [c.credential_name, c.resolves_to_destination]), + ); + note("resolves_to_destination", rtd); + + // ---------- ADMIN ---------- + await login(page, ADMIN_USER, ADMIN_PASS); + await gotoDestinations(page); + + const url = page.url(); + check("N1", "Destinations page renders at path route /ui/logging-and-alerts", url.includes("/ui/logging-and-alerts"), url, "/ui/logging-and-alerts"); + + const tabVisible = await page.getByText("Logging Callbacks", { exact: true }).first().isVisible().catch(() => false); + check("N2", "Logging Callbacks tab is exposed", tabVisible, String(tabVisible), "true"); + + const headingVisible = await page.getByText("Active Logging Callbacks", { exact: true }).first().isVisible().catch(() => false); + check("T2", '"Active Logging Callbacks" heading present', headingVisible, String(headingVisible), "true"); + + let scraped = await scrape(page); + note("adminScrape", scraped); + await page.screenshot({ path: "/tmp/uitrack_admin_table.png", fullPage: true }).catch(() => {}); + + const hdr = scraped.headers.map((h) => h.trim()); + check("T1", "Table has Callback Name / Mode / Scope columns", hdr[0] === "Callback Name" && hdr[1] === "Mode" && hdr[2] === "Scope", JSON.stringify(hdr), '["Callback Name","Mode","Scope",...]'); + check("T3", "Actions column header is screen-reader-only (renders empty)", hdr.length === 4 && hdr[3] === "Actions", JSON.stringify(hdr), '4 headers, last "Actions"'); + + const byName = {}; + for (const r of scraped.rows) { + const k = r.name + "|" + r.mode; + byName[k] = r; + } + const destRows = scraped.rows.filter((r) => r.mode === "—"); + const cfgRows = scraped.rows.filter((r) => r.mode !== "—"); + check("T4", "Config callbacks and destinations share one table", scraped.tableCount >= 1 && destRows.length > 0 && cfgRows.length > 0, `tables=${scraped.tableCount} dest=${destRows.length} cfg=${cfgRows.length}`, "one table containing both kinds"); + check("T5", "All 10 destination fixtures render as rows", destRows.length === DEST_COUNT, `${destRows.length} destination rows: ${destRows.map((r) => r.name).join(",")}`, "10"); + + // per-fixture checks + for (const [name, exp] of Object.entries(EXPECTED)) { + const row = destRows.find((r) => r.name === name); + const idBase = "R-" + name; + if (!row) { + check(idBase + "-name", `Row "${name}" present with its own name`, false, "row not found", name); + check(idBase + "-sub", `Row "${name}" backend sub-line`, false, "row not found", exp.sub); + check(idBase + "-mode", `Row "${name}" Mode is a dash`, false, "row not found", "—"); + check(idBase + "-scope", `Row "${name}" Scope`, false, "row not found", exp.scope); + continue; + } + check(idBase + "-name", `Row "${name}" present with its own name`, row.name === name, row.name, name); + check(idBase + "-sub", `Row "${name}" backend sub-line`, row.sub === exp.sub, row.sub, exp.sub); + check(idBase + "-mode", `Row "${name}" Mode is a dash (destinations have no mode)`, row.mode === "—", row.mode, "—"); + + let scopeOk; + let scopeExp = exp.scope; + if (exp.scope === "FOUR_PLUS_MORE") { + const badges = row.scope.split("\n").map((s) => s.trim()).filter(Boolean); + const teamBadges = badges.filter((b) => b.startsWith("team:")); + scopeOk = teamBadges.length === 4 && badges.includes("+2 more"); + scopeExp = "4 team badges + '+2 more'"; + } else if (exp.scope.includes("|")) { + scopeOk = exp.scope.split("|").every((part) => row.scope.includes(part)); + } else { + scopeOk = row.scope.replace(/\s+/g, " ").trim() === exp.scope; + } + check(idBase + "-scope", `Row "${name}" Scope cell`, scopeOk, JSON.stringify(row.scope), scopeExp); + } + + // Cross-check: Not active <=> resolves_to_destination === false + for (const name of Object.keys(EXPECTED)) { + const row = destRows.find((r) => r.name === name); + if (!row) continue; + const uiNotActive = row.scope.includes("Not active"); + const backendFalse = rtd[name] === false; + check( + "X-rtd-" + name, + `"Not active" for ${name} matches GET /credentials resolves_to_destination`, + uiNotActive === backendFalse, + `ui_not_active=${uiNotActive} resolves_to_destination=${rtd[name]}`, + "the two agree", + ); + } + + // Cross-check against /team/info + /organization/info resolved_logging_exporters + const teamsList = await apiGet("/team/list"); + const teamIdByAlias = Object.fromEntries((teamsList || []).map((t) => [t.team_alias, t.team_id])); + note("teamIdByAlias", teamIdByAlias); + const teamExpect = { + "uitrk-team-1": ["uitrk-global", "uitrk-team", "uitrk-many-teams", "datadog"], + "uitrk-team-2": ["uitrk-global", "uitrk-team-and-org", "uitrk-many-teams"], + "uitrk-team-3": ["uitrk-global", "uitrk-many-teams"], + "uitrk-team-7": ["uitrk-global"], + }; + for (const [alias, expNames] of Object.entries(teamExpect)) { + const info = await apiGet(`/team/info?team_id=${teamIdByAlias[alias]}`); + const actual = ((info.team_info || info).resolved_logging_exporters || []).slice().sort(); + // what the UI Scope column claims for this team: rows whose scope shows Global access, + // or a team badge naming this alias (expanded badges only) + const uiClaim = destRows + .filter((r) => r.scope.includes("Global access") || r.scope.includes(`team: ${alias}`)) + .map((r) => r.name) + .sort(); + const missingFromUi = actual.filter((n) => !uiClaim.includes(n)); + // uitrk-many-teams collapses behind "+N more" for teams 5/6; allow the collapse + const collapsed = missingFromUi.filter((n) => n !== "uitrk-many-teams"); + check( + "X-team-" + alias, + `Scope column agrees with /team/info resolved_logging_exporters for ${alias}`, + collapsed.length === 0, + `backend=${JSON.stringify(actual)} ui=${JSON.stringify(uiClaim)}`, + "no backend-granted destination missing from the UI scope (modulo +N more collapse)", + ); + } + const orgList = await apiGet("/organization/list"); + const orgIdByAlias = Object.fromEntries((orgList || []).map((o) => [o.organization_alias, o.organization_id])); + for (const alias of ["uitrk-org-1", "uitrk-org-2"]) { + const info = await apiGet(`/organization/info?organization_id=${orgIdByAlias[alias]}`); + const actual = (info.resolved_logging_exporters || []).slice().sort(); + const uiClaim = destRows + .filter((r) => r.scope.includes("Global access") || r.scope.includes(`org: ${alias}`)) + .map((r) => r.name) + .sort(); + check( + "X-org-" + alias, + `Scope column agrees with /organization/info resolved_logging_exporters for ${alias}`, + JSON.stringify(actual) === JSON.stringify(uiClaim), + `backend=${JSON.stringify(actual)} ui=${JSON.stringify(uiClaim)}`, + "identical sets", + ); + } + + // Check 3: the destination named datadog vs the real datadog config callback + const ddDest = destRows.find((r) => r.name === "datadog"); + const ddCfg = cfgRows.find((r) => r.name.toLowerCase() === "datadog"); + check("DD1", 'Destination named "datadog" keeps its own lowercase name', ddDest && ddDest.name === "datadog", ddDest ? ddDest.name : "missing", "datadog"); + check("DD2", "Real datadog config callback row also present", !!ddCfg, ddCfg ? ddCfg.name : "missing", 'a config-callback row named Datadog'); + check("DD3", "The two datadog rows are distinguishable (different displayed name)", !!ddDest && !!ddCfg && ddDest.name !== ddCfg.name, `dest="${ddDest && ddDest.name}" cfg="${ddCfg && ddCfg.name}"`, "different strings"); + check("DD4", "datadog config-callback row shows a Mode badge, destination shows a dash", !!ddCfg && ddCfg.mode !== "—" && !!ddDest && ddDest.mode === "—", `cfg mode="${ddCfg && ddCfg.mode}" dest mode="${ddDest && ddDest.mode}"`, "cfg has a mode, dest has —"); + check("DD5", "datadog config-callback Scope is a dash (not a destination)", !!ddCfg && ddCfg.scope === "—", ddCfg ? ddCfg.scope : "missing", "—"); + check("DD6", "datadog config-callback row has no backend sub-line", !!ddCfg && ddCfg.sub === "", ddCfg ? `"${ddCfg.sub}"` : "missing", '""'); + check("DD7", "The two datadog rows have distinct action triggers", !!ddDest && !!ddCfg && ddDest.triggerTestId !== ddCfg.triggerTestId, `${ddDest && ddDest.triggerTestId} vs ${ddCfg && ddCfg.triggerTestId}`, "distinct data-testids"); + + // Actions menus (admin) + const destTrigger = destRows.find((r) => r.name === "uitrk-global").triggerTestId; + const destItems = await openMenu(page, destTrigger); + note("adminDestMenu", destItems); + const destTexts = (destItems || []).map((i) => i.text); + check("A1", "Destination menu offers Edit scope", destTexts.includes("Edit scope"), JSON.stringify(destTexts), 'includes "Edit scope"'); + check("A2", "Destination menu offers Delete", destTexts.includes("Delete"), JSON.stringify(destTexts), 'includes "Delete"'); + check("A3", "Destination menu offers no Test", !destTexts.includes("Test"), JSON.stringify(destTexts), 'no "Test"'); + check("A4", "Destination menu offers no plain Edit", !destTexts.includes("Edit"), JSON.stringify(destTexts), 'no bare "Edit"'); + check("A5", "Destination menu has exactly 2 items", (destItems || []).length === 2, JSON.stringify(destTexts), "2 items"); + await closeMenu(page); + + const cfgTrigger = ddCfg.triggerTestId; + const cfgItems = await openMenu(page, cfgTrigger); + note("adminCfgMenu", cfgItems); + const cfgTexts = (cfgItems || []).map((i) => i.text); + check("A6", "Config-callback menu offers Test", cfgTexts.includes("Test"), JSON.stringify(cfgTexts), 'includes "Test"'); + check("A7", "Config-callback menu offers Edit", cfgTexts.includes("Edit"), JSON.stringify(cfgTexts), 'includes "Edit"'); + check("A8", "Config-callback menu offers Delete", cfgTexts.includes("Delete"), JSON.stringify(cfgTexts), 'includes "Delete"'); + check("A9", "Config-callback menu offers no Edit scope", !cfgTexts.includes("Edit scope"), JSON.stringify(cfgTexts), 'no "Edit scope"'); + check("A10", "Config-callback menu has exactly 3 items", (cfgItems || []).length === 3, JSON.stringify(cfgTexts), "3 items"); + await closeMenu(page); + + // Every destination row has a trigger for a full admin + const allDestHaveTrigger = destRows.every((r) => r.hasTrigger); + check("A11", "Every destination row exposes an actions trigger for a full admin", allDestHaveTrigger, JSON.stringify(destRows.filter((r) => !r.hasTrigger).map((r) => r.name)), "all have triggers"); + + // Menu on a "Not active" destination still offers Edit scope + Delete + const naTrigger = destRows.find((r) => r.name === "uitrk-adapter-reject").triggerTestId; + const naItems = await openMenu(page, naTrigger); + const naTexts = (naItems || []).map((i) => i.text); + check("A12", 'A "Not active" destination still offers Edit scope + Delete', naTexts.includes("Edit scope") && naTexts.includes("Delete"), JSON.stringify(naTexts), '["Edit scope","Delete"]'); + await closeMenu(page); + + // ---------- Delete dialog ---------- + const delTrigger = destRows.find((r) => r.name === "uitrk-team-and-org").triggerTestId; + await openMenu(page, delTrigger); + await page.locator('[data-testid="destination-action-delete"]').click(); + await page.waitForTimeout(900); + const dlg = await page.evaluate(() => { + const dialogs = [...document.querySelectorAll('.ant-modal-wrap, [role="dialog"]')].filter((d) => d.offsetParent !== null || d.getBoundingClientRect().height > 0); + const d = dialogs[dialogs.length - 1]; + return d ? { text: d.innerText, html: d.innerHTML.length } : null; + }); + note("deleteDialog", dlg); + await page.screenshot({ path: "/tmp/uitrack_delete_dialog.png", fullPage: true }).catch(() => {}); + const dtext = (dlg && dlg.text) || ""; + check("D1", "Delete dialog is titled for a destination", dtext.includes("Delete Destination"), JSON.stringify(dtext.slice(0, 200)), '"Delete Destination"'); + check("D2", "Delete dialog names the destination", dtext.includes("uitrk-team-and-org"), JSON.stringify(dtext.slice(0, 400)), '"uitrk-team-and-org"'); + check("D3", "Delete dialog names the backend", dtext.includes("Langfuse OTEL"), JSON.stringify(dtext.slice(0, 400)), '"Langfuse OTEL · https://cloud.langfuse.com"'); + check("D4", "Delete dialog invents no Mode value", !/\bMode\b/.test(dtext) && !/\bsuccess\b/.test(dtext), JSON.stringify(dtext.slice(0, 400)), "no Mode row, no invented 'success'"); + check("D5", "Delete dialog warns stored collector credentials go with it", dtext.includes("stored collector credentials are deleted with it"), JSON.stringify(dtext.slice(0, 400)), "collector-credentials warning"); + check("D6", 'Delete dialog uses the "Destination Information" section title', dtext.includes("Destination Information"), JSON.stringify(dtext.slice(0, 400)), '"Destination Information"'); + check("D7", "Delete dialog says the action cannot be undone", dtext.includes("cannot be undone"), JSON.stringify(dtext.slice(0, 400)), '"cannot be undone"'); + // cancel + const cancelBtn = page.getByRole("button", { name: /^Cancel$/ }).last(); + await cancelBtn.click().catch(async () => { + await page.keyboard.press("Escape"); + }); + await page.waitForTimeout(800); + const afterCancel = await scrape(page); + check("D8", "Cancelling the delete dialog deletes nothing", afterCancel.rows.some((r) => r.name === "uitrk-team-and-org"), `${afterCancel.rows.length} rows`, "destination still present"); + + // Delete dialog for a config callback, for contrast + await openMenu(page, cfgTrigger); + await page.locator('[data-testid="callback-action-delete"]').click(); + await page.waitForTimeout(900); + const dlg2 = await page.evaluate(() => { + const dialogs = [...document.querySelectorAll('.ant-modal-wrap, [role="dialog"]')].filter((d) => d.getBoundingClientRect().height > 0); + const d = dialogs[dialogs.length - 1]; + return d ? d.innerText : null; + }); + note("deleteDialogCfg", dlg2); + check("D9", "Config-callback delete dialog is titled Delete Callback (not Destination)", !!dlg2 && dlg2.includes("Delete Callback") && !dlg2.includes("Delete Destination"), JSON.stringify((dlg2 || "").slice(0, 200)), '"Delete Callback"'); + check("D10", "Config-callback delete dialog does show a Mode", !!dlg2 && /Mode/.test(dlg2), JSON.stringify((dlg2 || "").slice(0, 300)), "Mode row present"); + const dlg2Mode = (dlg2 || "").split("\n").map((s) => s.trim()); + const modeIdx = dlg2Mode.indexOf("Mode"); + const dlg2ModeValue = modeIdx >= 0 ? dlg2Mode[modeIdx + 1] : null; + check( + "D11", + "Config-callback delete dialog's Mode agrees with the Mode the row renders", + !!dlg2ModeValue && dlg2ModeValue.toLowerCase().replace(/[^a-z]/g, "") === ddCfg.mode.toLowerCase().replace(/[^a-z]/g, ""), + `dialog Mode="${dlg2ModeValue}" but the row's Mode badge is "${ddCfg.mode}"`, + "the two agree", + ); + check("D12", "Config-callback delete dialog names the callback", !!dlg2 && dlg2.includes("datadog"), JSON.stringify((dlg2 || "").slice(0, 300)), '"datadog"'); + await page.getByRole("button", { name: /^Cancel$/ }).last().click().catch(async () => await page.keyboard.press("Escape")); + await page.waitForTimeout(600); + + // Delete dialog for a "Not active" destination + const naDelTrigger = destRows.find((r) => r.name === "uitrk-no-description").triggerTestId; + await openMenu(page, naDelTrigger); + await page.locator('[data-testid="destination-action-delete"]').click(); + await page.waitForTimeout(900); + const dlg3 = await page.evaluate(() => { + const dialogs = [...document.querySelectorAll('.ant-modal-wrap, [role="dialog"]')].filter((d) => d.getBoundingClientRect().height > 0); + const d = dialogs[dialogs.length - 1]; + return d ? d.innerText : null; + }); + note("deleteDialogNotActive", dlg3); + check("D13", 'Delete dialog for a "Not active" destination is still titled for a destination', !!dlg3 && dlg3.includes("Delete Destination"), JSON.stringify((dlg3 || "").slice(0, 200)), '"Delete Destination"'); + check("D14", 'Delete dialog for a destination with no backend renders a Backend row (placeholder "-")', !!dlg3 && /Backend/.test(dlg3), JSON.stringify((dlg3 || "").slice(0, 300)), "Backend row present"); + await page.getByRole("button", { name: /^Cancel$/ }).last().click().catch(async () => await page.keyboard.press("Escape")); + await page.waitForTimeout(600); + + // "Not active" badge carries an explanation + const naRow = destRows.find((r) => r.name === "uitrk-adapter-reject"); + check("S1", '"Not active" badge carries an explanatory tooltip', !!naRow && (naRow.scopeTitle || "").includes("cannot be built"), naRow ? JSON.stringify(naRow.scopeTitle) : "row missing", "tooltip explaining it receives no traces"); + const noDescRow = destRows.find((r) => r.name === "uitrk-no-description"); + check("S2", 'A destination with no backend renders a placeholder sub-line rather than a blank/undefined one', !!noDescRow && noDescRow.sub !== "" && !/undefined|null/i.test(noDescRow.sub), noDescRow ? JSON.stringify(noDescRow.sub) : "row missing", "a non-empty, non-undefined placeholder"); + + // The +N more collapse hides real grants (documented behavior) + const t5 = await apiGet(`/team/info?team_id=${teamIdByAlias["uitrk-team-5"]}`); + const t5names = ((t5.team_info || t5).resolved_logging_exporters || []); + const manyRow = destRows.find((r) => r.name === "uitrk-many-teams"); + check( + "S3", + 'The "+N more" collapse hides team-5/6 grants the backend does report (intended collapse, not a wrong verdict)', + t5names.includes("uitrk-many-teams") && manyRow.scope.includes("+2 more") && !manyRow.scope.includes("uitrk-team-5"), + `backend for uitrk-team-5=${JSON.stringify(t5names)} ui scope=${JSON.stringify(manyRow.scope)}`, + "backend grants it, UI collapses the last two behind +2 more", + ); + + // ---------- Edit scope dialog ---------- + const esTrigger = destRows.find((r) => r.name === "uitrk-many-teams").triggerTestId; + await openMenu(page, esTrigger); + await page.locator('[data-testid="destination-action-edit-access"]').click(); + await page.waitForTimeout(1200); + const es = await page.evaluate(() => { + const m = [...document.querySelectorAll(".ant-modal")].filter((d) => d.getBoundingClientRect().height > 0).pop(); + if (!m) return null; + return { + title: m.querySelector(".ant-modal-title") ? m.querySelector(".ant-modal-title").innerText.trim() : null, + text: m.innerText, + labels: [...m.querySelectorAll("label")].map((l) => l.innerText.trim()), + hasSwitch: !!m.querySelector(".ant-switch"), + selects: m.querySelectorAll(".ant-select").length, + selectedTags: [...m.querySelectorAll(".ant-select-selection-item")].map((e) => e.innerText.trim()), + }; + }); + note("editScopeDialog", es); + await page.screenshot({ path: "/tmp/uitrack_editscope_dialog.png", fullPage: true }).catch(() => {}); + check("E1", "Edit-scope dialog opens", !!es, String(!!es), "true"); + check("E2", "Edit-scope dialog names the destination", !!es && (es.title || "").includes("uitrk-many-teams"), es ? es.title : "no dialog", "Edit scope — uitrk-many-teams"); + check("E3", "Edit-scope dialog offers Global", !!es && es.labels.includes("Global"), es ? JSON.stringify(es.labels) : "no dialog", 'includes "Global"'); + check("E4", "Edit-scope dialog offers Teams", !!es && es.labels.includes("Teams"), es ? JSON.stringify(es.labels) : "no dialog", 'includes "Teams"'); + check("E5", "Edit-scope dialog offers Organizations", !!es && es.labels.includes("Organizations"), es ? JSON.stringify(es.labels) : "no dialog", 'includes "Organizations"'); + check("E6", "Edit-scope dialog has a Global toggle and two multi-selects", !!es && es.hasSwitch && es.selects >= 2, es ? `switch=${es.hasSwitch} selects=${es.selects}` : "no dialog", "switch + 2 selects"); + check("E7", "Edit-scope dialog pre-seeds the destination's current teams", !!es && es.selectedTags.filter((t) => t.startsWith("uitrk-team-")).length === 6, es ? JSON.stringify(es.selectedTags) : "no dialog", "6 team tags pre-selected"); + await page.getByRole("button", { name: /^Cancel$/ }).last().click().catch(async () => await page.keyboard.press("Escape")); + await page.waitForTimeout(600); + + // ---------- Navigation ---------- + await page.goto(`${BASE}/ui/teams`, { waitUntil: "domcontentloaded" }); + await page.waitForLoadState("networkidle"); + await page.waitForTimeout(1000); + const onTeams = page.url().includes("/ui/teams"); + check("N3", "Can navigate away to /ui/teams", onTeams, page.url(), "/ui/teams"); + await gotoDestinations(page); + const back = await scrape(page); + check("N4", "Destinations page survives navigating away and back", back.rows.filter((r) => r.mode === "—").length === DEST_COUNT, `${back.rows.filter((r) => r.mode === "—").length} destination rows`, String(DEST_COUNT)); + check("N5", "Scope verdicts are stable across the round trip", JSON.stringify(back.rows.map((r) => [r.name, r.scope])) === JSON.stringify(scraped.rows.map((r) => [r.name, r.scope])), "scope cells differ after round trip", "identical"); + + const tabClickable = await page.getByText("Logging Callbacks", { exact: true }).first().isVisible().catch(() => false); + check("N6", "Logging Callbacks tab still exposed after the round trip", tabClickable, String(tabClickable), "true"); + + // The legacy query form is not a route for this page + await page.goto(`${BASE}/ui?page=logging-callbacks`, { waitUntil: "domcontentloaded" }); + await page.waitForLoadState("networkidle"); + await page.waitForTimeout(1500); + const qHasTable = await page.getByText("Active Logging Callbacks", { exact: true }).first().isVisible().catch(() => false); + note("queryFormUrl", page.url()); + check("N7", "The query form ?page=logging-callbacks is NOT a route to this page (path route is the only one)", !qHasTable, `?page=logging-callbacks rendered the destinations table: ${qHasTable}`, "does not render the destinations table"); + await gotoDestinations(page); + + const addBtn = await page.getByRole("button", { name: "Add Callback" }).first().isVisible().catch(() => false); + check("A13", 'Full admin sees the "Add Callback" button', addBtn, String(addBtn), "true"); + + note("adminConsoleErrors", consoleErrors.slice(0, 30)); + check("C1", "No uncaught page errors while rendering the destinations table as admin", !consoleErrors.some((e) => e.startsWith("PAGEERROR")), JSON.stringify(consoleErrors.filter((e) => e.startsWith("PAGEERROR")).slice(0, 5)), "none"); + + await ctx.close(); + + // ---------- READ-ONLY ADMIN ---------- + consoleErrors = []; + const ctx2 = await browser.newContext({ viewport: { width: 1600, height: 1200 } }); + const p2 = await ctx2.newPage(); + p2.on("console", (m) => { + if (m.type() === "error") consoleErrors.push(m.text().slice(0, 300)); + }); + p2.on("pageerror", (e) => consoleErrors.push("PAGEERROR: " + String(e).slice(0, 300))); + + await login(p2, VIEWER_USER, VIEWER_PASS); + note("viewerUrlAfterLogin", p2.url()); + await gotoDestinations(p2); + const vScraped = await scrape(p2); + note("viewerScrape", vScraped); + await p2.screenshot({ path: "/tmp/uitrack_viewer_table.png", fullPage: true }).catch(() => {}); + + const vDest = vScraped.rows.filter((r) => r.mode === "—"); + const vCfg = vScraped.rows.filter((r) => r.mode !== "—"); + check("V1", "Read-only admin can load the destinations page", p2.url().includes("/ui/logging-and-alerts"), p2.url(), "/ui/logging-and-alerts"); + check("V2", "Read-only admin sees all 10 destinations", vDest.length === DEST_COUNT, `${vDest.length}: ${vDest.map((r) => r.name).join(",")}`, String(DEST_COUNT)); + check( + "V3", + "Read-only admin sees identical Scope verdicts to the full admin", + JSON.stringify(vDest.map((r) => [r.name, r.scope])) === JSON.stringify(destRows.map((r) => [r.name, r.scope])), + JSON.stringify(vDest.map((r) => [r.name, r.scope])), + JSON.stringify(destRows.map((r) => [r.name, r.scope])), + ); + const vAdd = await p2.getByRole("button", { name: "Add Callback" }).first().isVisible().catch(() => false); + check("V4", 'Read-only admin has no "Add Callback" button', !vAdd, String(vAdd), "false"); + const vDestWithTrigger = vDest.filter((r) => r.hasTrigger).map((r) => r.name); + check("V5", "Read-only admin gets NO actions trigger on any destination row", vDestWithTrigger.length === 0, JSON.stringify(vDestWithTrigger), "[]"); + const vCfgRow = vCfg.find((r) => r.name.toLowerCase() === "datadog"); + check("V6", "Read-only admin keeps the actions trigger on the config-callback row", !!vCfgRow && vCfgRow.hasTrigger, vCfgRow ? String(vCfgRow.hasTrigger) : "config row missing", "true"); + if (vCfgRow && vCfgRow.hasTrigger) { + const vItems = await openMenu(p2, vCfgRow.triggerTestId); + note("viewerCfgMenu", vItems); + const vTexts = (vItems || []).map((i) => i.text); + check("V7", "Read-only admin's config-callback menu offers Test", vTexts.includes("Test"), JSON.stringify(vTexts), '["Test"]'); + check("V8", "Read-only admin's config-callback menu has no Edit", !vTexts.includes("Edit"), JSON.stringify(vTexts), 'no "Edit"'); + check("V9", "Read-only admin's config-callback menu has no Delete", !vTexts.includes("Delete"), JSON.stringify(vTexts), 'no "Delete"'); + check("V10", "Read-only admin's config-callback menu has exactly 1 item", (vItems || []).length === 1, JSON.stringify(vTexts), "1"); + await closeMenu(p2); + } else { + check("V7", "Read-only admin's config-callback menu offers Test", false, "no trigger", "menu with Test"); + check("V8", "Read-only admin's config-callback menu has no Edit", false, "no trigger", "no Edit"); + check("V9", "Read-only admin's config-callback menu has no Delete", false, "no trigger", "no Delete"); + check("V10", "Read-only admin's config-callback menu has exactly 1 item", false, "no trigger", "1"); + } + check("V11", "Read-only admin sees the destination named datadog under its own name", vDest.some((r) => r.name === "datadog"), JSON.stringify(vDest.map((r) => r.name)), 'includes "datadog"'); + check("V12", "Read-only admin sees the same Not-active verdicts", JSON.stringify(vDest.filter((r) => r.scope.includes("Not active")).map((r) => r.name).sort()) === JSON.stringify(["uitrk-adapter-reject", "uitrk-no-description"]), JSON.stringify(vDest.filter((r) => r.scope.includes("Not active")).map((r) => r.name).sort()), '["uitrk-adapter-reject","uitrk-no-description"]'); + note("viewerConsoleErrors", consoleErrors.slice(0, 30)); + check("C2", "No uncaught page errors for the read-only admin", !consoleErrors.some((e) => e.startsWith("PAGEERROR")), JSON.stringify(consoleErrors.filter((e) => e.startsWith("PAGEERROR")).slice(0, 5)), "none"); + + await ctx2.close(); + await browser.close(); + + fs.writeFileSync(OUT, JSON.stringify({ results, notes }, null, 2)); + const pass = results.filter((r) => r.status === "PASS").length; + const fail = results.filter((r) => r.status === "FAIL").length; + console.log(`TOTAL ${results.length} PASS ${pass} FAIL ${fail}`); + for (const r of results) { + if (r.status === "FAIL") console.log(`FAIL ${r.id}: ${r.desc}\n actual: ${r.actual}\n expected: ${r.expected}`); + } +})().catch((e) => { + console.error("RUNNER CRASH", e); + fs.writeFileSync(OUT, JSON.stringify({ results, notes, crash: String(e && e.stack) }, null, 2)); + process.exit(1); +}); diff --git a/tests/e2e/ui/ui_track_smoke.mjs b/tests/e2e/ui/ui_track_smoke.mjs new file mode 100644 index 00000000000..ff4d5f887b6 --- /dev/null +++ b/tests/e2e/ui/ui_track_smoke.mjs @@ -0,0 +1,33 @@ +import { chromium } from "playwright"; + +const BASE = "http://127.0.0.1:21501"; +const EXEC = + "/Users/yucheng/Library/Caches/ms-playwright/chromium_headless_shell-1234/chrome-headless-shell-mac-arm64/chrome-headless-shell/chrome-headless-shell"; + +const browser = await chromium.launch({ + executablePath: + "/Users/yucheng/Library/Caches/ms-playwright/chromium_headless_shell-1234/chrome-headless-shell-mac-arm64/chrome-headless-shell", + headless: true, +}); +const ctx = await browser.newContext({ viewport: { width: 1600, height: 1200 } }); +const page = await ctx.newPage(); +const errs = []; +page.on("console", (m) => errs.push(`[${m.type()}] ${m.text().slice(0, 500)}`)); +page.on("pageerror", (e) => errs.push("PAGEERROR: " + (e && e.stack ? e.stack.slice(0, 1200) : String(e)))); + +await page.goto(`${BASE}/ui/login`, { waitUntil: "domcontentloaded" }); +await page.getByPlaceholder("Enter your username").fill("admin"); +await page.getByPlaceholder("Enter your password").fill("sk-uitrack-21501"); +await page.getByRole("button", { name: "Login", exact: true }).click(); +await page.waitForLoadState("networkidle"); +console.log("after login url:", page.url()); + +await page.goto(`${BASE}/ui/logging-and-alerts`, { waitUntil: "domcontentloaded" }); +await page.waitForLoadState("networkidle"); +await page.waitForTimeout(4000); +console.log("url:", page.url()); +console.log("body:", (await page.evaluate(() => document.body.innerText)).slice(0, 1500)); +console.log("tables:", await page.evaluate(() => document.querySelectorAll("table").length)); +console.log("---- console/pageerrors ----"); +for (const e of errs) console.log(e); +await browser.close(); diff --git a/tests/e2e/ui/ui_track_test_action.mjs b/tests/e2e/ui/ui_track_test_action.mjs new file mode 100644 index 00000000000..324c033577a --- /dev/null +++ b/tests/e2e/ui/ui_track_test_action.mjs @@ -0,0 +1,38 @@ +import { chromium } from "playwright"; + +const BASE = "http://127.0.0.1:21501"; +const EXEC = + "/Users/yucheng/Library/Caches/ms-playwright/chromium_headless_shell-1234/chrome-headless-shell-mac-arm64/chrome-headless-shell"; +const who = process.argv[2] || "viewer"; +const creds = who === "viewer" ? ["uitrk-viewer@example.com", "uitrk-viewer-pw"] : ["admin", "sk-uitrack-21501"]; + +const browser = await chromium.launch({ executablePath: EXEC, headless: true }); +const ctx = await browser.newContext({ viewport: { width: 1600, height: 1200 } }); +const page = await ctx.newPage(); +const calls = []; +page.on("response", (r) => { + if (r.url().includes("/health/services")) calls.push({ url: r.url(), status: r.status() }); +}); + +await page.goto(`${BASE}/ui/login`, { waitUntil: "domcontentloaded" }); +await page.getByPlaceholder("Enter your username").fill(creds[0]); +await page.getByPlaceholder("Enter your password").fill(creds[1]); +await page.getByRole("button", { name: "Login", exact: true }).click(); +await page.waitForLoadState("networkidle"); +await page.goto(`${BASE}/ui/logging-and-alerts`, { waitUntil: "domcontentloaded" }); +await page.waitForLoadState("networkidle"); +await page.waitForTimeout(3000); + +await page.locator('[data-testid="callback-actions-datadog-success_and_failure"]').click(); +await page.waitForTimeout(500); +await page.locator('[data-testid="callback-action-test"]').click(); +await page.waitForTimeout(4000); +console.log(who, "health/services calls:", JSON.stringify(calls)); +const toast = await page.evaluate(() => + [...document.querySelectorAll('[class*="toast"], .ant-message, [role="status"], [role="alert"]')] + .map((e) => e.innerText.trim()) + .filter(Boolean) + .slice(0, 5), +); +console.log(who, "toast:", JSON.stringify(toast)); +await browser.close(); diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 6abc40eb28e..f430dcd462c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -12,9 +12,7 @@ from fastapi.testclient import TestClient from litellm._uuid import uuid -sys.path.insert( - 0, os.path.abspath("../../../") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../")) # Adds the parent directory to the system path from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( LiteLLM_BudgetTableFull, @@ -78,9 +76,7 @@ mock_prisma_client.db.litellm_teamtable.update = AsyncMock() # Fixture to provide the mock prisma client @pytest.fixture(autouse=True) def mock_db_client(): - with patch( - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ): # Mock in both places if necessary + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): # Mock in both places if necessary yield mock_prisma_client mock_prisma_client.reset_mock() @@ -118,27 +114,17 @@ async def test_validate_team_org_change_same_org_id(): organization.organization_id = org_id organization.models = [] organization.litellm_budget_table = MagicMock() - organization.litellm_budget_table.max_budget = ( - 50.0 # This would normally fail validation - ) - organization.litellm_budget_table.tpm_limit = ( - 500 # This would normally fail validation - ) - organization.litellm_budget_table.rpm_limit = ( - 50 # This would normally fail validation - ) + organization.litellm_budget_table.max_budget = 50.0 # This would normally fail validation + organization.litellm_budget_table.tpm_limit = 500 # This would normally fail validation + organization.litellm_budget_table.rpm_limit = 50 # This would normally fail validation organization.members = [] # Mock Router mock_router = MagicMock(spec=Router) # Use patch to ensure the model access check is never called - with patch( - "litellm.proxy.management_endpoints.team_endpoints.can_org_access_model" - ) as mock_access_check: - result = validate_team_org_change( - team=team, organization=organization, llm_router=mock_router - ) + with patch("litellm.proxy.management_endpoints.team_endpoints.can_org_access_model") as mock_access_check: + result = validate_team_org_change(team=team, organization=organization, llm_router=mock_router) # Assert the function returns True without checking anything assert result is True @@ -190,9 +176,7 @@ async def test_validate_team_org_change_members_in_org(): mock_router = MagicMock(spec=Router) # Test should pass - all team members are in org members - result = validate_team_org_change( - team=team, organization=organization, llm_router=mock_router - ) + result = validate_team_org_change(team=team, organization=organization, llm_router=mock_router) assert result is True @@ -244,9 +228,7 @@ async def test_validate_team_org_change_member_not_in_org(): # Test should fail - user_id_not_in_org is not in org members with pytest.raises(HTTPException) as exc_info: - validate_team_org_change( - team=team, organization=organization, llm_router=mock_router - ) + validate_team_org_change(team=team, organization=organization, llm_router=mock_router) assert exc_info.value.status_code == 403 assert "not a member of the organization" in str(exc_info.value.detail) @@ -290,10 +272,7 @@ async def test_get_team_permissions_list_success(mock_db_client, mock_admin_auth assert response.status_code == 200 response_data = response.json() assert response_data["team_id"] == test_team_id - assert ( - response_data["team_member_permissions"] - == mock_team_data["team_member_permissions"] - ) + assert response_data["team_member_permissions"] == mock_team_data["team_member_permissions"] assert ( response_data["all_available_permissions"] == TeamMemberPermissionChecks.get_all_available_team_member_permissions() @@ -356,9 +335,7 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): return_value=mock_existing_team_row, ): # Mock the database update function - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team_row - ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team_row) # Override the dependency for this test app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth @@ -455,17 +432,13 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): mock_db_client.db = MagicMock() # Mock object permission table creation - mock_object_perm_create = AsyncMock( - return_value=MagicMock(object_permission_id="objperm123") - ) + mock_object_perm_create = AsyncMock(return_value=MagicMock(object_permission_id="objperm123")) mock_db_client.db.litellm_objectpermissiontable = MagicMock() mock_db_client.db.litellm_objectpermissiontable.create = mock_object_perm_create # Mock model table creation mock_db_client.db.litellm_modeltable = MagicMock() - mock_db_client.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Capture team table creation team_create_result = MagicMock( @@ -481,9 +454,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create mock_db_client.db.litellm_teamtable.count = mock_team_count - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() @@ -552,9 +523,7 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut # Mock model table mock_db_client.db.litellm_modeltable = MagicMock() - mock_db_client.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model456") - ) + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model456")) # Mock team table team_create_result = MagicMock( @@ -566,13 +535,9 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut "object_permission_id": "objperm_team_mcp_456", } mock_db_client.db.litellm_teamtable = MagicMock() - mock_db_client.db.litellm_teamtable.create = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() @@ -631,20 +596,14 @@ def test_should_auto_add_team_creator(user_role, user_id, flag_value, expected): _should_auto_add_team_creator, ) - general_settings = ( - {} if flag_value is None else {"disable_auto_add_proxy_admin_to_teams": flag_value} - ) + general_settings = {} if flag_value is None else {"disable_auto_add_proxy_admin_to_teams": flag_value} auth = UserAPIKeyAuth(user_role=user_role, user_id=user_id) assert _should_auto_add_team_creator(auth, general_settings) is expected @pytest.mark.asyncio -@pytest.mark.parametrize( - "disable_flag,expect_creator_added", [(True, False), (False, True)] -) -async def test_new_team_disable_auto_add_proxy_admin_flag( - mock_db_client, disable_flag, expect_creator_added -): +@pytest.mark.parametrize("disable_flag,expect_creator_added", [(True, False), (False, True)]) +async def test_new_team_disable_auto_add_proxy_admin_flag(mock_db_client, disable_flag, expect_creator_added): """ When general_settings.disable_auto_add_proxy_admin_to_teams is True, a proxy admin calling /team/new must NOT be auto-added to the team's members. When @@ -659,9 +618,7 @@ async def test_new_team_disable_auto_add_proxy_admin_flag( team_create_result = MagicMock(team_id="team-789") team_create_result.model_dump.return_value = {"team_id": "team-789"} mock_db_client.db.litellm_teamtable = MagicMock() - mock_db_client.db.litellm_teamtable.create = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -671,17 +628,18 @@ async def test_new_team_disable_auto_add_proxy_admin_flag( from litellm.proxy._types import NewTeamRequest from litellm.proxy.management_endpoints.team_endpoints import new_team - admin_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user-1" - ) + admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user-1") - with patch( - "litellm.proxy.proxy_server.general_settings", - {"disable_auto_add_proxy_admin_to_teams": disable_flag}, - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", - new_callable=AsyncMock, - ) as mock_add_members: + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"disable_auto_add_proxy_admin_to_teams": disable_flag}, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + new_callable=AsyncMock, + ) as mock_add_members, + ): await new_team( data=NewTeamRequest(team_alias="flag-test-team"), http_request=MagicMock(spec=Request), @@ -730,16 +688,12 @@ async def test_team_update_object_permissions_existing_permission(monkeypatch): "vector_stores": ["old_store_1", "old_store_2"], } - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_object_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_object_permission) # Mock upsert operation updated_permission = MagicMock() updated_permission.object_permission_id = "existing_perm_id_123" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=updated_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=updated_permission) # Test data with new object permission data_json = { @@ -795,21 +749,17 @@ async def test_team_update_object_permissions_no_existing_permission(monkeypatch ) # Mock find_unique to return None (no existing permission) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "new_perm_id_456" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["brand_new_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["brand_new_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "team_alias": "updated_team_2", } @@ -855,21 +805,17 @@ async def test_team_update_object_permissions_missing_permission_record(monkeypa ) # Mock find_unique to return None (permission record not found) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "recreated_perm_id_789" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["recreated_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["recreated_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "team_alias": "updated_team_3", } @@ -975,14 +921,10 @@ async def test_add_team_member_budget_table_success(): mock_budget_record.budget_id = "budget-123" mock_budget_record.max_budget = 1000.0 - mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=mock_budget_record - ) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=mock_budget_record) # Create team info response object - team_info_response = TeamInfoResponseObjectTeamTable( - team_id="test-team-123", team_alias="Test Team" - ) + team_info_response = TeamInfoResponseObjectTeamTable(team_id="test-team-123", team_alias="Test Team") # Call the function result = await _add_team_member_budget_table( @@ -993,15 +935,11 @@ async def test_add_team_member_budget_table_success(): # Verify the result assert result.team_member_budget_table == mock_budget_record - assert result == team_info_response.model_copy( - update={"team_member_budget_table": mock_budget_record} - ) + assert result == team_info_response.model_copy(update={"team_member_budget_table": mock_budget_record}) assert team_info_response.team_member_budget_table is None # Verify database call was made correctly - mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( - where={"budget_id": "budget-123"} - ) + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with(where={"budget_id": "budget-123"}) @pytest.mark.asyncio @@ -1021,14 +959,10 @@ async def test_add_team_member_budget_table_exception_handling(): ) # Create team info response object - team_info_response = TeamInfoResponseObjectTeamTable( - team_id="test-team-456", team_alias="Test Team 2" - ) + team_info_response = TeamInfoResponseObjectTeamTable(team_id="test-team-456", team_alias="Test Team 2") # Mock the verbose_proxy_logger to capture log calls - with patch( - "litellm.proxy.management_endpoints.team_endpoints.verbose_proxy_logger" - ) as mock_logger: + with patch("litellm.proxy.management_endpoints.team_endpoints.verbose_proxy_logger") as mock_logger: # Call the function result = await _add_team_member_budget_table( team_member_budget_id="nonexistent-budget-456", @@ -1040,10 +974,7 @@ async def test_add_team_member_budget_table_exception_handling(): assert result == team_info_response # Verify team_member_budget_table is not set when exception occurs - assert ( - not hasattr(result, "team_member_budget_table") - or result.team_member_budget_table is None - ) + assert not hasattr(result, "team_member_budget_table") or result.team_member_budget_table is None # Verify the error was logged mock_logger.info.assert_called_once_with( @@ -1072,9 +1003,7 @@ async def test_add_team_member_budget_table_budget_not_found(): mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) # Create team info response object - team_info_response = TeamInfoResponseObjectTeamTable( - team_id="test-team-789", team_alias="Test Team 3" - ) + team_info_response = TeamInfoResponseObjectTeamTable(team_id="test-team-789", team_alias="Test Team 3") # Call the function result = await _add_team_member_budget_table( @@ -1291,9 +1220,7 @@ async def test_available_team_self_join_blocks_other_user_id(): await _validate_team_member_add_permissions( user_api_key_dict=user, complete_team_data=team, - data=_make_team_member_add_request( - member_user_id="bob-victim", role="user" - ), + data=_make_team_member_add_request(member_user_id="bob-victim", role="user"), ) assert exc_info.value.status_code == 403 @@ -1731,9 +1658,7 @@ async def test_update_team_members_list_duplicate_prevention(): # Create mock team with existing members mock_team = MagicMock(spec=LiteLLM_TeamTable) - mock_team.members_with_roles = [ - Member(user_id="existing-user", user_email="existing@example.com", role="admin") - ] + mock_team.members_with_roles = [Member(user_id="existing-user", user_email="existing@example.com", role="admin")] # Try to add the same member again duplicate_member = Member(user_id="existing-user", role="user") @@ -1882,9 +1807,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): ) mock_request = Mock(spec=Request) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") existing_team = MagicMock() existing_team.model_dump.return_value = { @@ -1921,12 +1844,8 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): new_callable=AsyncMock, ) as mock_cache_team, ): - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=updated_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) mock_prisma_client.db.execute_raw = AsyncMock(return_value=None) if endpoint_name == "team_model_add": @@ -1968,15 +1887,11 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): # "no team-level restriction" and stop enforcing the team's # search-tool allowlist on key issuance. assert call_kwargs["team_table"].object_permission is not None - assert call_kwargs["team_table"].object_permission.search_tools == [ - "allowed-tool-A" - ] + assert call_kwargs["team_table"].object_permission.search_tools == ["allowed-tool-A"] # Pin the Prisma call shape too — the regression is in *what the # update returns*, so the contract that the update asks for # `object_permission` belongs in this test. - update_call_kwargs = ( - mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs - ) + update_call_kwargs = mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs assert update_call_kwargs.get("include", {}).get("object_permission") is True @@ -1998,9 +1913,7 @@ async def test_update_team_team_member_budget_not_passed_to_db(): # Mock dependencies mock_request = Mock(spec=Request) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") with ( patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, @@ -2008,9 +1921,7 @@ async def test_update_team_team_member_budget_not_passed_to_db(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" - ) as mock_cache_team, + patch("litellm.proxy.management_endpoints.team_endpoints._cache_team_object") as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget, @@ -2022,20 +1933,14 @@ async def test_update_team_team_member_budget_not_passed_to_db(): "team_alias": "test_team", "metadata": {"team_member_budget_id": "budget_123"}, } - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock the update return value mock_updated_team = MagicMock() mock_updated_team.team_id = "test_team_id" mock_updated_team.model_dump.return_value = {"team_id": "test_team_id"} - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) - mock_prisma_client.jsonify_team_object = MagicMock( - side_effect=lambda db_data: db_data - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) + mock_prisma_client.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Mock budget upsert to return updated_kv without team_member_budget def mock_upsert_side_effect( @@ -2073,14 +1978,14 @@ async def test_update_team_team_member_budget_not_passed_to_db(): update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data - assert ( - "team_member_budget" not in update_data - ), f"team_member_budget should not be in update data, but found: {update_data}" + assert "team_member_budget" not in update_data, ( + f"team_member_budget should not be in update data, but found: {update_data}" + ) # Verify other fields are present (team_alias should be there) - assert "team_alias" in update_data or "team_id" in str( - call_args - ), "Expected team update fields should be present" + assert "team_alias" in update_data or "team_id" in str(call_args), ( + "Expected team update fields should be present" + ) # Reset mock for second test mock_prisma_client.db.litellm_teamtable.update.reset_mock() @@ -2106,9 +2011,9 @@ async def test_update_team_team_member_budget_not_passed_to_db(): update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data - assert ( - "team_member_budget" not in update_data - ), f"team_member_budget should not be in update data, but found: {update_data}" + assert "team_member_budget" not in update_data, ( + f"team_member_budget should not be in update data, but found: {update_data}" + ) # Test Case 3: No team_member_budget field at all (excluded from request) mock_prisma_client.db.litellm_teamtable.update.reset_mock() @@ -2133,14 +2038,12 @@ async def test_update_team_team_member_budget_not_passed_to_db(): update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data - assert ( - "team_member_budget" not in update_data - ), f"team_member_budget should not be in update data, but found: {update_data}" - - print( - "✅ All test cases passed: team_member_budget is properly excluded from database update operations" + assert "team_member_budget" not in update_data, ( + f"team_member_budget should not be in update data, but found: {update_data}" ) + print("✅ All test cases passed: team_member_budget is properly excluded from database update operations") + def test_clean_team_member_fields(): """ @@ -2202,9 +2105,7 @@ async def test_create_team_member_budget_table(): TeamMemberBudgetHandler, ) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") data = NewTeamRequest( team_id="test_team_id", @@ -2271,9 +2172,7 @@ async def test_create_team_member_budget_table_without_team_alias(): TeamMemberBudgetHandler, ) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") data = NewTeamRequest(team_id="test_team_id") new_team_data_json = { @@ -2317,9 +2216,7 @@ async def test_upsert_team_member_budget_table_existing_budget(): TeamMemberBudgetHandler, ) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") team_table = MagicMock(spec=LiteLLM_TeamTable) team_table.metadata = {"team_member_budget_id": "existing_budget_123"} @@ -2378,9 +2275,7 @@ async def test_upsert_team_member_budget_table_no_existing_budget(): TeamMemberBudgetHandler, ) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") team_table = MagicMock(spec=LiteLLM_TeamTable) team_table.metadata = {} @@ -2431,9 +2326,7 @@ async def test_update_team_with_team_member_budget_duration(): from litellm.proxy.management_endpoints.team_endpoints import update_team mock_request = Mock(spec=Request) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id") with ( patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, @@ -2441,9 +2334,7 @@ async def test_update_team_with_team_member_budget_duration(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" - ) as mock_cache_team, + patch("litellm.proxy.management_endpoints.team_endpoints._cache_team_object") as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget, @@ -2455,19 +2346,13 @@ async def test_update_team_with_team_member_budget_duration(): "metadata": {"team_member_budget_id": "budget_123"}, } mock_existing_team.metadata = {"team_member_budget_id": "budget_123"} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_updated_team = MagicMock() mock_updated_team.team_id = "test_team_id" mock_updated_team.model_dump.return_value = {"team_id": "test_team_id"} - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) - mock_prisma_client.jsonify_team_object = MagicMock( - side_effect=lambda db_data: db_data - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) + mock_prisma_client.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) def mock_upsert_side_effect( team_table, @@ -2534,9 +2419,7 @@ async def test_backfill_team_member_budget_entries_creates_missing_memberships() existing_membership.user_id = "user-A" mock_prisma = MagicMock() - mock_prisma.db.litellm_teammembership.find_many = AsyncMock( - return_value=[existing_membership] - ) + mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[existing_membership]) mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) mock_prisma.db.litellm_teammembership.update_many = AsyncMock(return_value=0) @@ -2554,9 +2437,7 @@ async def test_backfill_team_member_budget_entries_creates_missing_memberships() ) # find_many should have been called to fetch existing memberships - mock_prisma.db.litellm_teammembership.find_many.assert_awaited_once_with( - where={"team_id": team_id} - ) + mock_prisma.db.litellm_teammembership.find_many.assert_awaited_once_with(where={"team_id": team_id}) # create_many should only create an entry for user-B (user-A already has one) mock_prisma.db.litellm_teammembership.create_many.assert_awaited_once_with( @@ -2609,9 +2490,7 @@ async def test_backfill_team_member_budget_entries_no_op_when_all_exist(): existing_b.user_id = "user-B" mock_prisma = MagicMock() - mock_prisma.db.litellm_teammembership.find_many = AsyncMock( - return_value=[existing_a, existing_b] - ) + mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[existing_a, existing_b]) mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) mock_prisma.db.litellm_teammembership.update_many = AsyncMock(return_value=0) @@ -2656,9 +2535,7 @@ async def test_backfill_team_member_budget_entries_populates_null_budget_id_on_e existing_b.user_id = "user-B" mock_prisma = MagicMock() - mock_prisma.db.litellm_teammembership.find_many = AsyncMock( - return_value=[existing_a, existing_b] - ) + mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[existing_a, existing_b]) mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) mock_prisma.db.litellm_teammembership.update_many = AsyncMock(return_value=2) @@ -2847,9 +2724,7 @@ async def test_bulk_team_member_add_batch_size_limit(): from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add # Create more than 500 members (the max batch size) - large_member_list = [ - Member(user_email=f"user{i}@example.com", role="user") for i in range(501) - ] + large_member_list = [Member(user_email=f"user{i}@example.com", role="user") for i in range(501)] bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", @@ -2904,9 +2779,7 @@ async def test_bulk_team_member_add_all_users_flag(): ) as mock_team_member_add, ): # Mock the database find_many call - mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=mock_db_users - ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=mock_db_users) mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @@ -2916,9 +2789,7 @@ async def test_bulk_team_member_add_all_users_flag(): ) # Verify that find_many was called to get all users - mock_prisma.db.litellm_usertable.find_many.assert_called_once_with( - order={"created_at": "desc"} - ) + mock_prisma.db.litellm_usertable.find_many.assert_called_once_with(order={"created_at": "desc"}) # Verify team_member_add was called with users from database mock_team_member_add.assert_called_once() @@ -3041,9 +2912,7 @@ async def test_list_team_v2_security_check_non_admin_user(): ) assert exc_info.value.status_code == 401 - assert "Only admin users can query all teams/other teams" in str( - exc_info.value.detail - ) + assert "Only admin users can query all teams/other teams" in str(exc_info.value.detail) assert LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.detail) @@ -3091,9 +2960,7 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): ) assert exc_info.value.status_code == 401 - assert "Only admin users can query all teams/other teams" in str( - exc_info.value.detail - ) + assert "Only admin users can query all teams/other teams" in str(exc_info.value.detail) @pytest.mark.asyncio @@ -3243,17 +3110,11 @@ async def test_list_team_v2_with_status_deleted(): mock_prisma_client.db = mock_db # Mock deleted teams - mock_deleted_team1 = Mock( - model_dump=lambda: {"team_id": "team_1", "team_alias": "Deleted Team 1"} - ) - mock_deleted_team2 = Mock( - model_dump=lambda: {"team_id": "team_2", "team_alias": "Deleted Team 2"} - ) + mock_deleted_team1 = Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Deleted Team 1"}) + mock_deleted_team2 = Mock(model_dump=lambda: {"team_id": "team_2", "team_alias": "Deleted Team 2"}) # Mock deleted teams table (should be called) - mock_db.litellm_deletedteamtable.find_many = AsyncMock( - return_value=[mock_deleted_team1, mock_deleted_team2] - ) + mock_db.litellm_deletedteamtable.find_many = AsyncMock(return_value=[mock_deleted_team1, mock_deleted_team2]) mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=2) # Mock regular teams table (should NOT be called) @@ -3442,9 +3303,7 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): "organization_id": "org_A", "members_with_roles": [{"user_id": "other_user", "role": "user"}], } - mock_db.litellm_teamtable.find_many = AsyncMock( - return_value=[mock_team_1, mock_team_2] - ) + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team_1, mock_team_2]) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) @@ -3540,10 +3399,7 @@ async def test_list_team_v2_org_admin_cannot_view_other_orgs(): ) assert exc_info.value.status_code == 403 - assert ( - "only view teams within your organizations" - in str(exc_info.value.detail).lower() - ) + assert "only view teams within your organizations" in str(exc_info.value.detail).lower() @pytest.mark.asyncio @@ -3703,9 +3559,7 @@ async def test_list_team_v2_search_builds_or_clause(): from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 mock_request = Mock(spec=Request) - mock_admin = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" - ) + mock_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user") with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: mock_db = Mock() @@ -3750,9 +3604,7 @@ async def test_list_team_v2_search_team_id_match_prefix(): from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 mock_request = Mock(spec=Request) - mock_admin = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" - ) + mock_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user") with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: mock_db = Mock() @@ -3803,9 +3655,7 @@ async def test_list_team_v2_search_composes_with_user_id_filter(): from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 mock_request = Mock(spec=Request) - mock_user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="member_user" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="member_user") mock_user = LiteLLM_UserTable( user_id="member_user", @@ -3990,9 +3840,7 @@ async def test_list_team_v2_keys_count_skipped_for_deleted_status(): "team_alias": "Deleted Team", } - mock_db.litellm_deletedteamtable.find_many = AsyncMock( - return_value=[mock_deleted] - ) + mock_db.litellm_deletedteamtable.find_many = AsyncMock(return_value=[mock_deleted]) mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=1) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) @@ -4025,9 +3873,7 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a mock_team_row = MagicMock() mock_team_row.model_dump.return_value = { "team_id": test_team_id, - "members_with_roles": [ - {"user_id": test_user_id, "user_email": None, "role": "user"} - ], + "members_with_roles": [{"user_id": test_user_id, "user_email": None, "role": "user"}], "team_member_permissions": [], "metadata": {}, "models": [], @@ -4035,32 +3881,24 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a } # Configure DB mocks used by team_member_delete - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team_row - ) + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team_row) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) # User row to allow removal from user's teams list mock_user_row = MagicMock() mock_user_row.user_id = test_user_id mock_user_row.teams = [test_team_id] - mock_db_client.db.litellm_usertable.find_many = AsyncMock( - return_value=[mock_user_row] - ) + mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[mock_user_row]) mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) # Membership deletion should be called mock_db_client.db.litellm_teammembership = MagicMock() - mock_db_client.db.litellm_teammembership.delete_many = AsyncMock( - return_value=MagicMock() - ) + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) # Verification token deletion should be called mock_db_client.db.litellm_verificationtoken = MagicMock() mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock( - return_value=MagicMock() - ) + mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) # Execute await team_member_delete( @@ -4075,9 +3913,7 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a @pytest.mark.asyncio -async def test_team_member_delete_cleans_verification_tokens( - mock_db_client, mock_admin_auth -): +async def test_team_member_delete_cleans_verification_tokens(mock_db_client, mock_admin_auth): from litellm.proxy._types import TeamMemberDeleteRequest from litellm.proxy.management_endpoints.team_endpoints import team_member_delete @@ -4087,38 +3923,28 @@ async def test_team_member_delete_cleans_verification_tokens( mock_team_row = MagicMock() mock_team_row.model_dump.return_value = { "team_id": test_team_id, - "members_with_roles": [ - {"user_id": test_user_id, "user_email": None, "role": "user"} - ], + "members_with_roles": [{"user_id": test_user_id, "user_email": None, "role": "user"}], "team_member_permissions": [], "metadata": {}, "models": [], "spend": 0.0, } - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team_row - ) + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team_row) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) mock_user_row = MagicMock() mock_user_row.user_id = test_user_id mock_user_row.teams = [test_team_id] - mock_db_client.db.litellm_usertable.find_many = AsyncMock( - return_value=[mock_user_row] - ) + mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[mock_user_row]) mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) mock_db_client.db.litellm_teammembership = MagicMock() - mock_db_client.db.litellm_teammembership.delete_many = AsyncMock( - return_value=MagicMock() - ) + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) mock_db_client.db.litellm_verificationtoken = MagicMock() mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock( - return_value=MagicMock() - ) + mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) await team_member_delete( data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id), @@ -4166,9 +3992,7 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4196,9 +4020,7 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): # ProxyException stores status_code in 'code' attribute assert exc_info.value.code == "400" assert "max budget higher than user max" in str(exc_info.value.message) - assert "100.0" in str( - exc_info.value.message - ) # User's user_max_budget should be mentioned + assert "100.0" in str(exc_info.value.message) # User's user_max_budget should be mentioned assert LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.message) @@ -4235,9 +4057,7 @@ async def test_new_team_max_budget_within_user_limit(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4269,18 +4089,12 @@ async def test_new_team_max_budget_within_user_limit(): "max_budget": 50.0, "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock( - return_value=mock_created_team - ) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_created_team - ) + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) # Mock model table mock_prisma.db.litellm_modeltable = MagicMock() - mock_prisma.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations for adding the creator as a member mock_user = MagicMock() @@ -4303,9 +4117,7 @@ async def test_new_team_max_budget_within_user_limit(): "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.create = AsyncMock( - return_value=mock_membership - ) + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception result = await new_team( @@ -4367,12 +4179,8 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, - patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, + patch("litellm.proxy.management_endpoints.team_endpoints.get_org_object") as mock_get_org, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4412,18 +4220,12 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): "organization_id": "test-org-123", "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock( - return_value=mock_created_team - ) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_created_team - ) + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) # Mock model table mock_prisma.db.litellm_modeltable = MagicMock() - mock_prisma.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() @@ -4446,9 +4248,7 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.create = AsyncMock( - return_value=mock_membership - ) + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception - the fix should bypass user budget validation for org-scoped teams result = await new_team( @@ -4500,9 +4300,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): # Create team request with models that are within org's allowed models but not user's team_request = NewTeamRequest( team_alias="org-scoped-models-team", - models=[ - "gpt-4" - ], # Within org's allowed models, but not in user's personal models + models=["gpt-4"], # Within org's allowed models, but not in user's personal models organization_id="test-org-456", # This makes it an org-scoped team ) @@ -4513,12 +4311,8 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, - patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, + patch("litellm.proxy.management_endpoints.team_endpoints.get_org_object") as mock_get_org, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4560,18 +4354,12 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): "models": ["gpt-4"], "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock( - return_value=mock_created_team - ) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_created_team - ) + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) # Mock model table mock_prisma.db.litellm_modeltable = MagicMock() - mock_prisma.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() @@ -4594,9 +4382,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.create = AsyncMock( - return_value=mock_membership - ) + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception - the fix should bypass user model validation for org-scoped teams result = await new_team( @@ -4656,9 +4442,7 @@ async def test_new_team_standalone_validates_against_user_models(monkeypatch): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4725,9 +4509,7 @@ async def test_new_team_standalone_validates_against_user_budget(): patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4752,9 +4534,7 @@ async def test_new_team_standalone_validates_against_user_budget(): # Verify exception details assert exc_info.value.code == "400" assert "max budget higher than user max" in str(exc_info.value.message) - assert "3.0" in str( - exc_info.value.message - ) # User's max_budget should be mentioned + assert "3.0" in str(exc_info.value.message) # User's max_budget should be mentioned @pytest.mark.asyncio @@ -4799,12 +4579,8 @@ async def test_new_team_org_scoped_budget_exceeds_org_limit(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, - patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, + patch("litellm.proxy.management_endpoints.team_endpoints.get_org_object") as mock_get_org, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4879,12 +4655,8 @@ async def test_new_team_org_scoped_models_not_in_org_models(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, - patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, + patch("litellm.proxy.management_endpoints.team_endpoints.get_org_object") as mock_get_org, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4908,10 +4680,7 @@ async def test_new_team_org_scoped_models_not_in_org_models(): # Verify exception details assert exc_info.value.code == "400" - assert ( - "claude-3-opus" in str(exc_info.value.message) - or "organization" in str(exc_info.value.message).lower() - ) + assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() @pytest.mark.asyncio @@ -4955,9 +4724,7 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" @@ -4968,13 +4735,9 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin(): "team_id": "standalone-team-123", "organization_id": None, "max_budget": 30.0, - "members_with_roles": [ - {"user_id": "non-admin-update-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "non-admin-update-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_cache.async_get_cache = AsyncMock(return_value=None) with pytest.raises(ProxyException) as exc_info: @@ -5021,9 +4784,7 @@ async def test_update_team_standalone_budget_raise_allowed_for_proxy_admin(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" @@ -5034,13 +4795,9 @@ async def test_update_team_standalone_budget_raise_allowed_for_proxy_admin(): "team_id": "standalone-team-123", "organization_id": None, "max_budget": 30.0, - "members_with_roles": [ - {"user_id": "proxy-admin-update-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "proxy-admin-update-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() @@ -5055,9 +4812,7 @@ async def test_update_team_standalone_budget_raise_allowed_for_proxy_admin(): "organization_id": None, "max_budget": 100.0, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -5109,9 +4864,7 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" @@ -5122,13 +4875,9 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin(): "team_id": "standalone-team-123", "organization_id": None, "max_budget": 500.0, - "members_with_roles": [ - {"user_id": "budget-removal-admin", "role": "admin"} - ], + "members_with_roles": [{"user_id": "budget-removal-admin", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_cache.async_get_cache = AsyncMock(return_value=None) with pytest.raises(ProxyException) as exc_info: @@ -5176,9 +4925,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-uncapped-123" @@ -5189,13 +4936,9 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(): "team_id": "standalone-uncapped-123", "organization_id": None, "max_budget": None, - "members_with_roles": [ - {"user_id": "uncapped-team-admin", "role": "admin"} - ], + "members_with_roles": [{"user_id": "uncapped-team-admin", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() @@ -5210,9 +4953,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(): "organization_id": None, "max_budget": 1000.0, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -5268,9 +5009,7 @@ async def test_update_team_standalone_unchanged_budget_allowed(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Mock existing standalone team (no organization_id) with budget=$500 mock_existing_team = MagicMock() @@ -5282,13 +5021,9 @@ async def test_update_team_standalone_unchanged_budget_allowed(): "team_id": "standalone-unchanged-budget-123", "organization_id": None, "max_budget": 500.0, - "members_with_roles": [ - {"user_id": "standalone-unchanged-budget-admin", "role": "admin"} - ], + "members_with_roles": [{"user_id": "standalone-unchanged-budget-admin", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data # User has a restrictive personal budget that is lower than the team's. @@ -5310,9 +5045,7 @@ async def test_update_team_standalone_unchanged_budget_allowed(): "max_budget": 500.0, "tpm_limit": 50000, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise - unchanged budget skips the personal-budget check. result = await update_team( @@ -5364,9 +5097,7 @@ async def test_update_team_standalone_lower_budget_allowed(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-lower-budget-123" @@ -5377,13 +5108,9 @@ async def test_update_team_standalone_lower_budget_allowed(): "team_id": "standalone-lower-budget-123", "organization_id": None, "max_budget": 500.0, - "members_with_roles": [ - {"user_id": "standalone-lower-budget-admin", "role": "admin"} - ], + "members_with_roles": [{"user_id": "standalone-lower-budget-admin", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_user_obj = LiteLLM_UserTable( @@ -5403,9 +5130,7 @@ async def test_update_team_standalone_lower_budget_allowed(): "organization_id": None, "max_budget": 300.0, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -5467,9 +5192,7 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -5484,13 +5207,9 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): "team_id": "org-team-456", "organization_id": "test-org-update", "max_budget": 80.0, - "members_with_roles": [ - {"user_id": "org-admin-update-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because new budget exceeds org's max_budget with pytest.raises(ProxyException) as exc_info: @@ -5502,10 +5221,7 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): # Verify exception details assert exc_info.value.code == "400" - assert ( - "organization" in str(exc_info.value.message).lower() - or "budget" in str(exc_info.value.message).lower() - ) + assert "organization" in str(exc_info.value.message).lower() or "budget" in str(exc_info.value.message).lower() @pytest.mark.asyncio @@ -5545,9 +5261,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() @@ -5559,13 +5273,9 @@ async def test_update_team_standalone_models_not_gated_by_user_limit(): "team_id": "standalone-team-models-123", "organization_id": None, "models": ["gpt-3.5-turbo"], - "members_with_roles": [ - {"user_id": "non-admin-update-models-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "non-admin-update-models-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() @@ -5579,9 +5289,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit(): "organization_id": None, "models": ["gpt-4"], } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -5643,9 +5351,7 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -5661,13 +5367,9 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): "team_id": "org-team-update-budget-123", "organization_id": "test-org-update-budget", "max_budget": 30.0, - "members_with_roles": [ - {"user_id": "org-admin-update-budget-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-budget-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data # Mock user cache to return user with restrictive budget @@ -5676,9 +5378,7 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): max_budget=3.0, # Restrictive personal budget ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) - mock_cache.async_set_cache = ( - AsyncMock() - ) # Mock cache set for _cache_team_object + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() @@ -5691,9 +5391,7 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): "organization_id": "test-org-update-budget", "max_budget": 50.0, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - bypass user budget validation for org-scoped teams result = await update_team( @@ -5753,9 +5451,7 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -5771,17 +5467,11 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(): "team_id": "org-team-update-models-123", "organization_id": "test-org-update-models", "models": ["gpt-3.5-turbo"], - "members_with_roles": [ - {"user_id": "org-admin-update-models-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-models-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data - mock_cache.async_set_cache = ( - AsyncMock() - ) # Mock cache set for _cache_team_object + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() @@ -5794,9 +5484,7 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(): "organization_id": "test-org-update-models", "models": ["gpt-4"], } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - bypass user models validation for org-scoped teams result = await update_team( @@ -5856,9 +5544,7 @@ async def test_update_team_org_scoped_models_not_in_org_models(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -5873,13 +5559,9 @@ async def test_update_team_org_scoped_models_not_in_org_models(): "team_id": "org-team-update-models-fail-123", "organization_id": "test-org-update-models-fail", "models": ["gpt-4"], - "members_with_roles": [ - {"user_id": "org-admin-update-models-fail-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-models-fail-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because claude-3-opus is not in org's allowed models with pytest.raises(ProxyException) as exc_info: @@ -5891,10 +5573,7 @@ async def test_update_team_org_scoped_models_not_in_org_models(): # Verify exception details assert exc_info.value.code == "400" - assert ( - "claude-3-opus" in str(exc_info.value.message) - or "organization" in str(exc_info.value.message).lower() - ) + assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() @pytest.mark.asyncio @@ -5944,9 +5623,7 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -5962,17 +5639,11 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): "team_id": "org-team-all-proxy-models-123", "organization_id": "test-org-all-proxy-models", "models": ["gpt-4"], - "members_with_roles": [ - {"user_id": "org-admin-all-proxy-models-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-all-proxy-models-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data - mock_cache.async_set_cache = ( - AsyncMock() - ) # Mock cache set for _cache_team_object + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() @@ -5993,9 +5664,7 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): "gpt-4o-mini-test", ], } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - 'all-proxy-models' allows all models result = await update_team( @@ -6051,9 +5720,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): # Mock existing standalone team mock_existing_team = MagicMock() @@ -6067,9 +5734,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit(): "tpm_limit": 500, "members_with_roles": [{"user_id": "tpm-limit-user", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() @@ -6083,9 +5748,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit(): "organization_id": None, "tpm_limit": 5000, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -6131,9 +5794,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), ): # Mock existing standalone team mock_existing_team = MagicMock() @@ -6147,9 +5808,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit(): "rpm_limit": 50, "members_with_roles": [{"user_id": "rpm-limit-user", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() @@ -6163,9 +5822,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit(): "organization_id": None, "rpm_limit": 500, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) result = await update_team( data=update_request, @@ -6385,9 +6042,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), patch("litellm.proxy.proxy_server._license_check") as mock_license, - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org), @@ -6417,12 +6072,8 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): "metadata": None, "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock( - return_value=mock_created_team - ) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_created_team - ) + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Should succeed - bypasses user limits since org-scoped @@ -6502,13 +6153,9 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit(): "team_id": "org-team-update-tpm-123", "organization_id": "test-org-update-tpm", "tpm_limit": 5000, - "members_with_roles": [ - {"user_id": "org-admin-update-tpm-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-tpm-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because TPM exceeds org limit with pytest.raises(ProxyException) as exc_info: @@ -6589,13 +6236,9 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit(): "team_id": "org-team-update-rpm-123", "organization_id": "test-org-update-rpm", "rpm_limit": 500, - "members_with_roles": [ - {"user_id": "org-admin-update-rpm-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-rpm-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because RPM exceeds org limit with pytest.raises(ProxyException) as exc_info: @@ -6683,13 +6326,9 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): "organization_id": "test-org-update-bypass", "tpm_limit": 5000, "rpm_limit": 500, - "members_with_roles": [ - {"user_id": "org-admin-update-bypass-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-update-bypass-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_cache.async_set_cache = AsyncMock() mock_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() @@ -6703,9 +6342,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): "tpm_limit": 10000, "rpm_limit": 1000, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Should succeed - bypasses user limits since org-scoped @@ -6791,9 +6428,7 @@ async def test_update_team_guardrails_with_org_id(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), patch( "litellm.proxy.proxy_server.premium_user", True, # Required for guardrails feature @@ -6818,20 +6453,14 @@ async def test_update_team_guardrails_with_org_id(): "max_budget": None, "tpm_limit": None, "rpm_limit": None, - "members_with_roles": [ - {"user_id": "org-admin-guardrails-test", "role": "admin"} - ], + "members_with_roles": [{"user_id": "org-admin-guardrails-test", "role": "admin"}], } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_cache.async_set_cache = AsyncMock() # Mock organization fetch - this is where the bug occurred # The fix ensures 'teams: True' is in the include clause - mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock( - return_value=mock_org - ) + mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(return_value=mock_org) # Destination-org guard in update_team queries for the caller's # ORG_ADMIN membership on the destination org. Return a match so @@ -6839,26 +6468,20 @@ async def test_update_team_guardrails_with_org_id(): mock_org_admin_membership = MagicMock() mock_org_admin_membership.user_id = "org-admin-guardrails-test" mock_org_admin_membership.organization_id = "test-org-guardrails" - mock_prisma.db.litellm_organizationmembership.find_many = AsyncMock( - return_value=[mock_org_admin_membership] - ) + mock_prisma.db.litellm_organizationmembership.find_many = AsyncMock(return_value=[mock_org_admin_membership]) # Mock team update mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team.team_id = "team-guardrails-123" mock_updated_team.organization_id = "test-org-guardrails" - mock_updated_team.metadata = { - "guardrails": ["aporia-pre-call", "aporia-post-call"] - } + mock_updated_team.metadata = {"guardrails": ["aporia-pre-call", "aporia-post-call"]} mock_updated_team.litellm_model_table = None mock_updated_team.model_dump.return_value = { "team_id": "team-guardrails-123", "organization_id": "test-org-guardrails", "metadata": {"guardrails": ["aporia-pre-call", "aporia-post-call"]}, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # async_get_cache must be an AsyncMock so `await` in get_org_object works mock_cache.async_get_cache = AsyncMock(return_value=None) @@ -6888,11 +6511,7 @@ async def test_update_team_guardrails_with_org_id(): assert mock_prisma.db.litellm_organizationtable.find_unique.call_count >= 1 # Get the first call (from fetch_and_validate_organization) - first_call_kwargs = ( - mock_prisma.db.litellm_organizationtable.find_unique.call_args_list[ - 0 - ].kwargs - ) + first_call_kwargs = mock_prisma.db.litellm_organizationtable.find_unique.call_args_list[0].kwargs # Verify that 'teams' is included in the fetch assert "include" in first_call_kwargs @@ -6943,9 +6562,7 @@ def test_transform_teams_to_deleted_records(): assert all("litellm_changed_by" in record for record in records) assert all(record["deleted_by"] == "user-123" for record in records) # UserAPIKeyAuth hashes the api_key, so we check against the hashed value - assert all( - record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records - ) + assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records) assert all(record["litellm_changed_by"] == "admin-user" for record in records) record1 = records[0] @@ -7087,9 +6704,7 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many_teams mock_create_many_keys = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many_keys - ) + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = mock_create_many_keys mock_find_many_keys = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys @@ -7223,9 +6838,7 @@ async def test_team_member_delete_persists_deleted_keys(monkeypatch): mock_prisma_client.db.litellm_verificationtoken.delete_many = mock_delete_keys mock_create_many_keys = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many_keys - ) + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = mock_create_many_keys monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", @@ -7379,9 +6992,7 @@ async def test_new_team_soft_budget_validation( patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server._license_check") as mock_license, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -7415,18 +7026,12 @@ async def test_new_team_soft_budget_validation( "max_budget": expected_max_budget, "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock( - return_value=mock_created_team - ) - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_created_team - ) + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) # Mock model table mock_prisma.db.litellm_modeltable = MagicMock() - mock_prisma.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() @@ -7449,9 +7054,7 @@ async def test_new_team_soft_budget_validation( "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.create = AsyncMock( - return_value=mock_membership - ) + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) if should_succeed: # Should NOT raise an exception @@ -7579,9 +7182,7 @@ async def test_update_team_soft_budget_validation( patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()) as mock_audit, ): # Mock existing team with existing budgets mock_existing_team = MagicMock() @@ -7595,9 +7196,7 @@ async def test_update_team_soft_budget_validation( "soft_budget": existing_soft_budget, "max_budget": existing_max_budget, } - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock user cache mock_user_obj = LiteLLM_UserTable( @@ -7607,14 +7206,8 @@ async def test_update_team_soft_budget_validation( mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Mock updated team - preserve existing values if not being updated - final_soft_budget = ( - update_soft_budget - if update_soft_budget is not None - else existing_soft_budget - ) - final_max_budget = ( - update_max_budget if update_max_budget is not None else existing_max_budget - ) + final_soft_budget = update_soft_budget if update_soft_budget is not None else existing_soft_budget + final_max_budget = update_max_budget if update_max_budget is not None else existing_max_budget mock_updated_team = MagicMock() mock_updated_team.team_id = "test-team-123" @@ -7627,13 +7220,9 @@ async def test_update_team_soft_budget_validation( "soft_budget": final_soft_budget, "max_budget": final_max_budget, } - mock_prisma.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) mock_prisma.jsonify_team_object = lambda db_data: db_data - mock_cache.async_set_cache = ( - AsyncMock() - ) # Mock cache set for _cache_team_object + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object if should_succeed: # Should NOT raise an exception @@ -7679,9 +7268,7 @@ async def test_new_team_positive_budgets_accepted(): from litellm.proxy._types import NewTeamRequest # Should not raise any errors - request = NewTeamRequest( - team_alias="test-team", max_budget=100.0, team_member_budget=50.0 - ) + request = NewTeamRequest(team_alias="test-team", max_budget=100.0, team_member_budget=50.0) assert request.max_budget == 100.0 assert request.team_member_budget == 50.0 @@ -7702,9 +7289,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): # Mock model table creation mock_db_client.db.litellm_modeltable = MagicMock() - mock_db_client.db.litellm_modeltable.create = AsyncMock( - return_value=MagicMock(id="model123") - ) + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Capture team table creation team_create_result = MagicMock( @@ -7718,9 +7303,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create mock_db_client.db.litellm_teamtable.count = mock_team_count - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() @@ -7782,9 +7365,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Create a non-admin user user_id = "test_user_123" team_id = "test_team_456" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -7816,9 +7397,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[user_api_key_1, user_api_key_2] - ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[user_api_key_1, user_api_key_2]) # Mock get_user_object with patch( @@ -7855,9 +7434,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Verify user's API keys were fetched mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() - api_key_call_kwargs = ( - mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] - ) + api_key_call_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] assert api_key_call_kwargs["where"] == {"user_id": user_id} @@ -7874,9 +7451,7 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) # Create a team admin user user_id = "test_admin_123" team_id = "test_team_456" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -7960,9 +7535,7 @@ async def test_get_team_daily_activity_member_with_permission_sees_all_spend( # Create a non-admin user user_id = "test_user_with_perm_123" team_id = "test_team_789" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -8029,9 +7602,7 @@ async def test_get_team_daily_activity_member_with_permission_sees_all_spend( hasattr(mock_db_client.db.litellm_verificationtoken, "find_many") and mock_db_client.db.litellm_verificationtoken.find_many.called ): - assert ( - False - ), "API keys should not be fetched for members with /team/daily/activity permission" + assert False, "API keys should not be fetched for members with /team/daily/activity permission" @pytest.mark.asyncio @@ -8049,9 +7620,7 @@ async def test_get_team_daily_activity_member_without_permission_filters_by_keys # Create a non-admin user user_id = "test_user_no_perm_123" team_id = "test_team_789" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -8085,9 +7654,7 @@ async def test_get_team_daily_activity_member_without_permission_filters_by_keys # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[user_api_key_1, user_api_key_2] - ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[user_api_key_1, user_api_key_2]) # Mock get_user_object with patch( @@ -8219,9 +7786,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Create a non-admin user user_id = "test_user_123" team_id = "test_team_456" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -8253,9 +7818,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[user_api_key_1, user_api_key_2] - ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[user_api_key_1, user_api_key_2]) # Mock get_user_object with patch( @@ -8292,9 +7855,7 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Verify user's API keys were fetched mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() - api_key_call_kwargs = ( - mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] - ) + api_key_call_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] assert api_key_call_kwargs["where"] == {"user_id": user_id} @@ -8311,9 +7872,7 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) # Create a team admin user user_id = "test_admin_123" team_id = "test_team_456" - user_api_key_dict = UserAPIKeyAuth( - user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER - ) + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) # Mock user info mock_user_info = LiteLLM_UserTable( @@ -8438,9 +7997,7 @@ async def test_validate_and_populate_member_user_info_only_email_provided(): mock_user_find_first.user_email = "test@example.com" # Mock find_first to return the user - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=mock_user_find_first - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user_find_first) # Mock get_data to return single user (no duplicates) mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) @@ -8496,9 +8053,7 @@ async def test_validate_and_populate_member_user_info_only_user_id_not_found(): assert result.role == "user" # Verify find_unique was called with correct parameters - mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( - where={"user_id": "nonexistent-user"} - ) + mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with(where={"user_id": "nonexistent-user"}) @pytest.mark.asyncio @@ -8594,9 +8149,7 @@ async def test_list_team_v1_batches_key_queries(): return [key3] return [key1, key2, key3] - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - side_effect=filtered_find_many - ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=filtered_find_many) result = await list_team( http_request=mock_request, @@ -8693,9 +8246,7 @@ class TestBatchResolveAccessGroupResources: fake_row.access_agent_ids = ["agent-1", "agent-2"] fake_prisma = MagicMock() - fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[fake_row] - ) + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[fake_row]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1"]) @@ -8724,9 +8275,7 @@ class TestBatchResolveAccessGroupResources: row2.access_agent_ids = ["agent-2"] fake_prisma = MagicMock() - fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[row1, row2] - ) + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1, row2]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) @@ -8748,9 +8297,7 @@ class TestBatchResolveAccessGroupResources: row1.access_agent_ids = [] fake_prisma = MagicMock() - fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[row1] - ) + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1", "ag-missing"]) @@ -8788,9 +8335,7 @@ class TestBatchResolveAccessGroupResources: fake_prisma.db.litellm_accessgrouptable.find_many = fake_find_many with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): - result = await _batch_resolve_access_group_resources( - ["ag-1", "ag-1", "ag-1"] - ) + result = await _batch_resolve_access_group_resources(["ag-1", "ag-1", "ag-1"]) # Should have been called with deduplicated list call_args = fake_find_many.call_args @@ -8827,9 +8372,7 @@ class TestResolveTeamAccessGroupResources: row2.access_agent_ids = ["agent-1"] fake_prisma = MagicMock() - fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[row1, row2] - ) + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1, row2]) team_info = TeamInfoResponseObjectTeamTable( team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"] @@ -8845,10 +8388,7 @@ class TestResolveTeamAccessGroupResources: ] assert resolved.access_group_mcp_server_ids == ["mcp-1"] assert resolved.access_group_agent_ids == ["agent-1"] - assert [ - (d.access_group_id, d.access_group_name, d.models) - for d in (resolved.access_group_details or []) - ] == [ + assert [(d.access_group_id, d.access_group_name, d.models) for d in (resolved.access_group_details or [])] == [ ("ag-1", "shared-models", ("gpt-4", "claude-3")), ("ag-2", "extra-models", ("claude-3", "gemini")), ] @@ -8941,9 +8481,7 @@ async def test_update_team_rejects_unauthorized_caller(): ], "organization_id": "org-456", } - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) update_request = UpdateTeamRequest( team_id="team-123", @@ -9025,9 +8563,7 @@ async def test_team_member_me_returns_caller_membership(mock_db_client): team_id = "team-me-1" caller_id = "alice@example.com" other_id = "bob@example.com" - caller_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id - ) + caller_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id) team = _build_team_for_me( team_id, @@ -9039,9 +8575,7 @@ async def test_team_member_me_returns_caller_membership(mock_db_client): membership = _build_membership_for_me(caller_id, team_id, spend=42.0) user = LiteLLM_UserTable(user_id=caller_id, user_email=caller_id, max_budget=None) - p_team, p_membership, p_user = _patch_member_me_helpers( - team=team, membership=membership, user=user - ) + p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=user) with p_team, p_membership as mock_get_membership, p_user: response = await team_member_me( http_request=MagicMock(spec=Request), @@ -9059,9 +8593,7 @@ async def test_team_member_me_returns_caller_membership(mock_db_client): # budget_reset_at must survive end-to-end — proves the BudgetTableFull # variant of the Union is selected (created_at is present), not the base # LiteLLM_BudgetTable which would silently strip this field. - assert response.litellm_budget_table.budget_reset_at == datetime( - 2026, 5, 1, tzinfo=timezone.utc - ) + assert response.litellm_budget_table.budget_reset_at == datetime(2026, 5, 1, tzinfo=timezone.utc) # Membership lookup must scope to caller_id, not just team_id — proves the # endpoint cannot return another member's row. @@ -9096,13 +8628,9 @@ async def test_team_member_me_matches_email_only_member(mock_db_client): [{"user_id": None, "user_email": caller_email, "role": "user"}], ) membership = _build_membership_for_me(caller_id, team_id, spend=7.0) - user = LiteLLM_UserTable( - user_id=caller_id, user_email=caller_email, max_budget=None - ) + user = LiteLLM_UserTable(user_id=caller_id, user_email=caller_email, max_budget=None) - p_team, p_membership, p_user = _patch_member_me_helpers( - team=team, membership=membership, user=user - ) + p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=user) with p_team, p_membership, p_user: response = await team_member_me( http_request=MagicMock(spec=Request), @@ -9124,9 +8652,7 @@ async def test_team_member_me_returns_404_for_non_member(mock_db_client): team_id = "team-me-2" caller_id = "outsider@example.com" - caller_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id - ) + caller_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id) team = _build_team_for_me( team_id, @@ -9145,9 +8671,7 @@ async def test_team_member_me_returns_404_for_non_member(mock_db_client): @pytest.mark.asyncio -async def test_team_member_me_returns_404_for_proxy_admin_not_in_team( - mock_db_client, mock_admin_auth -): +async def test_team_member_me_returns_404_for_proxy_admin_not_in_team(mock_db_client, mock_admin_auth): """ Proxy admins get 404 if they are not actually a member of the team. `me` only resolves for actual team members; admins use /team/info instead. @@ -9187,9 +8711,7 @@ async def test_team_member_me_returns_defaults_when_no_membership_row(mock_db_cl team_id = "team-me-4" caller_id = "newmember@example.com" - caller_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id - ) + caller_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id) team = _build_team_for_me( team_id, @@ -9235,18 +8757,12 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client): from litellm.proxy.management_endpoints.team_endpoints import team_member_me - caller_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice@example.com" - ) + caller_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice@example.com") # get_team_object raises 404 directly when the team is missing. with patch( "litellm.proxy.management_endpoints.team_endpoints.get_team_object", - AsyncMock( - side_effect=HTTPException( - status_code=404, detail={"error": "Team doesn't exist in db."} - ) - ), + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})), ): with pytest.raises(HTTPException) as exc_info: await team_member_me( @@ -9258,9 +8774,7 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client): @pytest.mark.asyncio -async def test_new_team_encrypts_callback_vars( - mock_db_client, mock_admin_auth, monkeypatch -): +async def test_new_team_encrypts_callback_vars(mock_db_client, mock_admin_auth, monkeypatch): """/team/new must encrypt callback_vars values before they reach the DB.""" from fastapi import Request @@ -9275,9 +8789,7 @@ async def test_new_team_encrypts_callback_vars( # actual JSON serialization production uses (catches non-serializable # ciphertext, missing fields, etc.). mock_db_client.jsonify_object = PrismaClient.jsonify_object.__get__(mock_db_client) - mock_db_client.jsonify_team_object = PrismaClient.jsonify_team_object.__get__( - mock_db_client - ) + mock_db_client.jsonify_team_object = PrismaClient.jsonify_team_object.__get__(mock_db_client) mock_db_client.get_data = AsyncMock(return_value=None) mock_db_client.db = MagicMock() mock_db_client.db.litellm_teamtable = MagicMock() @@ -9286,9 +8798,7 @@ async def test_new_team_encrypts_callback_vars( mock_team_create = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_teamtable.create = mock_team_create mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=team_create_result - ) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -9325,9 +8835,7 @@ async def test_new_team_encrypts_callback_vars( def _non_admin_auth(): - return UserAPIKeyAuth( - user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER - ) + return UserAPIKeyAuth(user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER) def test_check_passthrough_routes_caller_permission_team(): @@ -9343,12 +8851,8 @@ def test_check_passthrough_routes_caller_permission_team(): NewTeamRequest(allowed_passthrough_routes=["/foo/*"]), admin, entity="team" ) - _check_passthrough_routes_caller_permission( - NewTeamRequest(), non_admin, entity="team" - ) - _check_passthrough_routes_caller_permission( - NewTeamRequest(allowed_passthrough_routes=[]), non_admin, entity="team" - ) + _check_passthrough_routes_caller_permission(NewTeamRequest(), non_admin, entity="team") + _check_passthrough_routes_caller_permission(NewTeamRequest(allowed_passthrough_routes=[]), non_admin, entity="team") with pytest.raises(HTTPException) as exc: _check_passthrough_routes_caller_permission( @@ -9385,9 +8889,7 @@ async def test_new_team_blocks_non_admin_passthrough_routes(mock_db_client): ): with pytest.raises(ProxyException) as exc: await new_team( - data=NewTeamRequest( - team_alias="t", allowed_passthrough_routes=["/admin/*"] - ), + data=NewTeamRequest(team_alias="t", allowed_passthrough_routes=["/admin/*"]), http_request=MagicMock(spec=Request), user_api_key_dict=_non_admin_auth(), ) @@ -9414,9 +8916,7 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): ): with pytest.raises(ProxyException) as exc: await update_team( - data=UpdateTeamRequest( - team_id="t1", allowed_passthrough_routes=["/admin/*"] - ), + data=UpdateTeamRequest(team_id="t1", allowed_passthrough_routes=["/admin/*"]), http_request=MagicMock(spec=Request), user_api_key_dict=_non_admin_auth(), ) @@ -9760,16 +9260,12 @@ async def test_team_info_forwards_key_limit_to_get_data(): from litellm.proxy.management_endpoints import team_endpoints mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( - return_value=LiteLLM_TeamTable(team_id="team-1") - ) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")) mock_prisma.get_data = AsyncMock(return_value=[]) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object( - team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[]) - ), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), ): await team_endpoints.team_info( http_request=MagicMock(spec=Request), @@ -9807,9 +9303,7 @@ async def test_team_info_returns_model_aliases(): with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object( - team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[]) - ), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), ): response = await team_endpoints.team_info( http_request=MagicMock(spec=Request), @@ -9825,6 +9319,42 @@ async def test_team_info_returns_model_aliases(): assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} +@pytest.mark.asyncio +@pytest.mark.parametrize("access_group_ids", [None, ["ag-1"]]) +async def test_team_info_discloses_logging_exporters_for_access_group_teams(monkeypatch, access_group_ids): + """Regression: the disclosure must survive access-group resolution. + + ``_resolve_team_access_group_resources`` returns the same object when the team has + no access groups but a ``model_copy`` when it does, and the response is built from + that return value. Writing the names onto the original therefore reached the caller + only for teams without access groups; a team that inherits from one disclosed null. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable(team_id="team-1", access_group_ids=access_group_ids) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + monkeypatch.setattr(team_endpoints, "resolved_logging_exporter_names", lambda *_: ("dest-a", "dest-b")) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), + patch.object(team_endpoints, "_batch_resolve_access_group_resources", AsyncMock(return_value={})), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response["team_info"].resolved_logging_exporters == ("dest-a", "dest-b") + + @pytest.mark.asyncio async def test_update_model_table_clears_aliases_with_empty_map(): """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) @@ -9833,12 +9363,8 @@ async def test_update_model_table_clears_aliases_with_empty_map(): """ mock_prisma = MagicMock() mock_prisma.db.litellm_modeltable.create = AsyncMock() - mock_prisma.db.litellm_modeltable.upsert = AsyncMock( - return_value=MagicMock(id="model-123") - ) - user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" - ) + mock_prisma.db.litellm_modeltable.upsert = AsyncMock(return_value=MagicMock(id="model-123")) + user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") returned_model_id = await _update_model_table( data=UpdateTeamRequest(team_id="team-1", model_aliases={}), @@ -9886,9 +9412,7 @@ class TestEmitTeamMembersMetric: return LiteLLM_TeamTable( team_id="team-x", team_alias="X", - members_with_roles=[ - Member(user_id=f"u{i}", role="user") for i in range(member_count) - ], + members_with_roles=[Member(user_id=f"u{i}", role="user") for i in range(member_count)], ) def test_emits_with_team_when_logger_registered(self, restore_callbacks): @@ -9962,9 +9486,7 @@ async def test_new_team_rejects_reserved_ui_session_team_id(): await new_team( data=team_request, http_request=dummy_request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) assert exc_info.value.code == "400" @@ -10042,9 +9564,7 @@ async def _drive_team_write( new=AsyncMock(), ), ): - pc.db.litellm_teamtable.find_unique = AsyncMock( - return_value=None if find_returns_none else existing - ) + pc.db.litellm_teamtable.find_unique = AsyncMock(return_value=None if find_returns_none else existing) pc.db.litellm_teamtable.update = AsyncMock( return_value=LiteLLM_TeamTable(team_id=_PATCH_TEAM_ID, team_alias="t") ) @@ -10145,9 +9665,7 @@ _METADATA_MAPPING = [ _METADATA_MAPPING, ids=[row[0] for row in _METADATA_MAPPING], ) -async def test_post_vs_patch_metadata_write_mapping( - label, existing_metadata, body, expected_post, expected_patch -): +async def test_post_vs_patch_metadata_write_mapping(label, existing_metadata, body, expected_post, expected_patch): """Exhaustive map: POST replaces metadata wholesale, PATCH merges per RFC 7386.""" post_meta = await _written_metadata("post", existing_metadata, body) patch_meta = await _written_metadata("patch", existing_metadata, body) @@ -10275,9 +9793,7 @@ async def test_patch_team_not_found_returns_404(): # metadata present -> patch_team does its own existence check with pytest.raises(ProxyException) as exc: - await _drive_team_write( - "patch", raw_body={"metadata": {"cost_center": "1"}}, find_returns_none=True - ) + await _drive_team_write("patch", raw_body={"metadata": {"cost_center": "1"}}, find_returns_none=True) assert exc.value.code == "404" or exc.value.code == 404 # metadata absent -> existence check happens in the delegated update_team @@ -10294,9 +9810,7 @@ async def test_patch_enforces_team_access_via_delegation(): outsider = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="outsider") with pytest.raises(ProxyException) as exc: - await _drive_team_write( - "patch", raw_body={"tpm_limit": 5}, user=outsider - ) + await _drive_team_write("patch", raw_body={"tpm_limit": 5}, user=outsider) assert exc.value.code == "403" or exc.value.code == 403 @@ -10306,9 +9820,7 @@ async def test_patch_returns_full_team_object_not_wrapper(): {"team_id", "data"} envelope.""" from litellm.proxy._types import LiteLLM_TeamTable - result, _ = await _drive_team_write( - "patch", existing_metadata={"a": 1}, raw_body={"metadata": {"b": 2}} - ) + result, _ = await _drive_team_write("patch", existing_metadata={"a": 1}, raw_body={"metadata": {"b": 2}}) assert isinstance(result, LiteLLM_TeamTable) assert result.team_id == _PATCH_TEAM_ID @@ -10622,9 +10134,7 @@ def test_patch_body_reshaping_adds_no_keys_the_caller_did_not_send(body): @pytest.mark.asyncio async def test_patch_ignores_unknown_body_keys(): """Unknown keys were silently dropped by the previous construction; keep that.""" - _, update_mock = await _drive_team_write( - "patch", raw_body={"tpm_limit": 5, "not_a_team_field": "x"} - ) + _, update_mock = await _drive_team_write("patch", raw_body={"tpm_limit": 5, "not_a_team_field": "x"}) written = update_mock.call_args.kwargs["data"] assert written["tpm_limit"] == 5