Add PlanFilter to Organization admin

squarelet/organizations/admin.py

@@ -1,6 +1,6 @@
 # Django
 from django.contrib import admin, messages
-from django.db.models import Count, JSONField, Q, Sum
+from django.db.models import Count, Exists, JSONField, OuterRef, Q, Sum
 from django.forms.models import BaseInlineFormSet
 from django.forms.widgets import Textarea
 from django.http.response import HttpResponse
@@ -186,6 +186,26 @@ class MembershipsInline(admin.TabularInline):
     autocomplete_fields = ("from_organization",)
 
 
+class PlanFilter(admin.SimpleListFilter):
+    """Filter organizations by the plan they're subscribed to"""
+
+    title = "plan"
+    parameter_name = "plan"
+    template = "admin/dropdown_filter.html"
+
+    def lookups(self, request, model_admin):
+        plans = Plan.objects.order_by("name").values_list("pk", "name")
+        return [("none", "— No plan —"), *plans]
+
+    def queryset(self, request, queryset):
+        value = self.value()
+        if value is None:
+            return queryset
+        if value == "none":
+            return queryset.filter(subscriptions__isnull=True)
+        return queryset.filter(subscriptions__plan_id=value).distinct()
+
+
 class OverdueInvoiceFilter(admin.SimpleListFilter):
     """Filter organizations by whether they have overdue invoices"""
 
@@ -305,7 +325,16 @@ def export_organizations_as_csv(self, request, queryset):
         "get_subtypes",
         "get_outstanding_invoices",
     )
+
+    def get_list_display(self, request):
+        """Add cancelled column when filtering by plan"""
+        list_display = list(super().get_list_display(request))
+        if "plan" in request.GET and request.GET["plan"] != "none":
+            list_display.append("get_subscription_renews")
+        return list_display
+
     list_filter = (
+        PlanFilter,
         "individual",
         "private",
         "verified_journalist",
@@ -373,7 +402,7 @@ def export_organizations_as_csv(self, request, queryset):
     )
 
     def get_queryset(self, request):
-        return (
+        qs = (
             super()
             .get_queryset(request)
             .prefetch_related("subtypes", "domains")
@@ -386,6 +415,18 @@ def get_queryset(self, request):
                 ),
             )
         )
+        plan_value = request.GET.get("plan")
+        if plan_value and plan_value != "none":
+            qs = qs.annotate(
+                subscription_cancelled=Exists(
+                    Subscription.objects.filter(
+                        organization_id=OuterRef("pk"),
+                        plan_id=plan_value,
+                        cancelled=True,
+                    )
+                )
+            )
+        return qs
 
     def save_model(self, request, obj, form, change):
         super().save_model(request, obj, form, change)
@@ -437,6 +478,16 @@ def get_outstanding_invoices(self, obj):
     get_outstanding_invoices.short_description = "Outstanding Invoices"
     get_outstanding_invoices.admin_order_field = "outstanding_invoice_count"
 
+    def get_subscription_renews(self, obj):
+        cancelled = getattr(obj, "subscription_cancelled", None)
+        if cancelled is None:
+            return "-"
+        return not cancelled
+
+    get_subscription_renews.short_description = "Will Renew"
+    get_subscription_renews.boolean = True
+    get_subscription_renews.admin_order_field = "subscription_cancelled"
+
 
 @admin.register(Plan)
 class PlanAdmin(VersionAdmin):

squarelet/organizations/tests/test_admin.py

@@ -7,7 +7,7 @@
 import stripe
 
 # Squarelet
-from squarelet.organizations.admin import InvoiceAdmin, OrganizationAdmin
+from squarelet.organizations.admin import InvoiceAdmin, OrganizationAdmin, PlanFilter
 from squarelet.organizations.models import Invoice, Organization
 
 
@@ -340,3 +340,136 @@ def test_save_model_subscribes_when_verified_journalist_set(
         org_admin.save_model(request, org, mock_form, change=True)
 
         subscribe_mock.assert_called_once()
+
+
+class TestPlanFilter:
+    """Tests for the PlanFilter admin list filter."""
+
+    @pytest.fixture
+    def request_factory(self):
+        return RequestFactory()
+
+    def _filter(self, params):
+        return PlanFilter(
+            request=None,
+            params=params,
+            model=Organization,
+            model_admin=OrganizationAdmin(Organization, AdminSite()),
+        )
+
+    @pytest.mark.django_db
+    def test_lookups_include_no_plan_and_each_plan(self, request_factory, plan_factory):
+        plan_a = plan_factory(name="Plan A")
+        plan_b = plan_factory(name="Plan B")
+        request = request_factory.get("/")
+        lookups = self._filter({}).lookups(
+            request, OrganizationAdmin(Organization, AdminSite())
+        )
+        values = [value for value, _label in lookups]
+        labels = [label for _value, label in lookups]
+        assert "none" in values
+        assert plan_a.pk in values
+        assert plan_b.pk in values
+        assert "— No plan —" in labels
+
+    @pytest.mark.django_db
+    def test_filter_by_plan_returns_orgs_with_that_subscription(
+        self,
+        request_factory,
+        organization_factory,
+        plan_factory,
+        subscription_factory,
+    ):
+        pro_plan = plan_factory(name="Pro")
+        free_plan = plan_factory(name="Free")
+        pro_org = organization_factory()
+        free_org = organization_factory()
+        subscription_factory(organization=pro_org, plan=pro_plan)
+        subscription_factory(organization=free_org, plan=free_plan)
+
+        request = request_factory.get("/")
+        result = self._filter({"plan": [str(pro_plan.pk)]}).queryset(
+            request, Organization.objects.all()
+        )
+        assert pro_org in result
+        assert free_org not in result
+
+    @pytest.mark.django_db
+    def test_filter_by_plan_includes_cancelled_subscriptions(
+        self,
+        request_factory,
+        organization_factory,
+        plan_factory,
+        subscription_factory,
+    ):
+        """Cancelled subs still count as subscribed per design decision."""
+        pro_plan = plan_factory(name="Pro")
+        org = organization_factory()
+        subscription_factory(organization=org, plan=pro_plan, cancelled=True)
+
+        request = request_factory.get("/")
+        result = self._filter({"plan": [str(pro_plan.pk)]}).queryset(
+            request, Organization.objects.all()
+        )
+        assert org in result
+
+    @pytest.mark.django_db
+    def test_filter_by_none_returns_orgs_without_subscriptions(
+        self,
+        request_factory,
+        organization_factory,
+        plan_factory,
+        subscription_factory,
+    ):
+        plan = plan_factory(name="Pro")
+        subscribed = organization_factory()
+        unsubscribed = organization_factory()
+        subscription_factory(organization=subscribed, plan=plan)
+
+        request = request_factory.get("/")
+        result = self._filter({"plan": ["none"]}).queryset(
+            request, Organization.objects.all()
+        )
+        assert unsubscribed in result
+        assert subscribed not in result
+
+    @pytest.mark.django_db
+    def test_unset_filter_is_a_no_op(
+        self,
+        request_factory,
+        organization_factory,
+        plan_factory,
+        subscription_factory,
+    ):
+        plan = plan_factory(name="Pro")
+        subscribed = organization_factory()
+        unsubscribed = organization_factory()
+        subscription_factory(organization=subscribed, plan=plan)
+
+        request = request_factory.get("/")
+        result = self._filter({}).queryset(request, Organization.objects.all())
+        assert subscribed in result
+        assert unsubscribed in result
+
+    @pytest.mark.django_db
+    def test_filter_result_has_no_duplicate_rows(
+        self,
+        request_factory,
+        organization_factory,
+        plan_factory,
+        subscription_factory,
+    ):
+        plan = plan_factory(name="Pro")
+        other_plan = plan_factory(name="Other")
+        org = organization_factory()
+        subscription_factory(organization=org, plan=plan)
+        subscription_factory(organization=org, plan=other_plan)
+
+        request = request_factory.get("/")
+        result = self._filter({"plan": [str(plan.pk)]}).queryset(
+            request, Organization.objects.all()
+        )
+        assert list(result).count(org) == 1
+
+    def test_filter_is_registered_on_organization_admin(self):
+        assert PlanFilter in OrganizationAdmin.list_filter