diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx
index 0c353621654..9362cbb3260 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx
@@ -1,5 +1,6 @@
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { render, renderHook, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
import { Form } from "antd";
import type { UploadProps } from "antd/es/upload";
import { describe, expect, it, vi } from "vitest";
@@ -8,6 +9,14 @@ import type { CredentialItem } from "../networking";
import { Providers } from "../provider_info_helpers";
import AddModelTab from "./add_model_tab";
+vi.mock("../molecules/models/ProviderLogo", () => ({
+ ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => (
+
+ {provider}
+
+ ),
+}));
+
vi.mock("../networking", async () => {
const actual = await vi.importActual("../networking");
return {
@@ -258,4 +267,47 @@ describe("Add Model Tab", () => {
{ timeout: 10000 },
);
}, 15000); // 15 second timeout to allow waitFor to complete
+
+ it("should show team selection when team-only switch is enabled", async () => {
+ const props = createTestProps();
+ const queryClient = createQueryClient();
+
+ render(
+
+
+ ,
+ );
+
+ // Wait for component to load
+ await screen.findByText("Provider");
+
+ // Find the team-BYOK switch by its role
+ const teamSwitch = screen.getByRole("switch");
+ expect(teamSwitch).toBeInTheDocument();
+
+ // Initially, team selection should not be visible
+ expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+ // Click the switch to enable team-only mode
+ await userEvent.click(teamSwitch!);
+
+ // Now team selection should be visible
+ expect(await screen.findByText("Select Team")).toBeInTheDocument();
+ });
});
diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
index b2e1dec2827..cfe00af5071 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx
@@ -15,6 +15,7 @@ import {
tagListCall,
} from "../networking";
import { Providers, providerLogoMap } from "../provider_info_helpers";
+import { ProviderLogo } from "../molecules/models/ProviderLogo";
import { Tag } from "../tag_management/types";
import AddAutoRouterTab from "./add_auto_router_tab";
import { TEST_MODES } from "./add_model_modes";
@@ -191,6 +192,7 @@ const AddModelTab: React.FC = ({
labelAlign="left"
>
= ({
return (
- {logoSrc ? (
-

{
- const target = e.currentTarget as HTMLImageElement;
- const parent = target.parentElement;
- if (!parent || !parent.contains(target)) {
- return;
- }
-
- try {
- const fallbackDiv = document.createElement("div");
- fallbackDiv.className =
- "w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs";
- fallbackDiv.textContent = displayName.charAt(0);
- parent.replaceChild(fallbackDiv, target);
- } catch (error) {
- console.error("Failed to replace provider logo fallback:", error);
- }
- }}
- />
- ) : (
-
- {displayName.charAt(0)}
-
- )}
+
{displayName}