From e036e366f519c1b4b2cc638783be5aa3a3a64475 Mon Sep 17 00:00:00 2001
From: hoangquochung1110 <arsenalforeversince2007@gmail.com>
Date: Tue, 2 Apr 2024 16:07:09 +0700
Subject: [PATCH] wip

---
 auditlog/conf.py         |  1 -
 auditlog/models.py       |  2 +-
 auditlog/registry.py     | 12 +++++++-----
 auditlog_tests/models.py | 16 ----------------
 auditlog_tests/tests.py  | 14 ++++++--------
 5 files changed, 14 insertions(+), 31 deletions(-)

diff --git a/auditlog/conf.py b/auditlog/conf.py
index 910b4573..2969edde 100644
--- a/auditlog/conf.py
+++ b/auditlog/conf.py
@@ -1,6 +1,5 @@
 from django.conf import settings
 
-
 # Register all models when set to True
 settings.AUDITLOG_INCLUDE_ALL_MODELS = getattr(
     settings, "AUDITLOG_INCLUDE_ALL_MODELS", False
diff --git a/auditlog/models.py b/auditlog/models.py
index 9af4d1d0..dc591f60 100644
--- a/auditlog/models.py
+++ b/auditlog/models.py
@@ -445,7 +445,7 @@ def changes_display_dict(self):
         auditlog = AuditlogRegistry.get_auditlog(model=model)
 
         model_fields = None
-        if auditlog.contains(model._meta.model):
+        if auditlog:
             model_fields = auditlog.get_model_fields(model._meta.model)
 
         changes_display_dict = {}
diff --git a/auditlog/registry.py b/auditlog/registry.py
index 8a789c39..520b21aa 100644
--- a/auditlog/registry.py
+++ b/auditlog/registry.py
@@ -401,9 +401,9 @@ def get_auditlog(cls, model: ModelBase):
         Retrieve the auditlog object associated with a given model.
         """
         auditlogs = get_default_auditlogs()
-        return next(
-            (auditlog for auditlog in auditlogs if auditlog._registry.get(model)), None
-        )
+        for auditlog in auditlogs:
+            if auditlog.contains(model):
+                return auditlog
 
     @classmethod
     def get_model_fields(cls, model: ModelBase):
@@ -411,8 +411,10 @@ def get_model_fields(cls, model: ModelBase):
         Wrapper of AuditlogModelRegistry.get_model_fields().
         """
         auditlog = cls.get_auditlog(model)
-        model_fields = auditlog.get_model_fields(model)
-        return model_fields
+        if auditlog:
+            model_fields = auditlog.get_model_fields(model)
+            return model_fields
+        return {}
 
 
 auditlog_registry = AuditlogRegistry()
diff --git a/auditlog_tests/models.py b/auditlog_tests/models.py
index 72d85f55..1e16f933 100644
--- a/auditlog_tests/models.py
+++ b/auditlog_tests/models.py
@@ -330,22 +330,6 @@ class AutoManyRelatedModel(models.Model):
     related = models.ManyToManyField(SimpleModel)
 
 
-class CreateUpdateOnlyModel(models.Model):
-    """
-    A simple model to test create_only_auditlog and update_only_auditlog.
-    """
-
-    text = models.TextField(blank=True)
-    boolean = models.BooleanField(default=False)
-    integer = models.IntegerField(blank=True, null=True)
-    datetime = models.DateTimeField(auto_now=True)
-
-    history = AuditlogHistoryField(delete_related=True)
-
-    def __str__(self):
-        return self.text
-
-
 auditlog.register(AltPrimaryKeyModel)
 auditlog.register(UUIDPrimaryKeyModel)
 auditlog.register(ProxyModel)
diff --git a/auditlog_tests/tests.py b/auditlog_tests/tests.py
index 85b6a38e..4b9fde32 100644
--- a/auditlog_tests/tests.py
+++ b/auditlog_tests/tests.py
@@ -39,7 +39,6 @@
     AutoManyRelatedModel,
     CharfieldTextfieldModel,
     ChoicesFieldModel,
-    CreateUpdateOnlyModel,
     DateTimeFieldModel,
     JSONModel,
     ManyRelatedModel,
@@ -2650,15 +2649,18 @@ def test_get_changes_for_missing_model(self):
 class CustomAuditlogTest(TestCase):
     def setUp(self):
         super().setUp()
+        if auditlog.contains(SimpleModel):
+            auditlog.unregister(SimpleModel)
 
     def tearDown(self):
         for model in create_only_auditlog.get_models():
             create_only_auditlog.unregister(model)
         for model in update_only_auditlog.get_models():
             update_only_auditlog.unregister(model)
+        auditlog.register(SimpleModel)
 
     def make_object(self):
-        return CreateUpdateOnlyModel.objects.create(text="I am not difficult.")
+        return SimpleModel.objects.create(text="I am not difficult.")
 
     def update_object(self, obj):
         obj.text = "Changed"
@@ -2671,9 +2673,7 @@ def test_create_only_auditlog(self):
                 "auditlog_tests.test_registry.create_only_auditlog",
             ]
         ):
-            create_only_auditlog.register(
-                CreateUpdateOnlyModel, include_fields=["text"]
-            )
+            create_only_auditlog.register(SimpleModel, include_fields=["text"])
 
             # Get the object to work with
             obj = self.make_object()
@@ -2693,9 +2693,7 @@ def test_update_only_auditlog(self):
                 "auditlog_tests.test_registry.update_only_auditlog",
             ]
         ):
-            update_only_auditlog.register(
-                CreateUpdateOnlyModel, include_fields=["text"]
-            )
+            update_only_auditlog.register(SimpleModel, include_fields=["text"])
 
             # Get the object to work with
             obj = self.make_object()