[BUILD-407] Create a script to archive duplicated memberships (#2137)

[BUILD-407] Create a script to archive duplicated memberships (#2137)

  • create a script to archive duplicated memberships

  • refactor script

diff --git a/scripts/archive_duplicated_memberships.py b/scripts/archive_duplicated_memberships.py
new file mode 100644
index 0000000..05a778b
--- /dev/null
+++ b/scripts/archive_duplicated_memberships.py
@@ -0,0 +1,53 @@
+import stripe
+from django.conf import settings
+from django.db import transaction
+from django.db.models import Count, Q
+
+from ankihub.memberships.models import Membership, MembershipStatusChoices
+from ankihub.users.models import User
+
+
+def run(*args):
+    stripe.api_key = settings.STRIPE_API_KEY
+    users_with_multiple_memberships = User.objects.annotate(
+        num_memberships=Count(
+            "memberships",
+            filter=~Q(memberships__status=MembershipStatusChoices.ARCHIVED),
+        )
+    ).filter(num_memberships__gt=1)
+    print(
+        f"Found {users_with_multiple_memberships.count()} users with multiple memberships"
+    )
+    for user in users_with_multiple_memberships:
+        print(f"Archiving memberships for id={user}, username={user.username}")
+
+        with transaction.atomic():
+            memberships = user.memberships.filter(
+                ~Q(status=MembershipStatusChoices.ARCHIVED)
+            ).order_by("status", "-created")
+            print(f"Found {memberships.count()} memberships for {user}")
+
+            customer_ids = {mem.get_customer_id() for mem in memberships}
+            definitive_customer_id = memberships.first().get_customer_id()
+            customer_ids.discard(definitive_customer_id)
+            print(f"Found {len(customer_ids)} customer ids for {user.username}")
+
+            user.customer_id = definitive_customer_id
+            user.save(update_fields=["customer_id"])
+
+            to_archive = memberships.values_list("id", flat=True)[1:]
+            Membership.objects.filter(id__in=list(to_archive)).update(
+                status=MembershipStatusChoices.ARCHIVED
+            )
+            print(f"Archived {len(to_archive)} memberships for {user.username}")
+
+            assert (
+                user.memberships.filter(
+                    ~Q(status=MembershipStatusChoices.ARCHIVED)
+                ).count()
+                == 1
+            ), f"Fail to archive {user} memberships, still has multiple memberships"
+
+            for customer_id in set(customer_ids):
+                print(f"Deleting customer {customer_id} related to {user.username}")
+                stripe.Customer.delete(customer_id)
diff --git a/scripts/tests/test_scripts.py b/scripts/tests/test_scripts.py
index 02e0e20..ffd68f3 100644
--- a/scripts/tests/test_scripts.py
+++ b/scripts/tests/test_scripts.py
@@ -16,10 +16,11 @@ from pytest import MonkeyPatch
 from ankihub.conftest import assert_fields_match, fields_from_string_list
 from ankihub.decks.models import ChangeNoteSuggestion, DeckMedia, TagNode
 from ankihub.decks.tests.factories import DeckFactory, NoteFactory, NoteTagFactory
-from ankihub.memberships.models import Membership
+from ankihub.memberships.models import Membership, MembershipStatusChoices
 from ankihub.memberships.tests.factories import MembershipFactory, OrganizationFactory
 from ankihub.users.tests.factories import UserFactory
 from scripts import (
+    archive_duplicated_memberships,
     create_sponsored_users_and_memberships,
     save_deck_tags_tree,
     update_ome_fields,
@@ -594,3 +595,173 @@ def test_save_deck_tags_tree_from_file():
         ]
     )
     assert TagNode.objects.count() == 3
+
+
+@pytest.mark.django_db
+class TestArchiveDuplicatedMemberships:
+    def test_archive_duplicated_memberships(self, monkeypatch):
+        user = UserFactory()
+        customer_id_1 = "cus_1"
+        customer_id_2 = "cus_2"
+        customer_id_3 = "cus_3"
+        active_membership = MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.ACTIVE,
+            metadata={"customer": {"id": customer_id_1}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.INACTIVE,
+            metadata={"customer": {"id": customer_id_2}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.ARCHIVED,
+            metadata={"customer": {"id": customer_id_3}},
+        )
+
+        mock_delete_customer = Mock()
+        monkeypatch.setattr(
+            "scripts.archive_duplicated_memberships.stripe.Customer.delete",
+            mock_delete_customer,
+        )
+
+        assert not user.customer_id
+
+        archive_duplicated_memberships.run()
+        user.refresh_from_db()
+
+        assert (
+            Membership.objects.filter(
+                user=user, status=MembershipStatusChoices.ARCHIVED
+            ).count()
+            == 2
+        )
+        assert user.memberships.get_ankihub_membership() == active_membership
+        assert user.customer_id == customer_id_1
+        mock_delete_customer.assert_called_once_with(customer_id_2)
+
+    def test_archive_duplicated_memberships_does_not_delete_the_same_customer_id(
+        self, monkeypatch
+    ):
+        user = UserFactory()
+        customer_id_1 = "cus_1"
+        active_membership = MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.ACTIVE,
+            metadata={"customer": {"id": customer_id_1}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.INACTIVE,
+            metadata={"customer": {"id": customer_id_1}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.ARCHIVED,
+            metadata={"customer": {"id": customer_id_1}},
+        )
+
+        mock_delete_customer = Mock()
+        monkeypatch.setattr(
+            "scripts.archive_duplicated_memberships.stripe.Customer.delete",
+            mock_delete_customer,
+        )
+
+        assert not user.customer_id
+
+        archive_duplicated_memberships.run()
+        user.refresh_from_db()
+
+        assert (
+            Membership.objects.filter(
+                user=user, status=MembershipStatusChoices.ARCHIVED
+            ).count()
+            == 2
+        )
+        assert user.memberships.get_ankihub_membership() == active_membership
+        assert user.customer_id == customer_id_1
+        mock_delete_customer.assert_not_called()
+
+    def test_archive_duplicated_memberships_delete_all_spare_customer_ids(
+        self, monkeypatch
+    ):
+        user = UserFactory()
+        customer_id_1 = "cus_1"
+        customer_id_2 = "cus_2"
+        customer_id_3 = "cus_3"
+        active_membership = MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.ACTIVE,
+            metadata={"customer": {"id": customer_id_1}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.INACTIVE,
+            metadata={"customer": {"id": customer_id_2}},
+        )
+        MembershipFactory(
+            user=user,
+            status=MembershipStatusChoices.FAILED,
+            metadata={"customer": {"id": customer_id_3}},
+        )
+
+        mock_delete_customer = Mock()
+        monkeypatch.setattr(
+            "scripts.archive_duplicated_memberships.stripe.Customer.delete",
+            mock_delete_customer,
+        )
+
+        assert not user.customer_id
+
+        archive_duplicated_memberships.run()
+        user.refresh_from_db()
+
+        assert (
+            Membership.objects.filter(
+                user=user, status=MembershipStatusChoices.ARCHIVED
+            ).count()
+            == 2
+        )
+        assert user.memberships.get_ankihub_membership() == active_membership
+        assert user.customer_id == customer_id_1
+        mock_delete_customer.assert_has_calls(
+            [call(customer_id_2), call(customer_id_3)], any_order=True
+        )
+
+    def test_archive_duplicated_memberships_use_last_created(self, monkeypatch):

[... diff too long, it was truncated ...]

GitHub
sha: 5b5017db32fa340c40e1434e0a9b6309e9ad1c63