From 71afdd8b36b0dc013ece774225f248b456ab5d81 Mon Sep 17 00:00:00 2001
From: mpeltriaux <michel.peltriaux@sgdnord.rlp.de>
Date: Wed, 11 May 2022 10:16:34 +0200
Subject: [PATCH] JS Tree enhancement

* extends compensation state forms to match the new logic
* adds minor changes for tests
---
 compensation/forms/modalForms.py    | 18 +++++++++++++++---
 compensation/models/compensation.py |  6 +++++-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/compensation/forms/modalForms.py b/compensation/forms/modalForms.py
index d129459d..581a7b3a 100644
--- a/compensation/forms/modalForms.py
+++ b/compensation/forms/modalForms.py
@@ -157,7 +157,7 @@ class NewStateModalForm(BaseModalForm):
     What has been on this area before changes/compensations have been applied and what will be the result ('after')?
 
     """
-    biotope_type = forms.MultipleChoiceField(
+    biotope_type = forms.ChoiceField(
         label=_("Biotope Type"),
         label_suffix="",
         required=True,
@@ -200,6 +200,16 @@ class NewStateModalForm(BaseModalForm):
         super().__init__(*args, **kwargs)
         self.form_title = _("New state")
         self.form_caption = _("Insert data for the new state")
+        choices = KonovaCode.objects.filter(
+            code_lists__in=[CODELIST_BIOTOPES_ID],
+            is_archived=False,
+            is_leaf=True,
+        ).values_list("id", flat=True)
+        choices = [
+            (choice, choice)
+            for choice in choices
+        ]
+        self.fields["biotope_type"].choices = choices
 
     def save(self, is_before_state: bool = False):
         state = self.instance.add_state(self, is_before_state)
@@ -262,8 +272,9 @@ class EditCompensationStateModalForm(NewStateModalForm):
         self.state = kwargs.pop("state", None)
         super().__init__(*args, **kwargs)
         self.form_title = _("Edit state")
+        biotope_type_id = self.state.biotope_type.id if self.state.biotope_type else None
         form_data = {
-            "biotope_type": self.state.biotope_type.id,
+            "biotope_type": biotope_type_id,
             "biotope_extra": self.state.biotope_type_details.all(),
             "surface": self.state.surface,
         }
@@ -271,7 +282,8 @@ class EditCompensationStateModalForm(NewStateModalForm):
 
     def save(self, is_before_state: bool = False):
         state = self.state
-        state.biotope_type = self.cleaned_data.get("biotope_type", None)
+        biotope_type_id = self.cleaned_data.get("biotope_type", None)
+        state.biotope_type = KonovaCode.objects.get(id=biotope_type_id)
         state.biotope_type_details.set(self.cleaned_data.get("biotope_extra", []))
         state.surface = self.cleaned_data.get("surface", None)
         state.save()
diff --git a/compensation/models/compensation.py b/compensation/models/compensation.py
index 3deff50a..3727e0c2 100644
--- a/compensation/models/compensation.py
+++ b/compensation/models/compensation.py
@@ -8,6 +8,8 @@ Created on: 16.11.21
 import shutil
 
 from django.contrib import messages
+
+from codelist.models import KonovaCode
 from user.models import User, Team
 from django.db import models, transaction
 from django.db.models import QuerySet, Sum
@@ -142,8 +144,10 @@ class AbstractCompensation(BaseObject, GeoReferencedMixin):
         """
         form_data = form.cleaned_data
         with transaction.atomic():
+            biotope_type_id = form_data["biotope_type"]
+            code = KonovaCode.objects.get(id=biotope_type_id)
             state = CompensationState.objects.create(
-                biotope_type=form_data["biotope_type"],
+                biotope_type=code,
                 surface=form_data["surface"],
             )
             state_additional_types = form_data["biotope_extra"]