From 5ead81786dea2322b614118d0f4eac7c7843e3f0 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 1 Mar 2025 17:42:50 -0800
Subject: [PATCH] test_can_team_access_model

---
 tests/proxy_unit_tests/test_auth_checks.py | 36 +++++++++++++---------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py
index ec36823633b..0eb1a387558 100644
--- a/tests/proxy_unit_tests/test_auth_checks.py
+++ b/tests/proxy_unit_tests/test_auth_checks.py
@@ -394,7 +394,7 @@ async def test_virtual_key_max_budget_check(
 
 
 @pytest.mark.parametrize(
-    "model, team_models, expected_result",
+    "model, team_models, expect_to_work",
     [
         ("gpt-4", ["gpt-4"], True),  # exact match
         ("gpt-4", ["all-proxy-models"], True),  # all-proxy-models access
@@ -427,7 +427,7 @@ async def test_virtual_key_max_budget_check(
     ],
 )
 @pytest.mark.asyncio
-async def test_can_team_access_model(model, team_models, expected_result):
+async def test_can_team_access_model(model, team_models, expect_to_work):
     """
     Test cases for can_team_access_model:
     1. Exact model match
@@ -438,18 +438,26 @@ async def test_can_team_access_model(model, team_models, expected_result):
     6. Empty model list
     7. None model list
     """
-    team_object = LiteLLM_TeamTable(
-        team_id="test-team",
-        models=team_models,
-    )
-
-    result = await can_team_access_model(
-        model=model,
-        team_object=team_object,
-        llm_router=None,
-        team_model_aliases=None,
-    )
-    assert result == expected_result
+    try:
+        team_object = LiteLLM_TeamTable(
+            team_id="test-team",
+            models=team_models,
+        )
+        result = await can_team_access_model(
+            model=model,
+            team_object=team_object,
+            llm_router=None,
+            team_model_aliases=None,
+        )
+        if not expect_to_work:
+            pytest.fail(
+                f"Expected model access check to fail for model={model}, team_models={team_models}"
+            )
+    except Exception as e:
+        if expect_to_work:
+            pytest.fail(
+                f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {str(e)}"
+            )
 
 
 @pytest.mark.parametrize(