[Enhancement] Swappable Packages (#216)

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#216
Reviewed-by: liambrydon <lbry121@aucklanduni.ac.nz>
Reviewed-by: Tobias O <tobias.olenyi@envipath.com>

tests/test_enviformer.py

@@ -1,10 +1,15 @@
 from collections import defaultdict
 from datetime import datetime
 from tempfile import TemporaryDirectory
+
+from django.conf import settings as s
 from django.test import TestCase, tag
+
 from epdb.logic import PackageManager
-from epdb.models import User, EnviFormer, Package, Setting
-from epdb.tasks import predict_simple, predict
+from epdb.models import EnviFormer, Setting, User
+from epdb.tasks import predict, predict_simple
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 def measure_predict(mod, pathway_pk=None):
@@ -68,11 +73,15 @@ class EnviFormerTest(TestCase):
 
                 # Test pathway prediction
                 times = [measure_predict(mods[1], self.BBD_SUBSET.pathways[0].pk) for _ in range(5)]
-                print(f"First pathway prediction took {times[0]} seconds, subsequent ones took {times[1:]}")
+                print(
+                    f"First pathway prediction took {times[0]} seconds, subsequent ones took {times[1:]}"
+                )
 
                 # Test eviction by performing three prediction with every model, twice.
                 times = defaultdict(list)
-                for _ in range(2):  # Eviction should cause the second iteration here to have to reload the models
+                for _ in range(
+                    2
+                ):  # Eviction should cause the second iteration here to have to reload the models
                     for mod in mods:
                         for _ in range(3):
                             times[mod.pk].append(measure_predict(mod))

tests/test_model.py

@@ -1,10 +1,13 @@
 from tempfile import TemporaryDirectory
 
 import numpy as np
+from django.conf import settings as s
 from django.test import TestCase
 
 from epdb.logic import PackageManager
-from epdb.models import User, MLRelativeReasoning, Package, RuleBasedRelativeReasoning
+from epdb.models import MLRelativeReasoning, RuleBasedRelativeReasoning, User
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class ModelTest(TestCase):
@@ -85,7 +88,7 @@ class ModelTest(TestCase):
                 mod.build_model()
                 mod.evaluate_model(True, eval_packages_objs, n_splits=2)
 
-                results = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
+                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
 
     def test_rbrr(self):
         with TemporaryDirectory() as tmpdir:
@@ -103,12 +106,12 @@ class ModelTest(TestCase):
                     threshold=threshold,
                     min_count=5,
                     max_count=0,
-                    name='ECC - BBD - 0.5',
-                    description='Created MLRelativeReasoning in Testcase',
+                    name="ECC - BBD - 0.5",
+                    description="Created MLRelativeReasoning in Testcase",
                 )
 
                 mod.build_dataset()
                 mod.build_model()
                 mod.evaluate_model(True, eval_packages_objs, n_splits=2)
 
-                results = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
+                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")

tests/test_multigen_eval.py

@@ -1,8 +1,12 @@
+from django.conf import settings as s
 from django.test import TestCase
 from networkx.utils.misc import graphs_equal
+
 from epdb.logic import PackageManager, SPathway
-from epdb.models import Pathway, User, Package
-from utilities.ml import multigen_eval, pathway_edit_eval, graph_from_pathway
+from epdb.models import Pathway, User
+from utilities.ml import graph_from_pathway, multigen_eval, pathway_edit_eval
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class MultiGenTest(TestCase):

tests/test_simpleambitrule.py

@@ -1,9 +1,10 @@
-from unittest.mock import patch, MagicMock, PropertyMock
+from unittest.mock import MagicMock, PropertyMock, patch
 
+from django.conf import settings as s
 from django.test import TestCase
 
 from epdb.logic import PackageManager
-from epdb.models import User, SimpleAmbitRule
+from epdb.models import SimpleAmbitRule, User
 
 
 class SimpleAmbitRuleTest(TestCase):
@@ -209,7 +210,7 @@ class SimpleAmbitRuleTest(TestCase):
 
         self.assertEqual(rule.products_smarts, expected_products)
 
-    @patch("epdb.models.Package.objects")
+    @patch(f"{s.EPDB_PACKAGE_MODEL.replace('.', '.models.')}.objects")
     def test_related_reactions_property(self, mock_package_objects):
         """Test related_reactions property returns correct queryset."""
         mock_qs = MagicMock()

tests/views/test_model_views.py

@@ -1,9 +1,11 @@
+from django.conf import settings as s
 from django.test import TestCase, override_settings
 from django.urls import reverse
-from django.conf import settings as s
 
 from epdb.logic import UserManager
-from epdb.models import Package, User
+from epdb.models import User
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 @override_settings(MODEL_DIR=s.FIXTURE_DIRS[0] / "models", CELERY_TASK_ALWAYS_EAGER=True)

tests/views/test_package_views.py

@@ -5,14 +5,15 @@ from django.urls import reverse
 
 from epdb.logic import UserManager
 from epdb.models import (
-    Package,
-    UserPackagePermission,
-    Permission,
-    GroupPackagePermission,
     Group,
+    GroupPackagePermission,
     License,
+    Permission,
+    UserPackagePermission,
 )
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 class PackageViewTest(TestCase):
     fixtures = ["test_fixtures.jsonl.gz"]