BI-2979: ObservationUnitDAO Optimization: Utilize searchQuery params instead of full program OU grab and filter

src/main/java/org/breedinginsight/brapi/v2/dao/BrAPIObservationUnitDAO.java

@@ -200,28 +200,25 @@ public List<BrAPIObservationUnit> getObservationUnitsById(Collection<String> obs
     }
 
     public List<BrAPIObservationUnit> getObservationUnitsForStudyDbId(@NotNull String studyDbId, Program program) throws ApiException {
-        // TODO: Optimize and change this to search/get on studyDbId once cache is removed for studies [BI-2979]
-        return getProgramObservationUnits(program.getId()).stream()
-                .filter(ou -> ou.getStudyDbId().equals(studyDbId))
-                .collect(Collectors.toList());
+        BrAPIObservationUnitSearchRequest observationUnitSearchRequest = buildSearchRequest(program, null);
+        observationUnitSearchRequest.setStudyDbIds(List.of(studyDbId));
+        return searchObservationUnitsAndProcess(observationUnitSearchRequest, program, true);
     }
 
     public List<BrAPIObservationUnit> getObservationUnitsForTrialDbIds(@NotNull UUID programId, List<String> trialDbIds) throws ApiException {
         if (trialDbIds.isEmpty()) {
             return Collections.emptyList();
         }
 
-        // TODO: Optimize and change this to search on trialDbIds [BI-2979]
-        return getProgramObservationUnits(programId).stream()
-                .filter(ou -> trialDbIds.contains(ou.getTrialDbId()))
-                .collect(Collectors.toList());
+        Program program = programDAO.get(programId).stream().findFirst().orElseThrow();
+
+        BrAPIObservationUnitSearchRequest observationUnitSearchRequest = buildSearchRequest(program, null);
+        observationUnitSearchRequest.setTrialDbIds(trialDbIds);
+        return searchObservationUnitsAndProcess(observationUnitSearchRequest, program, true);
     }
 
     public List<BrAPIObservationUnit> getObservationUnitsForTrialDbId(@NotNull UUID programId, @NotNull String trialDbId) throws ApiException {
-        return getProgramObservationUnits(programId).stream()
-                // TODO: Optimize and change this to search/get on trialDbId [BI-2979]
-                .filter(ou -> ou.getTrialDbId().equals(trialDbId))
-                .collect(Collectors.toList());
+        return getObservationUnitsForTrialDbIds(programId, List.of(trialDbId));
     }
 
     public List<BrAPIObservationUnit> getObservationUnitsForDataset(@NotNull String datasetId, @NotNull Program program) throws ApiException {

src/test/java/org/breedinginsight/brapi/v2/dao/BrAPIObservationUnitDAOUnitTest.java

@@ -0,0 +1,142 @@
+/*
+ * See the NOTICE file distributed with this work for additional information
+ * regarding copyright ownership.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.breedinginsight.brapi.v2.dao;
+
+import io.reactivex.functions.Function;
+import io.reactivex.functions.Function3;
+import lombok.SneakyThrows;
+import org.brapi.client.v2.BrAPIClient;
+import org.brapi.v2.model.core.BrAPIProgram;
+import org.brapi.v2.model.pheno.request.BrAPIObservationUnitSearchRequest;
+import org.breedinginsight.brapi.v2.services.BrAPIGermplasmService;
+import org.breedinginsight.brapps.importer.daos.ImportDAO;
+import org.breedinginsight.daos.ProgramDAO;
+import org.breedinginsight.model.Program;
+import org.breedinginsight.services.ProgramService;
+import org.breedinginsight.services.brapi.BrAPIEndpointProvider;
+import org.breedinginsight.utilities.BrAPIDAOUtil;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.util.List;
+import java.util.UUID;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.*;
+
+public class BrAPIObservationUnitDAOUnitTest {
+
+    private static final String BRAPI_PROGRAM_DB_ID = "brapi-program-1";
+
+    private BrAPIObservationUnitDAO observationUnitDAO;
+    private ProgramDAO programDAO;
+    private BrAPIDAOUtil brAPIDAOUtil;
+    private Program program;
+
+    @BeforeEach
+    @SneakyThrows
+    void setup() {
+        UUID programId = UUID.randomUUID();
+
+        BrAPIProgram brapiProgram = new BrAPIProgram().programDbId(BRAPI_PROGRAM_DB_ID);
+
+        program = new Program();
+        program.setId(programId);
+        program.setKey("TEST");
+        program.setBrapiProgram(brapiProgram);
+
+        programDAO = mock(ProgramDAO.class);
+        brAPIDAOUtil = mock(BrAPIDAOUtil.class);
+        BrAPIGermplasmService germplasmService = mock(BrAPIGermplasmService.class);
+
+        when(programDAO.get(programId)).thenReturn(List.of(program));
+        when(programDAO.getProgramBrAPI(program)).thenReturn(brapiProgram);
+        when(programDAO.getCoreClient(programId)).thenReturn(mock(BrAPIClient.class));
+        when(germplasmService.getGermplasm(programId)).thenReturn(List.of());
+
+        when(brAPIDAOUtil.search(any(Function.class), any(Function3.class), any(BrAPIObservationUnitSearchRequest.class))).thenReturn(List.of());
+
+        observationUnitDAO = new BrAPIObservationUnitDAO(programDAO, mock(ImportDAO.class),
+                brAPIDAOUtil,
+                new BrAPIEndpointProvider(), germplasmService, mock(ProgramService.class),
+                "test-reference-source",
+                1000
+        );
+    }
+
+    @Test
+    @SneakyThrows
+    void getObservationUnitsForStudyDbIdUsesBrAPISearch() {
+        String studyDbId = "study-1";
+
+        observationUnitDAO.getObservationUnitsForStudyDbId(studyDbId, program);
+
+        ArgumentCaptor<BrAPIObservationUnitSearchRequest> requestCaptor = ArgumentCaptor.forClass(BrAPIObservationUnitSearchRequest.class);
+
+        verify(brAPIDAOUtil).search(any(Function.class), any(Function3.class), requestCaptor.capture());
+
+        BrAPIObservationUnitSearchRequest request = requestCaptor.getValue();
+
+        assertEquals(List.of(BRAPI_PROGRAM_DB_ID), request.getProgramDbIds());
+        assertEquals(List.of(studyDbId), request.getStudyDbIds());
+    }
+
+    @Test
+    @SneakyThrows
+    void getObservationUnitsForTrialDbIdsUsesBrAPISearch() {
+        List<String> trialDbIds = List.of("trial-1", "trial-2");
+
+        observationUnitDAO.getObservationUnitsForTrialDbIds(program.getId(), trialDbIds);
+
+        ArgumentCaptor<BrAPIObservationUnitSearchRequest> requestCaptor = ArgumentCaptor.forClass(BrAPIObservationUnitSearchRequest.class);
+
+        verify(brAPIDAOUtil).search(any(Function.class), any(Function3.class), requestCaptor.capture());
+
+        BrAPIObservationUnitSearchRequest request = requestCaptor.getValue();
+
+        assertEquals(List.of(BRAPI_PROGRAM_DB_ID), request.getProgramDbIds());
+        assertEquals(trialDbIds, request.getTrialDbIds());
+    }
+
+    @Test
+    @SneakyThrows
+    void getObservationUnitsForTrialDbIdUsesBrAPISearch() {
+        String trialDbId = "trial-1";
+
+        observationUnitDAO.getObservationUnitsForTrialDbId(program.getId(), trialDbId);
+
+        ArgumentCaptor<BrAPIObservationUnitSearchRequest> requestCaptor = ArgumentCaptor.forClass(BrAPIObservationUnitSearchRequest.class);
+
+        verify(brAPIDAOUtil).search(any(Function.class), any(Function3.class), requestCaptor.capture());
+
+        BrAPIObservationUnitSearchRequest request = requestCaptor.getValue();
+
+        assertEquals(List.of(BRAPI_PROGRAM_DB_ID), request.getProgramDbIds());
+        assertEquals(List.of(trialDbId), request.getTrialDbIds());
+    }
+
+    @Test
+    @SneakyThrows
+    void getObservationUnitsForEmptyTrialDbIdsDoesNotSearch() {
+        assertTrue(observationUnitDAO.getObservationUnitsForTrialDbIds(program.getId(), List.of()).isEmpty());
+
+        verify(brAPIDAOUtil, never()).search(any(Function.class), any(Function3.class), any(BrAPIObservationUnitSearchRequest.class));
+    }
+}