From 8b7ffad59846c962301106fdbd535d60261e181b Mon Sep 17 00:00:00 2001
From: Vadim Stepanov <vadimkerr@gmail.com>
Date: Thu, 11 Jan 2024 15:03:54 +0000
Subject: [PATCH] Add team filter for users endpoint (#3666)

# What this PR does

Adds `team` filter for `users` endpoint

## Checklist

- [x] Unit, integration, and e2e (if applicable) tests updated
- [x] Documentation added (or `pr:no public docs` PR label added if not
required)
- [x] `CHANGELOG.md` updated (or `pr:no changelog` PR label added if not
required)
---
 engine/apps/api/tests/test_user.py   | 42 ++++++++++++++++++++++++++++
 engine/apps/api/views/user.py        |  4 ++-
 engine/common/api_helpers/filters.py |  3 +-
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/engine/apps/api/tests/test_user.py b/engine/apps/api/tests/test_user.py
index 68e38a98..5d7dc538 100644
--- a/engine/apps/api/tests/test_user.py
+++ b/engine/apps/api/tests/test_user.py
@@ -287,6 +287,48 @@ def test_list_users_filtered_by_public_primary_key(
     assert returned_user_pks == [user1.public_primary_key]
 
 
+@pytest.mark.django_db
+def test_list_users_filtered_by_team(
+    make_organization,
+    make_team,
+    make_user_for_organization,
+    make_token_for_organization,
+    make_user_auth_headers,
+):
+    organization = make_organization()
+    user1 = make_user_for_organization(organization)
+    user2 = make_user_for_organization(organization)
+
+    team1 = make_team(organization)
+    team2 = make_team(organization)
+    team3 = make_team(organization)
+
+    user1.teams.add(team1)
+    user1.teams.add(team2)
+    user2.teams.add(team2)
+
+    _, token = make_token_for_organization(organization)
+    client = APIClient()
+    url = reverse("api-internal:user-list")
+
+    def _get_user_pks(teams):
+        response = client.get(
+            url,
+            data={"team": [team.public_primary_key for team in teams]},  # these are query params
+            **make_user_auth_headers(user1, token),
+        )
+        assert response.status_code == status.HTTP_200_OK
+        return [u["pk"] for u in response.json()["results"]]
+
+    assert _get_user_pks([team1]) == [user1.public_primary_key]
+    assert _get_user_pks([team1, team2]) == [user1.public_primary_key, user2.public_primary_key]
+    assert _get_user_pks([team3]) == []
+
+    # check non-existent team returns bad request
+    response = client.get(f"{url}?team=null", **make_user_auth_headers(user1, token))
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+
 @pytest.mark.django_db
 def test_notification_chain_verbal(
     make_organization,
diff --git a/engine/apps/api/views/user.py b/engine/apps/api/views/user.py
index 0716e0be..effeb313 100644
--- a/engine/apps/api/views/user.py
+++ b/engine/apps/api/views/user.py
@@ -64,6 +64,7 @@ from apps.telegram.client import TelegramClient
 from apps.telegram.models import TelegramVerificationCode
 from apps.user_management.models import Team, User
 from common.api_helpers.exceptions import Conflict
+from common.api_helpers.filters import ByTeamModelFieldFilterMixin, TeamModelMultipleChoiceFilter
 from common.api_helpers.mixins import FilterSerializerMixin, PublicPrimaryKeyMixin
 from common.api_helpers.paginators import HundredPageSizePaginator
 from common.api_helpers.utils import create_engine_url
@@ -113,7 +114,7 @@ class CurrentUserView(APIView):
         return Response(serializer.data)
 
 
-class UserFilter(filters.FilterSet):
+class UserFilter(ByTeamModelFieldFilterMixin, filters.FilterSet):
     """
     https://django-filter.readthedocs.io/en/master/guide/rest_framework.html
     """
@@ -122,6 +123,7 @@ class UserFilter(filters.FilterSet):
     # TODO: remove "roles" in next version
     roles = filters.MultipleChoiceFilter(field_name="role", choices=LegacyAccessControlRole.choices())
     permission = filters.ChoiceFilter(method="filter_by_permission", choices=ALL_PERMISSION_CHOICES)
+    team = TeamModelMultipleChoiceFilter(field_name="teams", null_label=None, null_value=None)
 
     class Meta:
         model = User
diff --git a/engine/common/api_helpers/filters.py b/engine/common/api_helpers/filters.py
index 7e12a140..7e174ae8 100644
--- a/engine/common/api_helpers/filters.py
+++ b/engine/common/api_helpers/filters.py
@@ -85,8 +85,7 @@ class ByTeamModelFieldFilterMixin:
         if null_team_lookup is not None:
             teams_lookup = teams_lookup | null_team_lookup
 
-        queryset = queryset.filter(teams_lookup)
-        return queryset
+        return queryset.filter(teams_lookup).distinct()
 
 
 def get_team_queryset(request):