Keep related on historical fields set to original model

Author: macro1

simple_history/models.py

@@ -281,8 +281,9 @@ def get_field(self, other, cls):
         return field
 
     def do_related_class(self, other, cls):
-        self.related = RelatedObject(other, cls, self)
         field = self.get_field(other, cls)
+        if not hasattr(self, 'related'):
+            self.related = RelatedObject(other, self.rel.field.model, self)
         transform_field(field)
         field.rel = None
 

simple_history/tests/tests/test_models.py

@@ -469,8 +469,8 @@ def test_invalid_bases(self):
 
     def test_import_related(self):
         field_object = HistoricalChoice._meta.get_field_by_name('poll_id')[0]
-        self.assertEqual(field_object.related.model, HistoricalChoice)
+        self.assertEqual(field_object.related.model, Choice)
 
     def test_string_related(self):
         field_object = HistoricalState._meta.get_field_by_name('library_id')[0]
-        self.assertEqual(field_object.related.model, HistoricalState)
+        self.assertEqual(field_object.related.model, State)