updated populate fixing-investigation command

care/facility/tests/test_management_commands.py

@@ -0,0 +1,48 @@
+import json
+from pathlib import Path
+
+from django.core.management import call_command
+from django.test import TestCase
+
+from care.facility.models import PatientInvestigation, PatientInvestigationGroup
+
+
+class LoadInvestigationsCommandTest(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        call_command("populate_investigations", verbosity=0)
+
+        with Path("data/investigations.json").open() as investigations_data:
+            cls.investigations = json.load(investigations_data)
+
+        with Path("data/investigation_groups.json").open() as investigation_groups_data:
+            cls.investigation_groups = json.load(investigation_groups_data)
+
+    def test_number_of_entries(self):
+        self.assertEqual(len(self.investigations), PatientInvestigation.objects.count())
+        self.assertEqual(
+            len(self.investigation_groups), PatientInvestigationGroup.objects.count()
+        )
+
+    def test_relation_between_investigations_and_groups(self):
+        # taking first and last 5 data to test it out
+        test_investigation_data = self.investigations[:5] + self.investigations[-5:]
+
+        # creating a dictionary to avoid looping of group data for each check
+        group_dict = {
+            int(group["id"]): group["name"] for group in self.investigation_groups
+        }
+
+        for investigation_data in test_investigation_data:
+            investigation = PatientInvestigation.objects.get(
+                name=investigation_data["name"]
+            )
+
+            group_values = list(investigation.groups.values_list("name", flat=True))
+
+            expected_groups = [
+                group_dict[category_id]
+                for category_id in investigation_data["category_id"]
+            ]
+
+            self.assertCountEqual(group_values, expected_groups)

care/users/management/commands/populate_investigations.py

@@ -1,82 +1,98 @@
 import json
 from pathlib import Path
 
-from django.core.management import BaseCommand
+from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from care.facility.models.patient_investigation import (
-    PatientInvestigation,
-    PatientInvestigationGroup,
-)
+from care.facility.models import PatientInvestigation, PatientInvestigationGroup
 
-with Path("data/investigations.json").open() as investigations_data:
-    investigations = json.load(investigations_data)
 
-with Path("data/investigation_groups.json").open() as investigation_groups_data:
-    investigation_groups = json.load(investigation_groups_data)
+def load_json(file_path):
+    with Path(file_path).open() as json_file:
+        return json.load(json_file)
 
 
-class Command(BaseCommand):
-    """
-    Populates Investigation seed data
-    """
+def load_investigation_group_data(investigation_groups):
+    investigation_group_dict = {}
+    groups_to_create = []
+
+    for investigation_group in investigation_groups:
+        current_obj = PatientInvestigationGroup.objects.filter(
+            name=investigation_group["name"]
+        ).first()
+        if not current_obj:
+            current_obj = PatientInvestigationGroup(name=investigation_group["name"])
+            groups_to_create.append(current_obj)
+
+        investigation_group_dict[str(investigation_group["id"])] = current_obj
+
+    if groups_to_create:
+        PatientInvestigationGroup.objects.bulk_create(groups_to_create)
+
+    return investigation_group_dict
+
+
+def extract_investigation_data(investigation):
+    return {
+        "name": investigation["name"],
+        "unit": investigation.get("unit", ""),
+        "ideal_value": investigation.get("ideal_value", ""),
+        "min_value": parse_float(investigation.get("min")),
+        "max_value": parse_float(investigation.get("max")),
+        "investigation_type": investigation.get("type", None),
+        "choices": investigation.get("choices", ""),
+    }
+
+
+def parse_float(value):
+    try:
+        return float(value)
+    except (ValueError, TypeError):
+        return None
+
+
+def update_existing_investigation(existing_obj, data):
+    for field, value in data.items():
+        setattr(existing_obj, field, value)
 
+
+def handle_investigations(investigations, investigation_group_dict):
+    bulk_create_data = []
+    bulk_update_data = []
+    investigations_to_update_groups = []
+
+    for investigation in investigations:
+        data = extract_investigation_data(investigation)
+        existing_obj = PatientInvestigation.objects.filter(name=data["name"]).first()
+
+        if existing_obj:
+            update_existing_investigation(existing_obj, data)
+            bulk_update_data.append(existing_obj)
+            investigations_to_update_groups.append(
+                (existing_obj, investigation["category_id"])
+            )
+        else:
+            new_obj = PatientInvestigation(**data)
+            bulk_create_data.append(new_obj)
+            investigations_to_update_groups.append(
+                (new_obj, investigation["category_id"])
+            )
+
+    return bulk_create_data, bulk_update_data, investigations_to_update_groups
+
+
+class Command(BaseCommand):
     help = "Seed Data for Investigations"
 
     def handle(self, *args, **kwargs):
-        investigation_group_dict = {}
-
-        investigation_groups_to_create = [
-            PatientInvestigationGroup(name=group.get("name"))
-            for group in investigation_groups
-            if group.get("id") not in investigation_group_dict
-        ]
-        created_groups = PatientInvestigationGroup.objects.bulk_create(
-            investigation_groups_to_create
-        )
-        investigation_group_dict.update({group.id: group for group in created_groups})
+        investigation_groups = load_json("data/investigation_groups.json")
+        investigations = load_json("data/investigations.json")
 
-        existing_objs = PatientInvestigation.objects.filter(
-            name__in=[investigation["name"] for investigation in investigations]
-        )
+        investigation_group_dict = load_investigation_group_data(investigation_groups)
 
-        bulk_create_data = []
-        bulk_update_data = []
-
-        for investigation in investigations:
-            data = {
-                "name": investigation["name"],
-                "unit": investigation.get("unit", ""),
-                "ideal_value": investigation.get("ideal_value", ""),
-                "min_value": (
-                    None
-                    if investigation.get("min_value") is None
-                    else float(investigation.get("min_value"))
-                ),
-                "max_value": (
-                    None
-                    if investigation.get("max_value") is None
-                    else float(investigation.get("max_value"))
-                ),
-                "investigation_type": investigation["type"],
-                "choices": investigation.get("choices", ""),
-            }
-
-            existing_obj = existing_objs.filter(name=data["name"]).first()
-            if existing_obj:
-                bulk_update_data.append(existing_obj)
-            else:
-                new_obj = PatientInvestigation(**data)
-                bulk_create_data.append(new_obj)
-
-            group_ids = investigation.get("category_ids", [])
-            groups_to_add = [
-                investigation_group_dict[category_id] for category_id in group_ids
-            ]
-            if existing_obj:
-                existing_obj.groups.set(groups_to_add)
-            else:
-                data["groups"] = groups_to_add
+        bulk_create_data, bulk_update_data, investigations_to_update_groups = (
+            handle_investigations(investigations, investigation_group_dict)
+        )
 
         with transaction.atomic():
             if bulk_create_data:
@@ -94,3 +110,15 @@ def handle(self, *args, **kwargs):
                         "choices",
                     ],
                 )
+
+            for investigation_obj, category_ids in investigations_to_update_groups:
+                groups_to_add = [
+                    investigation_group_dict.get(str(category_id))
+                    for category_id in category_ids
+                ]
+                investigation_obj.groups.set(groups_to_add)
+
+        if kwargs.get("verbosity", 1) > 0:
+            self.stdout.write(
+                self.style.SUCCESS("Successfully populated investigation data")
+            )

data/investigation_groups.json

@@ -26,5 +26,57 @@
     {
         "id": "7",
         "name": "Kidney Function test"
+    },
+    {
+        "id": "8",
+        "name": "Electrolytes"
+    },
+    {
+        "id": "9",
+        "name": "Lipid Profile"
+    },
+    {
+        "id": "10",
+        "name": "ABG"
+    },
+    {
+        "id": "11",
+        "name": "Cardiac Enzyme"
+    },
+    {
+        "id": "12",
+        "name": "RFT"
+    },
+    {
+        "id": "13",
+        "name": "Blood Glucose Level"
+    },
+    {
+        "id": "14",
+        "name": "Immunoglobulins"
+    },
+    {
+        "id": "15",
+        "name": "Other Body Fluids"
+    },
+    {
+        "id": "16",
+        "name": "Cerebro Spinal Fluid analysis"
+    },
+    {
+        "id": "17",
+        "name": "Gastric Analysis"
+    },
+    {
+        "id": "18",
+        "name": "Semen analysis"
+    },
+    {
+        "id": "19",
+        "name": "Stool Examination"
+    },
+    {
+        "id": "20",
+        "name": "Thyroid Function Test"
     }
 ]