Fix remove_user_contracts to maintain user filtering and only remove managed contracts (#3031)

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: blarghmatey <[email protected]>
Co-authored-by: jkachel <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

b2b/models.py

@@ -162,7 +162,7 @@ def add_user_contracts(self, user):
 
         return contracts_qs.count()
 
-    def remove_user_contracts(self, user):  # noqa: ARG002
+    def remove_user_contracts(self, user):
         """
         Remove managed contracts from the given user.
 
@@ -171,7 +171,13 @@ def remove_user_contracts(self, user):  # noqa: ARG002
         Returns:
         - int: number of contracts removed
         """
-        return 0
+
+        return user.b2b_contracts.through.objects.filter(
+            user_id=user.id,
+            contractpage_id__in=self.contracts.filter(
+                integration_type__in=CONTRACT_MEMBERSHIP_AUTOS
+            ).values_list("id", flat=True),
+        ).delete()
 
     def __str__(self):
         """Return a reasonable representation of the org as a string."""

b2b/models_test.py

@@ -4,7 +4,11 @@
 import pytest
 
 from b2b.factories import ContractPageFactory, OrganizationPageFactory
-from courses.factories import CourseRunFactory, ProgramFactory
+from courses.factories import (
+    CourseRunFactory,
+    ProgramFactory,
+)
+from users.factories import UserFactory
 
 pytestmark = [pytest.mark.django_db]
 FAKE = faker.Faker()
@@ -94,3 +98,93 @@ def test_organization_page_slug_not_overwritten_if_set():
     # The slug should still be the custom one
     assert org.slug == "custom-slug"
     assert org.title == "Test Org Updated"
+
+
+def test_remove_user_contracts_only_affects_specified_user():
+    """Test that remove_user_contracts only removes contracts for the specified user."""
+
+    # Create an organization and contracts
+    organization = OrganizationPageFactory.create()
+    contract1 = ContractPageFactory.create(
+        organization=organization,
+        membership_type="auto",
+        integration_type="auto",
+    )
+    contract2 = ContractPageFactory.create(
+        organization=organization,
+        membership_type="managed",
+        integration_type="managed",
+    )
+
+    # Create two users and add them both to the contracts
+    user1 = UserFactory.create()
+    user2 = UserFactory.create()
+
+    user1.b2b_contracts.add(contract1, contract2)
+    user2.b2b_contracts.add(contract1, contract2)
+
+    # Verify both users have the contracts
+    assert user1.b2b_contracts.count() == 2
+    assert user2.b2b_contracts.count() == 2
+
+    # Remove contracts from user1
+    organization.remove_user_contracts(user1)
+
+    # Verify user1's contracts are removed
+    user1.refresh_from_db()
+    assert user1.b2b_contracts.count() == 0
+
+    # Verify user2's contracts are NOT affected
+    user2.refresh_from_db()
+    assert user2.b2b_contracts.count() == 2
+    assert user2.b2b_contracts.filter(id=contract1.id).exists()
+    assert user2.b2b_contracts.filter(id=contract2.id).exists()
+
+
+def test_remove_user_contracts_only_removes_managed_contracts():
+    """Test that remove_user_contracts only removes automatically managed contracts."""
+
+    # Create an organization with both managed and non-managed contracts
+    organization = OrganizationPageFactory.create()
+
+    # Automatically managed contracts (should be removed)
+    auto_contract = ContractPageFactory.create(
+        organization=organization,
+        membership_type="auto",
+        integration_type="auto",
+    )
+    managed_contract = ContractPageFactory.create(
+        organization=organization,
+        membership_type="managed",
+        integration_type="managed",
+    )
+    sso_contract = ContractPageFactory.create(
+        organization=organization,
+        membership_type="sso",
+        integration_type="sso",
+    )
+
+    # Non-managed contract (should NOT be removed)
+    code_contract = ContractPageFactory.create(
+        organization=organization,
+        membership_type="code",
+        integration_type="code",
+    )
+
+    # Create a user and add all contracts
+    user = UserFactory.create()
+    user.b2b_contracts.add(auto_contract, managed_contract, sso_contract, code_contract)
+
+    # Verify user has all 4 contracts
+    assert user.b2b_contracts.count() == 4
+
+    # Remove managed contracts from user
+    organization.remove_user_contracts(user)
+
+    # Verify only managed contracts are removed, code contract remains
+    user.refresh_from_db()
+    assert user.b2b_contracts.count() == 1
+    assert not user.b2b_contracts.filter(id=auto_contract.id).exists()
+    assert not user.b2b_contracts.filter(id=managed_contract.id).exists()
+    assert not user.b2b_contracts.filter(id=sso_contract.id).exists()
+    assert user.b2b_contracts.filter(id=code_contract.id).exists()