From a31987475c8eec7b05de7aedecf514ae7884cc83 Mon Sep 17 00:00:00 2001
From: nazarfil <nfilipchuk@bluesquarehub.com>
Date: Thu, 10 Oct 2024 15:14:58 +0200
Subject: [PATCH] feat: support both urgUnits and orgunits from IASO

---
 openhexa/toolbox/iaso/iaso.py            |  7 ++++++-
 tests/iaso/fixtures/iaso_api_fixtures.py | 16 ++++++++++++++++
 tests/iaso/test_iaso.py                  | 13 +++++++++++--
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/openhexa/toolbox/iaso/iaso.py b/openhexa/toolbox/iaso/iaso.py
index 3bff007..99b8b10 100644
--- a/openhexa/toolbox/iaso/iaso.py
+++ b/openhexa/toolbox/iaso/iaso.py
@@ -58,7 +58,12 @@ def get_org_units(self, page: int = 1, limit: int = 10, **kwargs) -> dict:
         params = kwargs
         params.update({"page": page, "limit": limit})
         response = self.api_client.get("/api/orgunits", params=params)
-        return response.json().get("orgunits")
+        json_response = response.json()
+        if "orgunits" in json_response:
+            return json_response.get("orgunits")
+        elif "orgUnits" in json_response:
+            return json_response.get("orgUnits")
+        return []
 
     def get_form_instances(
         self,
diff --git a/tests/iaso/fixtures/iaso_api_fixtures.py b/tests/iaso/fixtures/iaso_api_fixtures.py
index b246c73..8230a50 100644
--- a/tests/iaso/fixtures/iaso_api_fixtures.py
+++ b/tests/iaso/fixtures/iaso_api_fixtures.py
@@ -43,6 +43,22 @@
 }
 
 iaso_mocked_orgunits = {
+    "orgUnits": [
+        {
+            "name": "ACEH",
+            "id": 1978297,
+            "parent_id": 1978331,
+            "org_unit_type_id": 781,
+            "org_unit_type_name": "Province",
+            "validation_status": "VALID",
+            "created_at": 1712825023.085615,
+            "updated_at": 1712828860.665764,
+        }
+    ]
+}
+
+
+iaso_mocked_orgunits_with_params = {
     "orgunits": [
         {
             "name": "ACEH",
diff --git a/tests/iaso/test_iaso.py b/tests/iaso/test_iaso.py
index 931b1a4..d37b41c 100644
--- a/tests/iaso/test_iaso.py
+++ b/tests/iaso/test_iaso.py
@@ -9,7 +9,7 @@
     iaso_mocked_orgunits,
     iaso_mocked_refreshed_auth_token,
     iaso_mocked_projects,
-    iaso_mocked_instances,
+    iaso_mocked_instances, iaso_mocked_orgunits_with_params,
 )
 
 
@@ -50,7 +50,16 @@ def test_get_org_units(self, mock_responses):
         iaso = IASO("https://iaso-staging.bluesquare.org", "username", "password")
         r = iaso.get_org_units()
         assert len(r) > 0
-
+    def test_get_org_units_with_params(self, mock_responses):
+        mock_responses.add(
+            responses.POST, "https://iaso-staging.bluesquare.org/api/token/", json=iaso_mocked_auth_token, status=200
+        )
+        mock_responses.add(
+            responses.GET, "https://iaso-staging.bluesquare.org/api/orgunits/", json=iaso_mocked_orgunits_with_params, status=200
+        )
+        iaso = IASO("https://iaso-staging.bluesquare.org", "username", "password")
+        r = iaso.get_org_units()
+        assert len(r) > 0
     def test_get_forms(self, mock_responses):
         mock_responses.add(
             responses.POST, "https://iaso-staging.bluesquare.org/api/token/", json=iaso_mocked_auth_token, status=200