diff --git a/db-scripts/src/main/resources/cgds.sql b/db-scripts/src/main/resources/cgds.sql
index c0b20e8b6cd..2971b3d1be6 100644
--- a/db-scripts/src/main/resources/cgds.sql
+++ b/db-scripts/src/main/resources/cgds.sql
@@ -755,5 +755,5 @@ CREATE TABLE `resource_study` (
 );
 
 -- THIS MUST BE KEPT IN SYNC WITH db.version PROPERTY IN pom.xml
-INSERT INTO info VALUES ('2.13.0', NULL);
+INSERT INTO info VALUES ('2.13.1', NULL);
 
diff --git a/db-scripts/src/main/resources/migration.sql b/db-scripts/src/main/resources/migration.sql
index dfec6a97381..78c0d25ea14 100644
--- a/db-scripts/src/main/resources/migration.sql
+++ b/db-scripts/src/main/resources/migration.sql
@@ -1021,3 +1021,10 @@ ALTER TABLE `mutation_event` CHANGE COLUMN `ONCOTATOR_PROTEIN_POS_START` `PROTEI
 ALTER TABLE `mutation_event` CHANGE COLUMN `ONCOTATOR_PROTEIN_POS_END` `PROTEIN_POS_END` int(11);
 UPDATE `info` SET `DB_SCHEMA_VERSION`="2.13.0";
 
+
+##version: 2.13.1
+ALTER TABLE `clinical_event_data` MODIFY COLUMN `VALUE` varchar(3000) NOT NULL;
+CREATE INDEX idx_clinical_event_key ON clinical_event_data (`KEY`);
+CREATE INDEX idx_clinical_event_value ON clinical_event_data (`VALUE`);
+CREATE INDEX idx_sample_stable_id ON sample (`STABLE_ID`);
+UPDATE `info` SET `DB_SCHEMA_VERSION`="2.13.1";
\ No newline at end of file
diff --git a/persistence/persistence-api/src/main/java/org/cbioportal/persistence/TreatmentRepository.java b/persistence/persistence-api/src/main/java/org/cbioportal/persistence/TreatmentRepository.java
index 83fb9e1fedb..bc2e0705b16 100644
--- a/persistence/persistence-api/src/main/java/org/cbioportal/persistence/TreatmentRepository.java
+++ b/persistence/persistence-api/src/main/java/org/cbioportal/persistence/TreatmentRepository.java
@@ -14,20 +14,18 @@ public interface TreatmentRepository {
     public Map<String, List<Treatment>> getTreatmentsByPatientId(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key);
 
     @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Map<String, List<ClinicalEventSample>> getSamplesByPatientId(List<String> sampleIds, List<String> studyIds);
+    public List<Treatment> getTreatments(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key);
 
     @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Map<String, List<ClinicalEventSample>> getShallowSamplesByPatientId(List<String> sampleIds, List<String> studyIds);
+    public Map<String, List<ClinicalEventSample>> getSamplesByPatientId(List<String> sampleIds, List<String> studyIds);
 
     @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Set<String> getAllUniqueTreatments(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key);
+    public Map<String, List<ClinicalEventSample>> getShallowSamplesByPatientId(List<String> sampleIds, List<String> studyIds);
 
     @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Integer getTreatmentCount(List<String> studies, String key);
+    public Boolean hasTreatmentData(List<String> studies, ClinicalEventKeyCode key);
 
     @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Integer getSampleCount(List<String> studies);
+    public Boolean hasSampleTimelineData(List<String> studies);
 
-    @Cacheable(cacheResolver = "generalRepositoryCacheResolver", condition = "@cacheEnabledConfig.getEnabled()")
-    public Boolean studyIdHasTreatments(String studyId, ClinicalEventKeyCode key);
 }
diff --git a/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMapper.java b/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMapper.java
index 6d748341dbd..44e4f856079 100644
--- a/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMapper.java
+++ b/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMapper.java
@@ -13,11 +13,7 @@ public interface TreatmentMapper {
     
     List<ClinicalEventSample> getAllShallowSamples(List<String> sampleIds, List<String> studyIds);
 
-    Set<String> getAllUniqueTreatments(List<String> sampleIds, List<String> studyIds, String key);
+    Boolean hasTreatmentData(List<String> sampleIds, List<String> studyIds, String key);
 
-    Integer getTreatmentCount(List<String> sampleIds, List<String> studyIds, String key);
-
-    Integer getSampleCount(List<String> sampleIds, List<String> studyIds);
-
-    Boolean studyIdHasTreatments(String studyId, String key);
+    Boolean hasSampleTimelineData(List<String> sampleIds, List<String> studyIds);
 }
diff --git a/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepository.java b/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepository.java
index efa6efb27a9..d44bdfc808a 100644
--- a/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepository.java
+++ b/persistence/persistence-mybatis/src/main/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepository.java
@@ -18,12 +18,17 @@ public class TreatmentMyBatisRepository implements TreatmentRepository {
     
     @Override
     public Map<String, List<Treatment>> getTreatmentsByPatientId(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key) {
+        return getTreatments(sampleIds, studyIds, key)
+            .stream()
+            .collect(groupingBy(Treatment::getPatientId));
+    }
+
+    @Override
+    public List<Treatment> getTreatments(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key) {
         return treatmentMapper.getAllTreatments(sampleIds, studyIds, key.getKey())
             .stream()
             .flatMap(treatment -> splitIfDelimited(treatment, key))
-            .collect(groupingBy(Treatment::getPatientId));
-        
-        
+            .collect(Collectors.toList());
     }
 
     private Stream<Treatment> splitIfDelimited(Treatment unsplitTreatment, ClinicalEventKeyCode key) {
@@ -61,31 +66,12 @@ public Map<String, List<ClinicalEventSample>> getShallowSamplesByPatientId(List<
     }
 
     @Override
-    public Set<String> getAllUniqueTreatments(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key) {
-        return treatmentMapper.getAllUniqueTreatments(sampleIds, studyIds, key.getKey())
-            .stream()
-            .flatMap(treatment -> {
-                if (key.isDelimited()) {
-                    return Arrays.stream(treatment.split(key.getDelimiter()));
-                } else {
-                    return Stream.of(treatment);
-                }
-            })
-            .collect(Collectors.toSet());
-    }
-
-    @Override
-    public Integer getTreatmentCount(List<String> studies, String key) {
-        return treatmentMapper.getTreatmentCount(null, studies, key);
-    }
-
-    @Override
-    public Integer getSampleCount(List<String> studies) {
-        return treatmentMapper.getSampleCount(null, studies);
+    public Boolean hasTreatmentData(List<String> studies, ClinicalEventKeyCode key) {
+        return treatmentMapper.hasTreatmentData(null, studies, key.getKey());
     }
 
     @Override
-    public Boolean studyIdHasTreatments(String studyId, ClinicalEventKeyCode key) {
-        return treatmentMapper.studyIdHasTreatments(studyId, key.getKey());
+    public Boolean hasSampleTimelineData(List<String> studies) {
+        return treatmentMapper.hasSampleTimelineData(null, studies);
     }
 }
diff --git a/persistence/persistence-mybatis/src/main/resources/org/cbioportal/persistence/mybatis/TreatmentMapper.xml b/persistence/persistence-mybatis/src/main/resources/org/cbioportal/persistence/mybatis/TreatmentMapper.xml
index a9976f44d3a..de633ab7f55 100644
--- a/persistence/persistence-mybatis/src/main/resources/org/cbioportal/persistence/mybatis/TreatmentMapper.xml
+++ b/persistence/persistence-mybatis/src/main/resources/org/cbioportal/persistence/mybatis/TreatmentMapper.xml
@@ -16,6 +16,7 @@
             INNER JOIN sample ON patient.INTERNAL_ID = sample.PATIENT_ID
             INNER JOIN cancer_study ON patient.CANCER_STUDY_ID = cancer_study.CANCER_STUDY_ID
         <include refid="where"/>
+        AND clinical_event.EVENT_TYPE = 'TREATMENT'
         AND clinical_event_data.KEY = #{key}
     </select>
 
@@ -48,9 +49,9 @@
         <include refid="where"/>
     </select>
 
-    <select id="getAllUniqueTreatments" resultType="java.lang.String">
-        SELECT 
-            DISTINCT clinical_event_data.VALUE
+    <select id="hasTreatmentData" resultType="java.lang.Boolean">
+        SELECT EXISTS(SELECT
+            *
         FROM
             clinical_event
             INNER JOIN clinical_event_data ON clinical_event.CLINICAL_EVENT_ID = clinical_event_data.CLINICAL_EVENT_ID
@@ -58,25 +59,14 @@
             INNER JOIN sample ON patient.INTERNAL_ID = sample.PATIENT_ID
             INNER JOIN cancer_study ON patient.CANCER_STUDY_ID = cancer_study.CANCER_STUDY_ID
         <include refid="where"/>
-            AND clinical_event_data.KEY = #{key}
+            AND clinical_event.EVENT_TYPE = 'TREATMENT'
+            AND clinical_event_data.KEY = #{key} LIMIT 1
+        )
     </select>
 
-    <select id="getTreatmentCount" resultType="java.lang.Integer">
-        SELECT
-            COUNT(clinical_event.PATIENT_ID)
-        FROM
-            clinical_event
-            INNER JOIN clinical_event_data ON clinical_event.CLINICAL_EVENT_ID = clinical_event_data.CLINICAL_EVENT_ID
-            INNER JOIN patient ON clinical_event.PATIENT_ID = patient.INTERNAL_ID
-            INNER JOIN sample ON patient.INTERNAL_ID = sample.PATIENT_ID
-            INNER JOIN cancer_study ON patient.CANCER_STUDY_ID = cancer_study.CANCER_STUDY_ID
-        <include refid="where"/>
-            AND clinical_event_data.KEY = #{key}
-    </select>
-
-    <select id="getSampleCount" resultType="java.lang.Integer">
-        SELECT
-            COUNT(sample.STABLE_ID)
+    <select id="hasSampleTimelineData" resultType="java.lang.Boolean">
+        SELECT EXISTS(SELECT
+            *
         FROM
             clinical_event
             INNER JOIN clinical_event_data ON clinical_event.CLINICAL_EVENT_ID = clinical_event_data.CLINICAL_EVENT_ID
@@ -85,20 +75,7 @@
             INNER JOIN cancer_study ON patient.CANCER_STUDY_ID = cancer_study.CANCER_STUDY_ID
         <include refid="where"/>
             AND clinical_event_data.KEY = 'SAMPLE_ID'
-            AND (clinical_event.EVENT_TYPE LIKE 'Sample Acquisition' OR clinical_event.EVENT_TYPE LIKE 'SPECIMEN')
-    </select>
-    
-    <select id="studyIdHasTreatments" resultType="java.lang.Boolean">
-        SELECT
-            count(cancer_study.CANCER_STUDY_IDENTIFIER) > 0
-        FROM
-            clinical_event
-            INNER JOIN clinical_event_data ON clinical_event.CLINICAL_EVENT_ID = clinical_event_data.CLINICAL_EVENT_ID
-            INNER JOIN patient ON clinical_event.PATIENT_ID = patient.INTERNAL_ID
-            INNER JOIN cancer_study ON patient.CANCER_STUDY_ID = cancer_study.CANCER_STUDY_ID
-        WHERE
-            clinical_event_data.KEY = #{key}
-            AND cancer_study.CANCER_STUDY_IDENTIFIER = #{studyId}
+            AND (clinical_event.EVENT_TYPE LIKE 'Sample Acquisition' OR clinical_event.EVENT_TYPE LIKE 'SPECIMEN') LIMIT 1)
     </select>
 
     <sql id="where">
diff --git a/persistence/persistence-mybatis/src/test/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepositoryTest.java b/persistence/persistence-mybatis/src/test/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepositoryTest.java
index 0ebf28a7405..e61e23cb807 100644
--- a/persistence/persistence-mybatis/src/test/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepositoryTest.java
+++ b/persistence/persistence-mybatis/src/test/java/org/cbioportal/persistence/mybatis/TreatmentMyBatisRepositoryTest.java
@@ -134,32 +134,20 @@ public void getShallowSamplesByPatientId() {
     }
 
     @Test
-    public void getAllUniqueTreatments() {
-        HashSet<String> expected = new HashSet<>(Collections.singletonList("Madeupanib"));
-        
-        Set<String> actual = treatmentRepository.getAllUniqueTreatments(
-            Collections.singletonList("TCGA-A1-A0SD-01"),
-            Collections.singletonList("study_tcga_pub"),
-            ClinicalEventKeyCode.Agent
-        );
-        
-        Assert.assertEquals(actual, expected);
-    }
+    public void hasTreatmentData() {
+
+        Assert.assertEquals(true, treatmentRepository.hasTreatmentData(Collections.singletonList("study_tcga_pub"), ClinicalEventKeyCode.Agent));
+
+        Assert.assertEquals(false, treatmentRepository.hasTreatmentData(Collections.singletonList("acc_tcga"), ClinicalEventKeyCode.Agent));
 
-    @Test
-    public void getTreatmentCount() {
-        Integer expected = 1;
-        Integer actual = treatmentRepository.getTreatmentCount(Collections.singletonList("study_tcga_pub"), ClinicalEventKeyCode.Agent.getKey());
-        
-        Assert.assertEquals(actual, expected);
     }
 
     @Test
-    public void getSampleCount() {
-        Integer expected = 2;
-        Integer actual = treatmentRepository.getSampleCount(Collections.singletonList("study_tcga_pub"));
-        
-        Assert.assertEquals(actual, expected);
+    public void hasSampleTimelineData() {
+
+        Assert.assertEquals(true, treatmentRepository.hasSampleTimelineData(Collections.singletonList("study_tcga_pub")));
+
+        Assert.assertEquals(false, treatmentRepository.hasSampleTimelineData(Collections.singletonList("acc_tcga")));
     }
     
 }
\ No newline at end of file
diff --git a/pom.xml b/pom.xml
index 56c2b78fed7..ce9685c8f5f 100644
--- a/pom.xml
+++ b/pom.xml
@@ -314,7 +314,7 @@
     <tomcat.session.timeout>720</tomcat.session.timeout>
 
     <!-- THIS SHOULD BE KEPT IN SYNC TO VERSION IN CGDS.SQL -->
-    <db.version>2.13.0</db.version>
+    <db.version>2.13.1</db.version>
   </properties>
 
   <modules>
diff --git a/service/src/main/java/org/cbioportal/service/impl/TreatmentServiceImpl.java b/service/src/main/java/org/cbioportal/service/impl/TreatmentServiceImpl.java
index 439f94ce20c..75784a01c15 100644
--- a/service/src/main/java/org/cbioportal/service/impl/TreatmentServiceImpl.java
+++ b/service/src/main/java/org/cbioportal/service/impl/TreatmentServiceImpl.java
@@ -12,8 +12,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import static java.util.stream.Collectors.toList;
-
 @Service
 public class TreatmentServiceImpl implements TreatmentService {
     @Autowired
@@ -25,7 +23,7 @@ private Pair<List<String>, List<String>> filterIds(List<String> sampleIds, List<
         }
         Set<String> studiesWithTreatments = studyIds.stream()
             .distinct()
-            .filter(studyId -> treatmentRepository.studyIdHasTreatments(studyId, key))
+            .filter(studyId -> treatmentRepository.hasTreatmentData(Collections.singletonList(studyId), key))
             .collect(Collectors.toSet());
         
         ArrayList<String> filteredSampleIds = new ArrayList<>();
@@ -40,16 +38,16 @@ private Pair<List<String>, List<String>> filterIds(List<String> sampleIds, List<
         }
         return new ImmutablePair<>(filteredSampleIds, filteredStudyIds);
     }
-    
+
     @Override
     public List<SampleTreatmentRow> getAllSampleTreatmentRows(List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key) {
         Pair<List<String>, List<String>> filteredIds = filterIds(sampleIds, studyIds, key);
         sampleIds = filteredIds.getLeft();
         studyIds = filteredIds.getRight();
 
-        Map<String, List<ClinicalEventSample>> samplesByPatient = 
+        Map<String, List<ClinicalEventSample>> samplesByPatient =
             treatmentRepository.getSamplesByPatientId(sampleIds, studyIds);
-        Map<String, List<Treatment>> treatmentsByPatient = 
+        Map<String, List<Treatment>> treatmentsByPatient =
             treatmentRepository.getTreatmentsByPatientId(sampleIds, studyIds, key);
 
         Stream<SampleTreatmentRow> rows = samplesByPatient.keySet().stream()
@@ -150,7 +148,6 @@ Stream<SampleTreatmentRow> toRows() {
             );
         }
     }
-
     @Override
     public List<PatientTreatmentRow> getAllPatientTreatmentRows(
         List<String> sampleIds, List<String> studyIds, ClinicalEventKeyCode key
@@ -158,63 +155,44 @@ public List<PatientTreatmentRow> getAllPatientTreatmentRows(
         Pair<List<String>, List<String>> filteredIds = filterIds(sampleIds, studyIds, key);
         sampleIds = filteredIds.getLeft();
         studyIds = filteredIds.getRight();
-        
-        Map<String, List<Treatment>> treatmentsByPatient = 
-            treatmentRepository.getTreatmentsByPatientId(sampleIds, studyIds, key);
-        Map<String, List<ClinicalEventSample>> samplesByPatient = treatmentRepository
-            .getShallowSamplesByPatientId(sampleIds, studyIds)
-            .entrySet()
-            .stream()
-            .filter(e -> treatmentsByPatient.containsKey(e.getKey()))
-            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-        Set<String> treatments = treatmentRepository.getAllUniqueTreatments(sampleIds, studyIds, key);
-
-        return treatments.stream()
-            .map(t -> createPatientTreatmentRowForTreatment(t, treatmentsByPatient, samplesByPatient))
-            .collect(Collectors.toList());
-    }
 
-    private PatientTreatmentRow createPatientTreatmentRowForTreatment(
-        String treatment,
-        Map<String, List<Treatment>> treatmentsByPatient,
-        Map<String, List<ClinicalEventSample>> samplesByPatient
-    ) {
-        // find all the patients that have received this treatment
-        List<Map.Entry<String, List<Treatment>>> matchingPatients = matchingPatients(treatment, treatmentsByPatient);
+        Map<String, List<ClinicalEventSample>> samplesByPatient = treatmentRepository
+            .getShallowSamplesByPatientId(sampleIds, studyIds);
 
-        // from those patients, extract the unique samples
-        Set<ClinicalEventSample> samples = matchingPatients
+        Map<String, List<Treatment>> treatmentSet = treatmentRepository.getTreatments(sampleIds, studyIds, key)
             .stream()
-            .map(Map.Entry::getKey)
-            .flatMap(patient -> samplesByPatient.getOrDefault(patient, new ArrayList<>()).stream())
-            .collect(Collectors.toSet());
-
-
-        return new PatientTreatmentRow(treatment, matchingPatients.size(), samples);
-    }
+            .collect(Collectors.groupingBy(Treatment::getTreatment));
 
-    private List<Map.Entry<String, List<Treatment>>> matchingPatients(
-        String treatment,
-        Map<String, List<Treatment>> treatmentsByPatient
-    ) {
-        return treatmentsByPatient.entrySet().stream()
-            .filter(p -> p.getValue().stream().anyMatch(t -> t.getTreatment().equals(treatment)))
-            .collect(toList());
+        /*
+            This logic transforms treatmentSet to list of PatientTreatmentRow. transformation steps:
+            - key in treatmentSet is going to be treatment
+            - get all unique patient ids -> this is going to give count
+            - get all clinicalEventSamples using above unique patient ids
+         */
+        return treatmentSet.entrySet()
+            .stream()
+            .map(entry -> {
+                String treatment = entry.getKey();
+                Set<String> patientIds = entry.getValue().stream().map(Treatment::getPatientId).collect(Collectors.toSet());
+                Set<ClinicalEventSample> clinicalEventSamples = patientIds
+                    .stream()
+                    .flatMap(patientId -> samplesByPatient.getOrDefault(patientId, new ArrayList<>()).stream())
+                    .collect(Collectors.toSet());
+                return new PatientTreatmentRow(treatment, patientIds.size(), clinicalEventSamples);
+            })
+            .collect(Collectors.toList());
     }
 
     @Override
     public Boolean containsTreatmentData(List<String> studies, ClinicalEventKeyCode key) {
-        return treatmentRepository.getTreatmentCount(studies, key.getKey()) > 0;
+        return treatmentRepository.hasTreatmentData(studies, key);
     }
 
     @Override
     public Boolean containsSampleTreatmentData(List<String> studyIds, ClinicalEventKeyCode key) {
         studyIds = studyIds.stream()
-            .filter(studyId -> treatmentRepository.studyIdHasTreatments(studyId, key))
+            .filter(studyId -> treatmentRepository.hasTreatmentData(Collections.singletonList(studyId), key))
             .collect(Collectors.toList());
-        Integer sampleCount = treatmentRepository.getSampleCount(studyIds);
-        Integer treatmentCount = treatmentRepository.getTreatmentCount(studyIds, key.getKey());
-        
-        return sampleCount * treatmentCount > 0;
+        return studyIds.size() > 0 && treatmentRepository.hasSampleTimelineData(studyIds);
     }
 }
diff --git a/service/src/test/java/org/cbioportal/service/impl/TreatmentServiceImplTest.java b/service/src/test/java/org/cbioportal/service/impl/TreatmentServiceImplTest.java
index 9b489f16690..8a86e465d3f 100644
--- a/service/src/test/java/org/cbioportal/service/impl/TreatmentServiceImplTest.java
+++ b/service/src/test/java/org/cbioportal/service/impl/TreatmentServiceImplTest.java
@@ -26,14 +26,11 @@ public class TreatmentServiceImplTest {
     private TreatmentRepository treatmentRepository;
     
     @Test
-    public void getAllSampleTreatmentsSingleRow() {
-        mockTreatmentsByPatient(
-            makeTreatment("madeupanib", "P0", 0, 10)
-        );
+    public void getAllPatientTreatmentRows() {
         mockSamplesByPatient(
             makeSample("S0", "P0", 5)
         );
-        mockAllTreatments("madeupanib");
+        mockTreatments(makeTreatment("madeupanib", "P0", 0, 10));
 
         PatientTreatmentRow rowA = makePatientRow("madeupanib", 1, Collections.singletonList("S0"), Collections.singletonList("P0"));
         List<PatientTreatmentRow> expected = Collections.singletonList(rowA);
@@ -43,15 +40,14 @@ public void getAllSampleTreatmentsSingleRow() {
     }
 
     @Test
-    public void getAllSampleTreatmentsOneSampleTwoTreatmentsOnePatient() {
-        mockTreatmentsByPatient(
+    public void getAllPatientTreatmentRowsOneSampleTwoTreatmentsOnePatient() {
+        mockTreatments(
             makeTreatment("madeupanib", "P0", 0, 10),
             makeTreatment("fakedrugazol", "P0", 0, 10)
         );
         mockSamplesByPatient(
             makeSample("S0", "P0", 5)
         );
-        mockAllTreatments("madeupanib", "fakedrugazol");
 
 
         PatientTreatmentRow rowA = makePatientRow("fakedrugazol", 1, Collections.singletonList("S0"), Collections.singletonList("P0"));
@@ -63,15 +59,14 @@ public void getAllSampleTreatmentsOneSampleTwoTreatmentsOnePatient() {
     }
 
     @Test
-    public void getAllSampleTreatmentsTwoSamplesOnePatientOneTreatment() {
-        mockTreatmentsByPatient(
+    public void getAllPatientTreatmentRowsTwoSamplesOnePatientOneTreatment() {
+        mockTreatments(
             makeTreatment("fakedrugazol", "P0", 0, 10)
         );
         mockSamplesByPatient(
             makeSample("S0", "P0", 5),
             makeSample("S1", "P0", 10)
         );
-        mockAllTreatments("fakedrugazol");
 
         // even though there are two samples, you expect a count of 1,
         // because this is from the patient level, and both samples are for the same patient
@@ -83,8 +78,8 @@ public void getAllSampleTreatmentsTwoSamplesOnePatientOneTreatment() {
     }
 
     @Test
-    public void getAllSampleTreatmentsTwoSamplesTwoPatientsTwoTreatments() {
-        mockTreatmentsByPatient(
+    public void getAllPatientTreatmentRowsTwoSamplesTwoPatientsTwoTreatments() {
+        mockTreatments(
             makeTreatment("fakedrugazol", "P0", 0, 10),
             makeTreatment("fakedrugazol", "P1", 0, 10)
         );
@@ -92,7 +87,6 @@ public void getAllSampleTreatmentsTwoSamplesTwoPatientsTwoTreatments() {
             makeSample("S0", "P0", 5),
             makeSample("S1", "P1", 10)
         );
-        mockAllTreatments("fakedrugazol");
 
         // now there are two patients, so you expect a count of two
         PatientTreatmentRow rowA = makePatientRow("fakedrugazol", 2, Arrays.asList("S0", "S1"), Arrays.asList("P0", "P1"));
@@ -103,8 +97,8 @@ public void getAllSampleTreatmentsTwoSamplesTwoPatientsTwoTreatments() {
     }
 
     @Test
-    public void getAllSampleTreatmentsTwoSamplesTwoPatientsTwoDifferentTreatments() {
-        mockTreatmentsByPatient(
+    public void getAllPatientTreatmentRowsTwoSamplesTwoPatientsTwoDifferentTreatments() {
+        mockTreatments(
             makeTreatment("fakedrugazol", "P0", 0, 10),
             makeTreatment("madeupanib", "P1", 0, 10)
         );
@@ -112,7 +106,6 @@ public void getAllSampleTreatmentsTwoSamplesTwoPatientsTwoDifferentTreatments()
             makeSample("S0", "P0", 5),
             makeSample("S1", "P1", 10)
         );
-        mockAllTreatments("fakedrugazol", "madeupanib");
 
         PatientTreatmentRow rowA = makePatientRow("fakedrugazol", 1, Collections.singletonList("S0"), Collections.singletonList("P0"));
         PatientTreatmentRow rowB = makePatientRow("madeupanib", 1, Collections.singletonList("S1"), Collections.singletonList("P1"));
@@ -252,6 +245,11 @@ private void mockTreatmentsByPatient(Treatment... treatments) {
             .thenReturn(treatmentsByPatient);
     }
 
+    private void mockTreatments(Treatment... treatments) {
+        Mockito.when(treatmentRepository.getTreatments(Mockito.any(), Mockito.any(), Mockito.any()))
+            .thenReturn(Arrays.stream(treatments).collect(Collectors.toList()));
+    }
+
     private void mockSamplesByPatient(ClinicalEventSample... samples) {
         Map<String, List<ClinicalEventSample>> samplesByPatient = Arrays.stream(samples)
             .collect(Collectors.groupingBy(ClinicalEventSample::getPatientId));
@@ -261,12 +259,6 @@ private void mockSamplesByPatient(ClinicalEventSample... samples) {
             .thenReturn(samplesByPatient);
     }
 
-    private void mockAllTreatments(String... treatments) {
-        Set<String> allTreatments = new HashSet<>(Arrays.asList(treatments));
-        Mockito.when(treatmentRepository.getAllUniqueTreatments(Mockito.any(), Mockito.any(), Mockito.any()))
-            .thenReturn(allTreatments);
-    }
-
     private Treatment makeTreatment(String treatment, String patientId, Integer start, Integer stop) {
         Treatment t = new Treatment();
         t.setTreatment(treatment);