Treatment migratrion: init empty treatement state

1 files changed, 14 insertions, 6 deletions

diff --git a/archaeological_finds/migrations/0050_auto_20181211_1509.py b/archaeological_finds/migrations/0050_auto_20181211_1509.py
index c9928d617..40461569c 100644
--- a/archaeological_finds/migrations/0050_auto_20181211_1509.py
+++ b/archaeological_finds/migrations/0050_auto_20181211_1509.py
@@ -7,16 +7,19 @@ from django.db import migrations, models
 
 def migrate_treatment_states(apps, schema):
     TreatmentState = apps.get_model('archaeological_finds', 'TreatmentState')
-    q = TreatmentState.objects.filter(txt_idx="completed")
-    if q.count():
-        t = q.all()[0]
-        t.executed = True
-        t.save()
+    Treatment = apps.get_model('archaeological_finds', 'Treatment')
+    for finished in ['completed', 'finished']:
+        q = TreatmentState.objects.filter(txt_idx=finished)
+        if q.count():
+            t = q.all()[0]
+            t.executed = True
+            t.save()
+            break
     for t in TreatmentState.objects.all():
         t.txt_idx = t.txt_idx.replace("_", "-")
         t.save()
     orders = [(10, "planned"), (20, "to-be-confirmed"), (30, "in-progress"),
-              (40, "completed"), (50, "cancelled"), (60, "unknown"),]
+              (40, finished), (50, "cancelled"), (60, "unknown"),]
     for order, txt_idx in orders:
         q = TreatmentState.objects.filter(txt_idx=txt_idx)
         if not q.count():
@@ -24,6 +27,11 @@ def migrate_treatment_states(apps, schema):
         t = q.all()[0]
         t.order = order
         t.save()
+    if TreatmentState.objects.filter(txt_idx=finished).count():
+        f = TreatmentState.objects.filter(txt_idx=finished)
+        for t in Treatment.objects.filter(treatment_state__isnull=True):
+            t.treatment_state = f
+            t.save()
 
 
 class Migration(migrations.Migration):