From cb6961539c3b5666c7c3b575b010e4dc7855da00 Mon Sep 17 00:00:00 2001
From: David Wobrock <david.wobrock@gmail.com>
Date: Sat, 8 Feb 2020 21:23:04 +0100
Subject: [PATCH] Handle for forward and backward data migration imports

---
 django_migration_linter/migration_linter.py | 47 ++++++++++++++++++++-
 tests/functional/test_data_migrations.py    | 22 ++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/django_migration_linter/migration_linter.py b/django_migration_linter/migration_linter.py
index fe28cfc..21e90be 100644
--- a/django_migration_linter/migration_linter.py
+++ b/django_migration_linter/migration_linter.py
@@ -421,5 +421,50 @@ class MigrationLinter(object):
             else:
                 warning.append(issue)
 
-        # TODO: test if importing model directly or correctly using "get_model()"
+        # Detect wrong model imports
+        # Forward
+        issues = self.get_data_migration_model_import_issues(runpython.code)
+        for issue in issues:
+            if issue["code"] in self.exclude_migration_tests:
+                ignored.append(issue)
+            else:
+                error.append(issue)
+
+        # Backward
+        if runpython.reversible:
+            issues = self.get_data_migration_model_import_issues(runpython.reverse_code)
+            for issue in issues:
+                if issue and issue["code"] in self.exclude_migration_tests:
+                    ignored.append(issue)
+                else:
+                    error.append(issue)
         return error, ignored, warning
+
+    @staticmethod
+    def get_data_migration_model_import_issues(code):
+        model_object_regex = re.compile(r"[^a-zA-Z]?([a-zA-Z0-9]+?)\.objects")
+
+        function_name = code.__name__
+        reverse_code_source = inspect.getsource(code)
+
+        called_models = model_object_regex.findall(reverse_code_source)
+        issues = []
+        for model in called_models:
+            has_get_model_call = (
+                re.search(
+                    r"get_model\(.*?,.*?{}.+?\)".format(model), reverse_code_source
+                )
+                is not None
+            )
+            if not has_get_model_call:
+                issues.append(
+                    {
+                        "code": "DATA_MIGRATION_MODEL_IMPORT",
+                        "msg": (
+                            "'{}': Could not find an 'apps.get_model(\"...\", \"{}\")' "
+                            "call. Importing the model directly is incorrect for "
+                            "data migrations."
+                        ).format(function_name, model),
+                    }
+                )
+        return issues
diff --git a/tests/functional/test_data_migrations.py b/tests/functional/test_data_migrations.py
index b9f1510..e3953a1 100644
--- a/tests/functional/test_data_migrations.py
+++ b/tests/functional/test_data_migrations.py
@@ -64,3 +64,25 @@ class DataMigrationDetectionTestCase(unittest.TestCase):
         self.assertEqual(0, self.linter.nb_warnings)
         self.assertEqual(1, self.linter.nb_erroneous)
         self.assertTrue(self.linter.has_errors)
+
+    def test_missing_get_model_import(self):
+        def incorrect_importing_model_forward(apps, schema_editor):
+            from tests.test_project.app_data_migrations.models import MyModel
+
+            MyModel.objects.filter(id=1).first()
+
+        issues = MigrationLinter.get_data_migration_model_import_issues(
+            incorrect_importing_model_forward
+        )
+        self.assertEqual(1, len(issues))
+
+    def test_correct_get_model_import(self):
+        def correct_importing_model_forward(apps, schema_editor):
+            MyModel = apps.get_model("app_data_migrations", "MyModel")
+
+            MyModel.objects.filter(id=1).first()
+
+        issues = MigrationLinter.get_data_migration_model_import_issues(
+            correct_importing_model_forward
+        )
+        self.assertEqual(0, len(issues))