Graph generation: fix non full generation with changing relation type - attach full to upper item

4 files changed, 43 insertions, 5 deletions

diff --git a/archaeological_context_records/models.py b/archaeological_context_records/models.py
index 1b81f3b63..b0e72c4a1 100644
--- a/archaeological_context_records/models.py
+++ b/archaeological_context_records/models.py
@@ -421,6 +421,7 @@ class ContextRecord(BulkUpdatedItem, DocumentItem, BaseHistorizedItem,
             pgettext_lazy("key for text search", "site"),
             'cached_label'),
     }
+    MAIN_UP_MODEL_QUERY = "operation"
     RELATIVE_SESSION_NAMES = [
         ('operation', 'operation__pk'),
         ('site', 'archaeological_site__pk'),
diff --git a/archaeological_context_records/tests.py b/archaeological_context_records/tests.py
index a0a76ba78..9e16603de 100644
--- a/archaeological_context_records/tests.py
+++ b/archaeological_context_records/tests.py
@@ -886,6 +886,8 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         self.create_context_record({"label": u"CR 1B"})
         self.create_context_record({"label": u"CR 2B"})
         self.create_context_record({"label": u"CR 3B"})
+        self.create_context_record({"label": u"CR 2C"})
+        self.create_context_record({"label": u"CR 2D"})
 
         cr_1 = self.context_records[0]
         self.cr_2 = cr_2 = self.context_records[1]
@@ -894,6 +896,8 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         cr_1B = self.context_records[4]
         cr_2B = self.context_records[5]
         cr_3B = self.context_records[6]
+        cr_2C = self.context_records[7]
+        cr_2D = self.context_records[8]
 
         models.RecordRelations.objects.create(
             left_record=cr_2, right_record=cr_2B,
@@ -915,8 +919,14 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         models.RecordRelations.objects.create(
             left_record=cr_3B, right_record=cr_2B,
             relation_type=self.rel_type_above)
+        models.RecordRelations.objects.create(
+            left_record=cr_3B, right_record=cr_2C,
+            relation_type=self.rel_type_above)
+        models.RecordRelations.objects.create(
+            left_record=cr_2D, right_record=cr_3B,
+            relation_type=self.rel_type_bellow)
 
-    def test_gen_relation_full(self):
+    def test_gen_relation(self):
         generate_relation_graph(self.cr_2)
         cr_2 = models.ContextRecord.objects.get(pk=self.cr_2.pk)
         self.assertIsNotNone(cr_2.relation_image)
@@ -930,6 +940,26 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         self.assertIn('"CR 3"', content)
         self.assertIn('"CR 3B"', content)
         self.assertIn('"CR 4"', content)
+        self.assertNotIn('"CR 2C"', content)
+        self.assertNotIn('"CR 2D"', content)
+
+    def test_gen_relation_full(self):
+        generate_relation_graph(self.cr_2, full=True)
+        cr_2 = models.ContextRecord.objects.get(pk=self.cr_2.pk)
+        operation = cr_2.operation
+        self.assertIsNotNone(operation.relation_image)
+        self.assertIsNotNone(operation.relation_bitmap_image)
+        self.assertIsNotNone(operation.relation_dot)
+        content = open(operation.relation_dot.path).read()
+        self.assertIn('"CR 1"', content)
+        self.assertIn('"CR 1B"', content)
+        self.assertIn('"CR 2B"', content)
+        self.assertIn('"CR 2",style=filled,fillcolor="#C6C0C0"', content)
+        self.assertIn('"CR 3"', content)
+        self.assertIn('"CR 3B"', content)
+        self.assertIn('"CR 4"', content)
+        self.assertIn('"CR 2C"', content)
+        self.assertIn('"CR 2D"', content)
 
     def test_gen_relation_above(self):
         generate_relation_graph(self.cr_2, render_above=False)
@@ -945,6 +975,8 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         self.assertNotIn('"CR 3B"', content)
         self.assertNotIn('"CR 3"', content)
         self.assertNotIn('"CR 4"', content)
+        self.assertNotIn('"CR 2C"', content)
+        self.assertNotIn('"CR 2D"', content)
 
     def test_gen_relation_bellow(self):
         generate_relation_graph(self.cr_2, render_bellow=False)
@@ -960,5 +992,7 @@ class GraphGenerationTest(ContextRecordInit, TestCase):
         self.assertIn('"CR 3"', content)
         self.assertIn('"CR 3B"', content)
         self.assertIn('"CR 4"', content)
+        self.assertNotIn('"CR 2C"', content)
+        self.assertNotIn('"CR 2D"', content)
 
 
diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index b2f235207..872f740e6 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -1201,6 +1201,7 @@ class RelationItem(models.Model):
     """
     Items with relation between them
     """
+    MAIN_UP_MODEL_QUERY = ""
     relation_image = models.FileField(
         _("Generated relation image (SVG)"), null=True, blank=True,
         upload_to=get_image_path, help_text=max_size_help()
diff --git a/ishtar_common/utils.py b/ishtar_common/utils.py
index ebfbaaa77..eb8414ced 100644
--- a/ishtar_common/utils.py
+++ b/ishtar_common/utils.py
@@ -1160,18 +1160,16 @@ def get_relations_for_graph(rel_model, obj_pk, above_relations=None,
                 other_record = left_record
             else:
                 other_record = right_record
-
-                render_bellow = False
             if get_above and get_bellow and not full and (is_bellow or
                                                           is_above):
-                if is_above:
+                if (is_above and not inverse) or (is_bellow and inverse):
                     ar, er, substyles = get_relations_for_graph(
                         rel_model, other_record, above_relations,
                         equal_relations, treated, styles,
                         render_above=True,
                         render_bellow=False
                     )
-                elif is_bellow:
+                else:
                     ar, er, substyles = get_relations_for_graph(
                         rel_model, other_record, above_relations,
                         equal_relations, treated, styles,
@@ -1278,6 +1276,10 @@ def generate_relation_graph(obj, highlight_current=True,
     else:
         suffix = ""
 
+    if full and obj.MAIN_UP_MODEL_QUERY and getattr(obj,
+                                                    obj.MAIN_UP_MODEL_QUERY):
+        obj = getattr(obj, obj.MAIN_UP_MODEL_QUERY)
+
     with open(dot_name, "r") as dot_file:
         django_file = File(dot_file)
         attr = "relation_dot" + suffix