diff --git a/litellm/proxy/common_utils/ldap_sync.py b/litellm/proxy/common_utils/ldap_sync.py index 8c86eb7c2e2..126063e0906 100644 --- a/litellm/proxy/common_utils/ldap_sync.py +++ b/litellm/proxy/common_utils/ldap_sync.py @@ -209,7 +209,11 @@ class LDAPSync: is not True ): return - await self._sync() + try: + await self._sync() + finally: + # Release on completion so a restart (e.g. dry_run -> apply) syncs immediately. + await self.pod_lock_manager.release_lock(cronjob_id=LDAP_SYNC_JOB_ID) def ldap_sync_settings(general_settings: Mapping[str, object]) -> LDAPSyncSettings | None: diff --git a/tests/unit/proxy/common_utils/test_ldap_directory.py b/tests/unit/proxy/common_utils/test_ldap_directory.py index d3ac73af8f8..aee31596142 100644 --- a/tests/unit/proxy/common_utils/test_ldap_directory.py +++ b/tests/unit/proxy/common_utils/test_ldap_directory.py @@ -58,6 +58,7 @@ class Connection: self.extend = self self.standard = self self.searches: list[tuple[str, str, str]] = [] + self.requested_attributes: list[tuple[str, ...]] = [] # mutable-ok: records attributes each search asked for self.closed = False self.tls_ok = tls_ok self.stop_after_search = stop_after_search @@ -86,6 +87,7 @@ class Connection: time_limit: int, generator: bool, ) -> object: + self.requested_attributes.append(tuple(attributes)) return self.groups def search( @@ -98,6 +100,7 @@ class Connection: time_limit: int, ) -> bool: self.searches.append((search_base, search_filter, search_scope)) + self.requested_attributes.append(tuple(attributes)) if search_base in self.raising: result_code, exception = self.raising[search_base] if result_code is not None: @@ -213,6 +216,31 @@ def test_standard_user_object_classes_resolve_as_members(object_class: str): assert read_directory(SETTINGS, connection) == (DirectoryGroup("project-a", "project-a", ("a@example.org",)),) +def test_configured_group_id_and_name_attributes_identify_the_team(): + settings = SETTINGS.model_copy(update={"group_id_attribute": "entryUUID", "group_name_attribute": "description"}) + connection = Connection( + [group(member=[ALICE], entryUUID=["0b6f-team"], description=["Project A"])], + {ALICE: [user("alice@example.org")]}, + ) + assert read_directory(settings, connection) == (DirectoryGroup("0b6f-team", "Project A", ("alice@example.org",)),) + assert connection.requested_attributes[0] == ("entryUUID", "description", "member") + + +def test_configured_user_match_attribute_supplies_member_identity(): + settings = SETTINGS.model_copy(update={"user_match_attribute": "uid"}) + connection = Connection([group(member=[ALICE])], {ALICE: [user("alice@example.org", uid=["alice"])]}) + assert read_directory(settings, connection) == (DirectoryGroup("project-a", "project-a", ("alice",)),) + assert connection.requested_attributes[1] == ("uid", "objectClass") + + +@pytest.mark.parametrize("member_attribute", ["uniqueMember", "roleOccupant"]) +def test_dn_valued_member_attributes_resolve_member_entries(member_attribute: str): + settings = SETTINGS.model_copy(update={"member_attribute": member_attribute}) + connection = Connection([group(**{member_attribute: [ALICE]})], {ALICE: [user("alice@example.org")]}) + assert read_directory(settings, connection) == (DirectoryGroup("project-a", "project-a", ("alice@example.org",)),) + assert connection.searches == [(ALICE, "(objectClass=*)", "BASE")] + + def test_duplicate_ids_are_rejected_even_when_one_group_cannot_resolve_members(): result = read_directory(SETTINGS, Connection([group(member=[]), group(member=[ALICE])])) assert result == DirectoryError("LDAP group IDs are not unique within the selected groups") diff --git a/tests/unit/proxy/common_utils/test_ldap_sync.py b/tests/unit/proxy/common_utils/test_ldap_sync.py index d496d2bbae8..8e1f9fdf15f 100644 --- a/tests/unit/proxy/common_utils/test_ldap_sync.py +++ b/tests/unit/proxy/common_utils/test_ldap_sync.py @@ -107,6 +107,9 @@ class LeaseRedis: async def async_get_cache(self, key: str) -> object | None: return self.values.get(key) + async def async_delete_cache(self, key: str) -> int: + return 1 if self.values.pop(key, None) is not None else 0 + def job_for( prisma: _FakePrisma, @@ -165,17 +168,55 @@ async def test_dry_run_leaves_users_teams_memberships_and_keys_untouched(team_ex @pytest.mark.asyncio -async def test_second_worker_cannot_sync_while_first_workers_lease_is_held(): +async def test_second_worker_cannot_sync_while_first_workers_run_holds_the_lock(monkeypatch: pytest.MonkeyPatch): + prisma: Final = _FakePrisma(teams=[_team()]) + redis: Final = LeaseRedis() + first: Final = job_for(prisma, Connection([group(TEAM, member=[])]), redis=redis) + second_connection: Final = Connection([group(TEAM, member=[])]) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + original: Final = first.sync_group + + async def held_sync_group(group: DirectoryGroup) -> None: + entered.set() + await release.wait() + await original(group) + + monkeypatch.setattr(first, "sync_group", held_sync_group) + running: Final = asyncio.create_task(first.run()) + await entered.wait() + await job_for(prisma, second_connection, redis=redis).run() + release.set() + await running + assert not second_connection.closed + assert redis.values == {} + + +@pytest.mark.asyncio +async def test_finished_run_releases_the_lock_so_a_restarted_worker_syncs_immediately(): prisma: Final = _FakePrisma(teams=[_team()]) redis: Final = LeaseRedis() first: Final = Connection([group(TEAM, member=[])]) - second: Final = Connection([group(TEAM, member=[])]) + restarted: Final = Connection([group(TEAM, member=[])]) await job_for(prisma, first, redis=redis).run() - transactions: Final = prisma.transactions - await job_for(prisma, second, redis=redis).run() + assert redis.values == {} + await job_for(prisma, restarted, redis=redis).run() assert first.closed - assert not second.closed - assert prisma.transactions == transactions + assert restarted.closed + + +@pytest.mark.asyncio +async def test_failed_run_still_releases_the_lock(monkeypatch: pytest.MonkeyPatch): + redis: Final = LeaseRedis() + job: Final = job_for(_FakePrisma(teams=[_team()]), Connection([group(TEAM, member=[])]), redis=redis) + + async def failing_sync() -> None: + raise RuntimeError("database unavailable") + + monkeypatch.setattr(job, "_sync", failing_sync) + with pytest.raises(RuntimeError): + await job.run() + assert redis.values == {} @pytest.mark.asyncio @@ -378,3 +419,34 @@ def test_initialize_requires_shared_redis(): UserApiKeyCache(), None, # pyright: ignore[reportArgumentType] # proxy logging is unused at registration ) + + +@pytest.mark.asyncio +async def test_configured_ldap_and_litellm_identity_fields_create_and_fill_the_team(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from tests.unit.proxy.management_helpers.test_team_roster_sync import ( + _UserRow, # pyright: ignore[reportPrivateUsage] # fake user row with extra columns + ) + + alice: Final = _UserRow.model_validate( + {"user_id": "alice", "user_email": "alice@example.com", "sso_user_id": "z-alice"} + ) + prisma: Final = _FakePrisma(users=[alice, _user("bob")]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + settings: Final = SETTINGS.model_copy( + update={ + "group_id_attribute": "entryUUID", + "group_name_attribute": "description", + "user_match_attribute": "uid", + "litellm_match_field": "sso_user_id", + } + ) + connection: Final = Connection( + [group(member=[ALICE], entryUUID=[TEAM], description=["Project A"])], + {ALICE: [user("alice@example.org", uid=["z-alice"])]}, + ) + await job_for(prisma, connection, settings).run() + created: Final = prisma.writer_db.litellm_teamtable.rows[TEAM] + assert created.team_alias == "Project A" + assert [member.user_id for member in created.members_with_roles] == ["alice"] + assert prisma.writer_db.litellm_usertable.rows["bob"].teams == []