From f9e6b78cf67015f0945edac0da73804ebb7e4eb2 Mon Sep 17 00:00:00 2001 From: sandyr Date: Wed, 20 Sep 2023 17:43:38 +0100 Subject: [PATCH] propagate suppression status through models using "suppression tree" --- .../commands/clean_empty_studies.py | 2 +- emgapi/migrations/0011_auto_20230914_0716.py | 44 +++++ emgapi/models.py | 137 +++++++++---- .../commands/sync_assemblies_with_ena.py | 6 +- .../management/commands/sync_runs_with_ena.py | 5 +- .../commands/sync_samples_with_ena.py | 3 - .../commands/sync_studies_with_ena.py | 3 - tests/ena/test_sync_assemblies_with_ena.py | 20 +- tests/ena/test_sync_runs_with_ena.py | 28 ++- tests/ena/test_sync_samples_with_ena.py | 18 +- tests/ena/test_sync_studies_with_ena.py | 27 ++- tests/test_utils/emg_fixtures.py | 181 +++++++++++++++++- 12 files changed, 400 insertions(+), 74 deletions(-) create mode 100644 emgapi/migrations/0011_auto_20230914_0716.py diff --git a/emgapi/management/commands/clean_empty_studies.py b/emgapi/management/commands/clean_empty_studies.py index 02fcc9e79..88ed4179e 100644 --- a/emgapi/management/commands/clean_empty_studies.py +++ b/emgapi/management/commands/clean_empty_studies.py @@ -56,5 +56,5 @@ def handle(self, *args, **kwargs): run.ena_study_accession = study.secondary_accession run.study = None Run.objects.bulk_update(runs, ["ena_study_accession", "study"]) - study.suppress() + study.suppress(propagate=False) logger.info(f"{study} suppressed") diff --git a/emgapi/migrations/0011_auto_20230914_0716.py b/emgapi/migrations/0011_auto_20230914_0716.py new file mode 100644 index 000000000..1bc01b352 --- /dev/null +++ b/emgapi/migrations/0011_auto_20230914_0716.py @@ -0,0 +1,44 @@ +# Generated by Django 3.2.18 on 2023-09-14 07:16 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('emgapi', '0010_runextraannotation'), + ] + + operations = [ + migrations.AlterField( + model_name='analysisjob', + name='suppression_reason', + field=models.IntegerField(blank=True, choices=[(1, 'Draft'), (3, 'Cancelled'), (5, 'Suppressed'), (6, 'Killed'), (7, 'Temporary Suppressed'), (8, 'Temporary Killed'), (100, 'Ancestor Suppressed')], db_column='SUPPRESSION_REASON', null=True), + ), + migrations.AlterField( + model_name='assembly', + name='study', + field=models.ForeignKey(blank=True, db_column='STUDY_ID', null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='assemblies', to='emgapi.study'), + ), + migrations.AlterField( + model_name='assembly', + name='suppression_reason', + field=models.IntegerField(blank=True, choices=[(1, 'Draft'), (3, 'Cancelled'), (5, 'Suppressed'), (6, 'Killed'), (7, 'Temporary Suppressed'), (8, 'Temporary Killed'), (100, 'Ancestor Suppressed')], db_column='SUPPRESSION_REASON', null=True), + ), + migrations.AlterField( + model_name='run', + name='suppression_reason', + field=models.IntegerField(blank=True, choices=[(1, 'Draft'), (3, 'Cancelled'), (5, 'Suppressed'), (6, 'Killed'), (7, 'Temporary Suppressed'), (8, 'Temporary Killed'), (100, 'Ancestor Suppressed')], db_column='SUPPRESSION_REASON', null=True), + ), + migrations.AlterField( + model_name='sample', + name='suppression_reason', + field=models.IntegerField(blank=True, choices=[(1, 'Draft'), (3, 'Cancelled'), (5, 'Suppressed'), (6, 'Killed'), (7, 'Temporary Suppressed'), (8, 'Temporary Killed'), (100, 'Ancestor Suppressed')], db_column='SUPPRESSION_REASON', null=True), + ), + migrations.AlterField( + model_name='study', + name='suppression_reason', + field=models.IntegerField(blank=True, choices=[(1, 'Draft'), (3, 'Cancelled'), (5, 'Suppressed'), (6, 'Killed'), (7, 'Temporary Suppressed'), (8, 'Temporary Killed'), (100, 'Ancestor Suppressed')], db_column='SUPPRESSION_REASON', null=True), + ), + ] diff --git a/emgapi/models.py b/emgapi/models.py index 52a7c46e8..dab01e622 100644 --- a/emgapi/models.py +++ b/emgapi/models.py @@ -16,6 +16,7 @@ import logging +from django.apps import apps from django.conf import settings from django.db import models from django.db.models import (CharField, Count, OuterRef, Prefetch, Q, @@ -73,7 +74,7 @@ class SuppressibleModel(models.Model): suppressible_descendants = [] # List of related_names from this model that should have their suppression status propagated from this. # E.g. Study.suppressible_descendants = ['samples'] to suppress a study's samples if the study is suppressed. - + class Reason(models.IntegerChoices): DRAFT = 1 CANCELLED = 3 @@ -88,6 +89,79 @@ class Reason(models.IntegerChoices): suppressed_at = models.DateTimeField(db_column='SUPPRESSED_AT', blank=True, null=True) suppression_reason = models.IntegerField(db_column='SUPPRESSION_REASON', blank=True, null=True, choices=Reason.choices) + def _get_suppression_descendant_tree(self, suppressing: bool = True): + """ + Recursively find all suppressible descendants of the calling suppressible model. + :param suppressing: True if looking for descendants that should be (i.e. are not currently) suppressed. False if opposite. + :return: Dict mapping model names to sets of model instances + """ + suppressibles = {} + + def __add_to_suppressibles(kls, additional_suppressibles): + suppressibles.setdefault( + kls, + set() + ).update(additional_suppressibles) + + logger.debug(f'Building suppression descendant tree for {self._meta.object_name}') + + for descendant_relation_name in self.suppressible_descendants: + descendant_relation = getattr( + self, + descendant_relation_name + ) + descendants_to_update = descendant_relation.filter( + is_suppressed=not suppressing, + ) + if not descendants_to_update.exists(): + logger.debug(f'No {descendant_relation_name} descendants to handle.') + continue + # Check whether the descendant might have other non-suppressed ancestors of the same type as this + # (If so, it shouldn't be suppressed). + relation_field = self._meta.get_field(descendant_relation_name) + logger.info(f'{relation_field = }') + if isinstance(relation_field, models.ManyToManyField): + logger.debug( + f"Descendant relation {descendant_relation_name} on {self.__class__} is a Many2Many." + f"Checking whether descendants have unsuppressed siblings.." + ) + logger.debug(f"Before filtering, had {descendants_to_update.count()} {descendant_relation_name}") + + descendant_ids_with_unsuppressed_alike_ancestors = descendant_relation.through.objects.filter( + **{ + f"{descendant_relation.target_field_name}__in": descendant_relation.all(), # e.g. sample in study.samples + f"{descendant_relation.source_field_name}__is_suppressed": False, # e.g. not study.is_suppressed + } + ).exclude( + **{ + f"{descendant_relation.source_field_name}": self, # e.g. study != self + } + ).values_list( + f"{descendant_relation.target_field_name}_id", + flat=True + ) + descendants_to_update = descendants_to_update.exclude( + pk__in=descendant_ids_with_unsuppressed_alike_ancestors + ) + + logger.debug(f"After filtering, had {descendants_to_update.count()} {descendant_relation_name}") + + __add_to_suppressibles( + descendant_relation.model._meta.object_name, + descendants_to_update + ) + + for descendant in descendants_to_update: + for kls, kls_suppressibles in descendant._get_suppression_descendant_tree( + suppressing + ).items(): + __add_to_suppressibles( + kls, + kls_suppressibles + ) + + return suppressibles + def suppress(self, suppression_reason=None, save=True, propagate=True): self.is_suppressed = True self.suppressed_at = timezone.now() @@ -95,57 +169,34 @@ def suppress(self, suppression_reason=None, save=True, propagate=True): if save: self.save() if propagate: - for descendant_relation in self.suppressible_descendants: - descendants_to_suppress: QuerySet = getattr( - self, - descendant_relation - ).filter( - is_suppressed=False - ) - for descendant in descendants_to_suppress: + descendant_tree = self._get_suppression_descendant_tree() + for descendants_object_type, descendants in descendant_tree.items(): + for descendant in descendants: descendant.is_suppressed = True descendant.suppression_reason = self.Reason.ANCESTOR_SUPPRESSED - descendants_to_suppress.bulk_update( - descendants_to_suppress, - [ - 'is_suppressed', - 'suppression_reason' - ] - ) + m: SuppressibleModel = apps.get_model(app_label='emgapi', model_name=descendants_object_type) + m.objects.bulk_update(descendants, fields=['is_suppressed', 'suppression_reason']) logger.info( - f'Propagated suppression of {self} ' - f'to {len(descendants_to_suppress)} {descendant_relation} descendants' + f'Propagated suppression of {self} to {len(descendants)} descendant {descendants_object_type}s' ) return self - def unsuppress(self, suppression_reason=None, save=True, propagate=True): + def unsuppress(self, save=True, propagate=True): self.is_suppressed = False self.suppressed_at = None self.suppression_reason = None if save: self.save() if propagate: - for descendant_relation in self.suppressible_descendants: - descendants_to_unsuppress: QuerySet = getattr( - self, - descendant_relation - ).filter( - is_suppressed=True, - suppression_reason=self.Reason.ANCESTOR_SUPPRESSED - ) - for descendant in descendants_to_unsuppress: + descendant_tree = self._get_suppression_descendant_tree(suppressing=False) + for descendants_object_type, descendants in descendant_tree.items(): + for descendant in descendants: descendant.is_suppressed = False descendant.suppression_reason = None - descendants_to_unsuppress.bulk_update( - descendants_to_unsuppress, - [ - 'is_suppressed', - 'suppression_reason' - ] - ) + m: SuppressibleModel = apps.get_model(app_label='emgapi', model_name=descendants_object_type) + m.objects.bulk_update(descendants, fields=['is_suppressed', 'suppression_reason']) logger.info( - f'Propagated unsuppression of {self} ' - f'to {len(descendants_to_unsuppress)} {descendant_relation} descendants' + f'Propagated unsuppression of {self} to {len(descendants)} descendant {descendants_object_type}s' ) return self @@ -154,9 +205,11 @@ class Meta: class ENASyncableModel(SuppressibleModel, PrivacyControlledModel): - def sync_with_ena_status(self, ena_model_status: ENAStatus): + def sync_with_ena_status(self, ena_model_status: ENAStatus, propagate=True): """Sync the model with the ENA status accordingly. Fields that are updated: is_suppressed, suppressed_at, reason and is_private + + :propagate: If True, propagate the ena status of this entity to entities that are derived from / children of it. """ if ena_model_status == ENAStatus.PRIVATE and not self.is_private: self.is_private = True @@ -197,7 +250,7 @@ def sync_with_ena_status(self, ena_model_status: ENAStatus): elif ena_model_status == ENAStatus.CANCELLED: reason = SuppressibleModel.Reason.CANCELLED - self.suppress(suppression_reason=reason, save=False) + self.suppress(suppression_reason=reason, save=False, propagate=propagate) logging.info( f"{self} was suppressed, status on ENA {ena_model_status}" @@ -1115,7 +1168,7 @@ class Meta: class SampleQuerySet(BaseQuerySet, SuppressQuerySet): - + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -1334,7 +1387,7 @@ def __str__(self): class RunQuerySet(BaseQuerySet, SuppressQuerySet): - + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -1757,7 +1810,7 @@ class VariableNames(models.Model): class Meta: db_table = 'VARIABLE_NAMES' unique_together = (('var_id', 'var_name'), ('var_id', 'var_name'),) - verbose_name = 'variable name' + verbose_name = 'variable name' def __str__(self): return self.var_name diff --git a/emgena/management/commands/sync_assemblies_with_ena.py b/emgena/management/commands/sync_assemblies_with_ena.py index 63e45b7aa..76a8378e6 100644 --- a/emgena/management/commands/sync_assemblies_with_ena.py +++ b/emgena/management/commands/sync_assemblies_with_ena.py @@ -71,11 +71,7 @@ def handle(self, *args, **kwargs): ) continue - # inherits the status of its study - if study.is_suppressed: - emg_assembly.suppress( - suppression_reason=study.suppression_reason - ) + # inherits the privacy status of its study emg_assembly.is_private = study.is_private continue elif ena_assembly.status_id is None: diff --git a/emgena/management/commands/sync_runs_with_ena.py b/emgena/management/commands/sync_runs_with_ena.py index bd19f2c96..3adafbd6d 100644 --- a/emgena/management/commands/sync_runs_with_ena.py +++ b/emgena/management/commands/sync_runs_with_ena.py @@ -15,11 +15,8 @@ # limitations under the License. import logging -import os -from django.db.models import Count from django.core.management import BaseCommand -from django.conf import settings from emgapi import models as emg_models from emgena import models as ena_models @@ -38,7 +35,7 @@ def handle(self, *args, **kwargs): logging.info(f"Total Runs on EMG {runs_count}") while offset < runs_count: - emg_runs_batch = emg_models.Run.objects.all()[offset : offset + batch_size] + emg_runs_batch = emg_models.Run.objects.all()[offset: offset + batch_size] ena_runs_batch = ena_models.Run.objects.using("era").filter( run_id__in=[run.accession for run in emg_runs_batch] ) diff --git a/emgena/management/commands/sync_samples_with_ena.py b/emgena/management/commands/sync_samples_with_ena.py index cd1f5e477..659326193 100644 --- a/emgena/management/commands/sync_samples_with_ena.py +++ b/emgena/management/commands/sync_samples_with_ena.py @@ -15,11 +15,8 @@ # limitations under the License. import logging -import os -from django.db.models import Count from django.core.management import BaseCommand -from django.conf import settings from emgapi import models as emg_models from emgena import models as ena_models diff --git a/emgena/management/commands/sync_studies_with_ena.py b/emgena/management/commands/sync_studies_with_ena.py index 8fec5b5b2..9b28b8d6f 100644 --- a/emgena/management/commands/sync_studies_with_ena.py +++ b/emgena/management/commands/sync_studies_with_ena.py @@ -15,11 +15,8 @@ # limitations under the License. import logging -import os -from django.db.models import Count from django.core.management import BaseCommand -from django.conf import settings from emgapi import models as emg_models from emgena import models as ena_models diff --git a/tests/ena/test_sync_assemblies_with_ena.py b/tests/ena/test_sync_assemblies_with_ena.py index 13d678b69..d1ae22789 100644 --- a/tests/ena/test_sync_assemblies_with_ena.py +++ b/tests/ena/test_sync_assemblies_with_ena.py @@ -21,7 +21,7 @@ from django.urls import reverse from django.core.management import call_command -from emgapi.models import Assembly +from emgapi.models import Assembly, AnalysisJob from test_utils.emg_fixtures import * # noqa @@ -117,3 +117,21 @@ def test_sync_assemblies_based_on_study( assembly.refresh_from_db() assert assembly.is_suppressed == False assert assembly.is_private == False + + + @patch("emgena.models.Assembly.objects") + def test_sync_assemblies_propagation( + self, ena_assembly_objs_mock, ena_suppression_propagation_assemblies + ): + ena_assembly_objs_mock.using("ena").filter.return_value = ena_suppression_propagation_assemblies + + assert Assembly.objects.filter(is_suppressed=True).count() == 0 + assert AnalysisJob.objects.filter(is_suppressed=True).count() == 0 + + call_command("sync_assemblies_with_ena") + + assert Assembly.objects.filter(is_suppressed=True).count() == 32 + assert AnalysisJob.objects.filter( + is_suppressed=True, + suppression_reason=AnalysisJob.Reason.ANCESTOR_SUPPRESSED + ).count() == 64 diff --git a/tests/ena/test_sync_runs_with_ena.py b/tests/ena/test_sync_runs_with_ena.py index 84a1c4136..85d7ccb6d 100644 --- a/tests/ena/test_sync_runs_with_ena.py +++ b/tests/ena/test_sync_runs_with_ena.py @@ -15,13 +15,11 @@ # limitations under the License. import pytest -import os from unittest.mock import patch -from django.urls import reverse from django.core.management import call_command -from emgapi.models import Run +from emgapi.models import Run, Assembly, AnalysisJob from test_utils.emg_fixtures import * # noqa @@ -63,7 +61,7 @@ def test_make_runs_public(self, ena_run_objs_mock, ena_public_runs): assert run.is_private == False @patch("emgena.models.Run.objects") - def test_suppress_studies(self, ena_run_objs_mock, ena_suppressed_runs): + def test_suppress_runs(self, ena_run_objs_mock, ena_suppressed_runs): ena_run_objs_mock.using("era").filter.return_value = ena_suppressed_runs suppress_runs = Run.objects.order_by("?").all()[0:5] @@ -87,3 +85,25 @@ def test_suppress_studies(self, ena_run_objs_mock, ena_suppressed_runs): ena_run.get_status_id_display().lower() == run.get_suppression_reason_display().lower() ) + + @patch("emgena.models.Run.objects") + def test_suppress_runs_propagation(self, ena_run_objs_mock, ena_suppression_propagation_runs): + ena_run_objs_mock.using("era").filter.return_value = ena_suppression_propagation_runs + + runs = Run.objects.order_by("?").all() + assemblies = Assembly.objects.all() + analyses = AnalysisJob.objects.all() + for run in runs: + assert not run.is_suppressed + for assembly in assemblies: + assert not assembly.is_suppressed + for analysis in analyses: + assert not analysis.is_suppressed + + call_command("sync_runs_with_ena") + assert Run.objects.filter(is_suppressed=True).count() == 2 + assert Run.objects.filter(is_suppressed=False).count() == 2 + assert Assembly.objects.filter(is_suppressed=True).count() == 4 + assert Assembly.objects.filter(is_suppressed=False).count() == 4 + assert AnalysisJob.objects.filter(is_suppressed=True).count() == 8 + assert AnalysisJob.objects.filter(is_suppressed=False).count() == 8 diff --git a/tests/ena/test_sync_samples_with_ena.py b/tests/ena/test_sync_samples_with_ena.py index 0653f1ba2..9c2d539b0 100644 --- a/tests/ena/test_sync_samples_with_ena.py +++ b/tests/ena/test_sync_samples_with_ena.py @@ -21,7 +21,7 @@ from django.urls import reverse from django.core.management import call_command -from emgapi.models import Sample +from emgapi.models import Sample, Assembly, AnalysisJob, Run from test_utils.emg_fixtures import * # noqa @@ -89,3 +89,19 @@ def test_suppress_samples(self, ena_sample_objs_mock, ena_suppressed_samples): ena_sample.get_status_id_display().lower() == sample.get_suppression_reason_display().lower() ) + + @patch("emgena.models.Sample.objects") + def test_suppress_samples_propagation(self, ena_sample_objs_mock, ena_suppression_propagation_samples): + ena_sample_objs_mock.using("era").filter.return_value = ena_suppression_propagation_samples + + assert Sample.objects.filter(is_suppressed=True).count() == 0 + assert Run.objects.filter(is_suppressed=True).count() == 0 + assert Assembly.objects.filter(is_suppressed=True).count() == 0 + assert AnalysisJob.objects.filter(is_suppressed=True).count() == 0 + + call_command("sync_samples_with_ena") + + assert Sample.objects.filter(is_suppressed=True).count() == 2 + assert Run.objects.filter(is_suppressed=True).count() == 4 + assert Assembly.objects.filter(is_suppressed=True).count() == 8 + assert AnalysisJob.objects.filter(is_suppressed=True).count() == 16 diff --git a/tests/ena/test_sync_studies_with_ena.py b/tests/ena/test_sync_studies_with_ena.py index b841c7440..392a32df6 100644 --- a/tests/ena/test_sync_studies_with_ena.py +++ b/tests/ena/test_sync_studies_with_ena.py @@ -1,7 +1,6 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- - -# Copyright 2017-2022 EMBL - European Bioinformatics Institute +# Copyright 2017-2023 EMBL - European Bioinformatics Institute # # Licensed under the Apache License, Version 2.0 (the 'License'); # you may not use this file except in compliance with the License. @@ -15,16 +14,16 @@ # limitations under the License. import pytest -import os from unittest.mock import patch -from django.urls import reverse from django.core.management import call_command from emgapi.models import Study from test_utils.emg_fixtures import * # noqa +from emgena.models import Status + @pytest.mark.django_db class TestSyncENAStudies: @@ -34,7 +33,6 @@ def test_make_studies_private(self, ena_study_objs_mock, ena_private_studies): all_studies = Study.objects.order_by("?").all() public_studies = all_studies[0:5] - reminder = all_studies[5 : len(all_studies)] for study in public_studies: study.is_private = False @@ -55,7 +53,6 @@ def test_make_studies_public(self, ena_study_objs_mock, ena_public_studies): all_studies = Study.objects.order_by("?").all() private_studies = all_studies[0:5] - reminder = all_studies[5 : len(all_studies)] for study in private_studies: study.is_private = True @@ -96,3 +93,21 @@ def test_suppress_studies(self, ena_study_objs_mock, ena_suppressed_studies): ena_study.get_study_status_display().lower() == study.get_suppression_reason_display().lower() ) + + @patch("emgena.models.Study.objects") + def test_suppress_studies_propagation(self, ena_study_objs_mock, ena_suppression_propagation_studies): + ena_study_objs_mock.using("era").filter.return_value = ena_suppression_propagation_studies + + call_command("sync_studies_with_ena") + + for ena_study in ena_suppression_propagation_studies: + emg_study = Study.objects.get(secondary_accession=ena_study.study_id) + if ena_study.study_status == Status.SUPPRESSED: + assert emg_study.is_suppressed + for descendant in ['runs', 'samples', 'assemblies', 'analyses']: + related_qs = getattr(emg_study, descendant) + assert related_qs.filter(is_suppressed=True).exists() + # 1 sample was shared with unsuppressed studies so should remain + assert related_qs.filter(is_suppressed=False).count() == (1 if descendant == 'samples' else 0) + assert (related_qs.filter(is_suppressed=True).count() == + related_qs.filter(suppression_reason=Study.Reason.ANCESTOR_SUPPRESSED).count()) diff --git a/tests/test_utils/emg_fixtures.py b/tests/test_utils/emg_fixtures.py index f6bb03a87..17e5b706a 100644 --- a/tests/test_utils/emg_fixtures.py +++ b/tests/test_utils/emg_fixtures.py @@ -1,7 +1,7 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -# Copyright 2019 EMBL - European Bioinformatics Institute +# Copyright 2019-2023 EMBL - European Bioinformatics Institute # # Licensed under the Apache License, Version 2.0 (the 'License'); # you may not use this file except in compliance with the License. @@ -18,7 +18,6 @@ import uuid from django.conf import settings -from django.contrib.auth import get_user_model from model_bakery import baker @@ -37,7 +36,8 @@ 'ena_private_studies', 'ena_suppressed_studies', 'ena_public_runs', 'ena_private_runs', 'ena_suppressed_runs', 'ena_public_samples', 'ena_private_samples', 'ena_suppressed_samples', 'ena_public_assemblies', 'ena_private_assemblies', 'ena_suppressed_assemblies', - 'assembly_extra_annotation', + 'assembly_extra_annotation', 'ena_suppression_propagation_studies', 'ena_suppression_propagation_runs', + 'ena_suppression_propagation_samples', 'ena_suppression_propagation_assemblies', ] @@ -688,6 +688,7 @@ def ena_run_study(): study.pubmed_id = "" return study + def make_suppresible_studies(quantity, emg_props=None, ena_props=None): emg_props = emg_props or {} ena_props = ena_props or {} @@ -700,7 +701,6 @@ def make_suppresible_studies(quantity, emg_props=None, ena_props=None): return ena_studies - @pytest.fixture def ena_public_studies(): return make_suppresible_studies( @@ -746,6 +746,63 @@ def ena_suppressed_studies(): ) return studies +@pytest.fixture +def ena_suppression_propagation_studies(experiment_type_assembly): + """Returns: + 2 Studies that are suppressed in ENA but not in EMG, with unsuppressed sample/assembly/run + 2 Studies that are unsuprressed in ENA and in EMG, with unsuppressed sample/assembly/run + One sample in the ENA-suppressed studies is ALSO in an unsuppressed study. + """ + ena_studies_to_be_suppressed = make_suppresible_studies( + 2, + emg_props={}, + ena_props={"study_status": ena_models.Status.SUPPRESSED} + ) + ena_studies_to_remain = make_suppresible_studies( + 2, + emg_props={}, + ena_props={} + ) + for study in emg_models.Study.objects.all(): + samples = baker.make( + emg_models.Sample, + _quantity=2, + studies=[study] + ) + for sample in samples: + runs = baker.make( + emg_models.Run, + _quantity=2, + study=study, + sample=sample + ) + for run in runs: + assemblies = baker.make( + emg_models.Assembly, + _quantity=2, + runs=[run], + samples=[run.sample], + study=study, + experiment_type=experiment_type_assembly + ) + for assembly in assemblies: + baker.make( + emg_models.AnalysisJob, + _quantity=2, + assembly=assembly, + sample=assembly.samples.first(), + study=study, + ) + emg_study_to_suppress = emg_models.Study.objects.get(secondary_accession=ena_studies_to_be_suppressed[0].study_id) + # Share one sample across all studies + sample = emg_study_to_suppress.samples.first() + for study in emg_models.Study.objects.all(): + if not study.samples.filter(sample_id=sample.sample_id).exists(): + emg_models.StudySample.objects.create(study=study, sample=sample) + + return ena_studies_to_be_suppressed + ena_studies_to_remain + + def make_suppresible_runs(quantity, emg_props=None, ena_props=None): emg_props = emg_props or {} ena_props = ena_props or {} @@ -804,6 +861,49 @@ def ena_suppressed_runs(): ) return runs + +@pytest.fixture +def ena_suppression_propagation_runs(experiment_type_assembly, study): + runs = make_suppresible_runs( + 2, + emg_props={}, + ena_props={"status_id": ena_models.Status.SUPPRESSED}, + ) + runs.extend( + make_suppresible_runs( + 2, + emg_props={}, + ena_props={"status_id": ena_models.Status.PUBLIC} + ) + ) + + for run in emg_models.Run.objects.all(): + sample = baker.make( + emg_models.Sample, + _quantity=1, + studies=[study] + )[0] + run.sample = sample + run.save() + assemblies = baker.make( + emg_models.Assembly, + _quantity=2, + runs=[run], + study=study, + samples=[run.sample], + experiment_type=experiment_type_assembly + ) + for assembly in assemblies: + baker.make( + emg_models.AnalysisJob, + _quantity=2, + assembly=assembly, + sample=assembly.samples.first(), + study=study, + ) + return runs + + def make_suppresible_samples(quantity, emg_props=None, ena_props=None): emg_props = emg_props or {} ena_props = ena_props or {} @@ -862,6 +962,50 @@ def ena_suppressed_samples(): ) return samples + +@pytest.fixture +def ena_suppression_propagation_samples(experiment_type_assembly, study): + samples = make_suppresible_samples( + 2, + emg_props={}, + ena_props={"status_id": ena_models.Status.SUPPRESSED}, + ) + samples.extend( + make_suppresible_samples( + 2, + emg_props={}, + ena_props={"status_id": ena_models.Status.PUBLIC}, + ) + ) + + for sample in emg_models.Sample.objects.all(): + runs = baker.make( + emg_models.Run, + _quantity=2, + study=study, + sample=sample + ) + for run in runs: + assemblies = baker.make( + emg_models.Assembly, + _quantity=2, + runs=[run], + samples=[run.sample], + study=study, + experiment_type=experiment_type_assembly + ) + for assembly in assemblies: + baker.make( + emg_models.AnalysisJob, + _quantity=2, + assembly=assembly, + sample=assembly.samples.first(), + study=study, + ) + + return samples + + def make_suppresible_assemblies(quantity, emg_props=None, ena_props=None): emg_props = emg_props or {} ena_props = ena_props or {} @@ -920,3 +1064,32 @@ def ena_suppressed_assemblies(): ) ) return assemblies + + +@pytest.fixture +def ena_suppression_propagation_assemblies(experiment_type_assembly, study): + assemblies = [] + emg_props = {} + assemblies.extend( + make_suppresible_assemblies( + 32, + emg_props=emg_props, + ena_props={"status_id": ena_models.Status.SUPPRESSED}, + ) + ) + assemblies.extend( + make_suppresible_assemblies( + 32, + emg_props=emg_props, + ena_props={"status_id": ena_models.Status.PUBLIC}, + ) + ) + for assembly in emg_models.Assembly.objects.all(): + baker.make( + emg_models.AnalysisJob, + _quantity=2, + assembly=assembly, + sample=assembly.samples.first(), + study=study, + ) + return assemblies