improve tests

osf/management/commands/migrate_preprint_affiliation.py

@@ -38,7 +38,7 @@ def handle(self, *args, **options):
         if dry_run:
             logger.info('Dry run mode activated.')
 
-        processed_count, updated_count, skipped_count = assign_affiliations_to_contributors(
+        processed_count, updated_count, skipped_count = assign_affiliations_to_preprint(
             exclude_guids=exclude_guids,
             dry_run=dry_run
         )
@@ -49,7 +49,7 @@ def handle(self, *args, **options):
         logger.info(f'Total run time: {finish_time - start_time}')
 
 
-def assign_affiliations_to_contributors(exclude_guids=None, dry_run=True):
+def assign_affiliations_to_preprint(exclude_guids=None, dry_run=True):
     exclude_guids = exclude_guids or set()
 
     # Define the permissions to filter contributors

osf_tests/management_commands/test_migrate_preprint_affiliations.py

@@ -1,10 +1,15 @@
 import pytest
-from osf.management.commands.migrate_preprint_affiliation import assign_creator_affiliations_to_preprints
-from osf_tests.factories import PreprintFactory, InstitutionFactory, AuthUserFactory
+from django.core.management import call_command
+from osf.management.commands.migrate_preprint_affiliation import assign_affiliations_to_preprint
+from osf_tests.factories import (
+    PreprintFactory,
+    InstitutionFactory,
+    AuthUserFactory,
+)
 
 
 @pytest.mark.django_db
-class TestAssignCreatorAffiliationsToPreprints:
+class TestAssignAffiliationsToContributors:
 
     @pytest.fixture()
     def institution(self):
@@ -22,32 +27,131 @@ def user_without_affiliation(self):
         return AuthUserFactory()
 
     @pytest.fixture()
-    def preprint_with_affiliated_creator(self, user_with_affiliation):
-        return PreprintFactory(creator=user_with_affiliation)
+    def preprint_with_affiliated_contributor(self, user_with_affiliation, institution):
+        preprint = PreprintFactory()
+        preprint.affiliated_institutions.add(institution)
+        preprint.save()
+        preprint.add_contributor(
+            user=user_with_affiliation,
+            permission='admin',
+            visible=True
+        )
+        return preprint
 
     @pytest.fixture()
-    def preprint_with_non_affiliated_creator(self, user_without_affiliation):
-        return PreprintFactory(creator=user_without_affiliation)
+    def preprint_with_non_affiliated_contributor(self, user_without_affiliation, institution):
+        preprint = PreprintFactory()
+        preprint.affiliated_institutions.add(institution)
+        preprint.save()
+        preprint.add_contributor(
+            user=user_without_affiliation,
+            permission='admin',
+            visible=True
+        )
+        return preprint
+
+    @pytest.fixture()
+    def opt_out_user(self, institution):
+        user = AuthUserFactory()
+        user.affiliated_institutions.add(institution)
+        user.save()
+        user.save()
+        return user
+
+    @pytest.fixture()
+    def preprint_with_opted_out_contributor(self, opt_out_user, institution):
+        preprint = PreprintFactory()
+        preprint.affiliated_institutions.add(institution)
+        preprint.save()
+        preprint.add_contributor(
+            user=opt_out_user,
+            permission='admin',
+            visible=True
+        )
+        return preprint
 
     @pytest.mark.parametrize('dry_run', [True, False])
-    def test_assign_affiliations_with_affiliated_creator(self, preprint_with_affiliated_creator, institution, dry_run):
-        assert preprint_with_affiliated_creator.affiliated_institutions.count() == 0
+    def test_assign_affiliations_with_affiliated_contributor(self, preprint_with_affiliated_contributor, dry_run):
+        user = preprint_with_affiliated_contributor.contributors.first()
+        user.affiliated_institutions.clear()
+        user.save()
 
-        assign_creator_affiliations_to_preprints(dry_run=dry_run)
+        assign_affiliations_to_preprint(dry_run=dry_run)
 
         if dry_run:
-            assert preprint_with_affiliated_creator.affiliated_institutions.count() == 0
+            assert not user.affiliated_institutions.exists()
         else:
-            assert institution in preprint_with_affiliated_creator.affiliated_institutions.all()
+            assert user.affiliated_institutions.exists()
+
+    @pytest.mark.parametrize('dry_run', [True, False])
+    def test_no_affiliations_for_non_affiliated_contributor(self, preprint_with_non_affiliated_contributor, dry_run):
+        user = preprint_with_non_affiliated_contributor.contributors.first()
+        assign_affiliations_to_preprint(dry_run=dry_run)
+        assert not user.affiliated_institutions.exists()
+
+    @pytest.mark.parametrize('dry_run', [True, False])
+    def test_exclude_contributor_by_guid(self, preprint_with_affiliated_contributor, dry_run):
+        user = preprint_with_affiliated_contributor.contributors.first()
+        user.affiliated_institutions.clear()
+        user.save()
+
+        exclude_guids = {user._id}
+        assign_affiliations_to_preprint(exclude_guids=exclude_guids, dry_run=dry_run)
+
+        assert not user.affiliated_institutions.exists()
 
     @pytest.mark.parametrize('dry_run', [True, False])
-    def test_no_affiliations_for_non_affiliated_creator(self, preprint_with_non_affiliated_creator, dry_run):
-        assign_creator_affiliations_to_preprints(dry_run=dry_run)
-        assert preprint_with_non_affiliated_creator.affiliated_institutions.count() == 0
+    def test_no_affiliations_for_opted_out_contributor(self, preprint_with_opted_out_contributor, dry_run):
+        user = preprint_with_opted_out_contributor.contributors.first()
+        user.affiliated_institutions.clear()
+        user.save()
+
+        assign_affiliations_to_preprint(dry_run=dry_run)
+
+        assert not user.affiliated_institutions.exists()
 
     @pytest.mark.parametrize('dry_run', [True, False])
-    def test_exclude_creator_by_guid(self, preprint_with_affiliated_creator, institution, dry_run):
-        exclude_guid = preprint_with_affiliated_creator.creator._id
-        assign_creator_affiliations_to_preprints(exclude_guids={exclude_guid}, dry_run=dry_run)
+    def test_affiliations_from_other_contributor(self, institution, dry_run):
+        # Create a preprint with two contributors: one opted out, one not
+        opt_out_user = AuthUserFactory()
+        opt_out_user.affiliated_institutions.add(institution)
+        opt_out_user.opted_out_affiliation = True
+        opt_out_user.save()
+
+        non_opt_out_user = AuthUserFactory()
+        non_opt_out_user.affiliated_institutions.add(institution)
+        non_opt_out_user.save()
+
+        preprint = PreprintFactory()
+        preprint.affiliated_institutions.clear()
+        preprint.save()
+        preprint.add_contributor(
+            user=opt_out_user,
+            permission='admin',
+            visible=True
+        )
+        preprint.add_contributor(
+            user=non_opt_out_user,
+            permission='write',
+            visible=True
+        )
+
+        assign_affiliations_to_preprint(dry_run=dry_run)
+
+        preprint.refresh_from_db()
+
+        if dry_run:
+            assert not preprint.affiliated_institutions.exists()
+        else:
+            # The preprint should get the institution from the non-opted-out contributor
+            assert institution in preprint.affiliated_institutions.all()
+
+    def test_command_execution(self, preprint_with_affiliated_contributor):
+        user = preprint_with_affiliated_contributor.contributors.first()
+        user.affiliated_institutions.clear()
+        user.save()
+
+        call_command('assign_affiliations_to_contributors')
 
-        assert preprint_with_affiliated_creator.affiliated_institutions.count() == 0
+        user.refresh_from_db()
+        assert user.affiliated_institutions.exists()