From 286f75cca0054a4290f8ad108a197ce37d6242eb Mon Sep 17 00:00:00 2001
From: Enrico Colasante <enrico@dhis2.org>
Date: Tue, 12 Nov 2024 13:34:21 +0100
Subject: [PATCH] fix: Fix enrollment queyr to get correct orgUnit
 [DHIS2-14829] (#19130)

---
 .../aggregates/query/EnrollmentQuery.java     |  2 +-
 .../TrackedEntityServiceTest.java             | 38 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/query/EnrollmentQuery.java b/dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/query/EnrollmentQuery.java
index dcd7aa5f84c7..9384db933f90 100644
--- a/dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/query/EnrollmentQuery.java
+++ b/dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/query/EnrollmentQuery.java
@@ -95,7 +95,7 @@ public static String getQuery() {
         + "join program p on en.programid = p.programid "
         + "join trackedentity te on en.trackedentityid = te.trackedentityid "
         + "join trackedentitytype tet on te.trackedentitytypeid = tet.trackedentitytypeid "
-        + "join organisationunit o on te.organisationunitid = o.organisationunitid "
+        + "join organisationunit o on en.organisationunitid = o.organisationunitid "
         + "where en.trackedentityid in (:ids) ";
   }
 
diff --git a/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/tracker/export/trackedentity/TrackedEntityServiceTest.java b/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/tracker/export/trackedentity/TrackedEntityServiceTest.java
index 3d19c64dd568..a4c81310c859 100644
--- a/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/tracker/export/trackedentity/TrackedEntityServiceTest.java
+++ b/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/tracker/export/trackedentity/TrackedEntityServiceTest.java
@@ -369,7 +369,7 @@ void setUp() {
     enrollmentA.setFollowup(true);
     manager.save(enrollmentA, false);
 
-    enrollmentB = createEnrollment(programB, trackedEntityA, orgUnitA);
+    enrollmentB = createEnrollment(programB, trackedEntityA, orgUnitB);
     manager.save(enrollmentB);
     trackedEntityA.getEnrollments().add(enrollmentB);
     manager.update(trackedEntityA);
@@ -642,6 +642,42 @@ void shouldReturnTrackedEntityIncludingAllAttributesEnrollmentsEventsRelationshi
                     .collect(Collectors.toSet())));
   }
 
+  @Test
+  void shouldReturnTrackedEntityIncludingAllEnrollments()
+      throws ForbiddenException, NotFoundException, BadRequestException {
+    TrackedEntityOperationParams operationParams =
+        TrackedEntityOperationParams.builder()
+            .organisationUnits(orgUnitA)
+            .orgUnitMode(SELECTED)
+            .trackedEntityType(trackedEntityTypeA)
+            .trackedEntityParams(TrackedEntityParams.TRUE)
+            .build();
+
+    final List<TrackedEntity> trackedEntities =
+        trackedEntityService.getTrackedEntities(operationParams);
+
+    assertContainsOnly(List.of(trackedEntityA.getUid()), uids(trackedEntities));
+    assertContainsOnly(
+        Set.of(enrollmentA.getUid(), enrollmentB.getUid()),
+        uids(trackedEntities.get(0).getEnrollments()));
+    assertEquals(
+        orgUnitA.getUid(),
+        trackedEntities.get(0).getEnrollments().stream()
+            .filter(e -> e.getUid().equals(enrollmentA.getUid()))
+            .findFirst()
+            .get()
+            .getOrganisationUnit()
+            .getUid());
+    assertEquals(
+        orgUnitB.getUid(),
+        trackedEntities.get(0).getEnrollments().stream()
+            .filter(e -> e.getUid().equals(enrollmentB.getUid()))
+            .findFirst()
+            .get()
+            .getOrganisationUnit()
+            .getUid());
+  }
+
   @Test
   void shouldReturnTrackedEntityIncludeSpecificProtectedProgram()
       throws ForbiddenException, NotFoundException, BadRequestException {