CNAM-410: Add new tests.

Author: Youcef Sebiat

src/main/scala/fr/polytechnique/cmap/cnam/etl/transformers/exposures/LimitedExposurePeriodAdder.scala

@@ -15,26 +15,26 @@ private class LimitedExposurePeriodAdder(data: DataFrame) extends ExposurePeriod
   private val orderedWindow = window.orderBy(col(Start))
 
   /** *
-    * This strategy works as the following:
-    * 1. Each DrugPurchase will have a corresponding Exposure.
-    * 2. Each Exposure has one or multiple DrugPurchases.
-    * 3. An Exposure is defined recursively as follows:
-    *   A. The first DrugPurchase defines a new Exposure.
-    *   B. If there is a DrugPurchase within the defined window of the first DrugPurchase, then expand the current
-    * Exposure with the DrugPurchase.
-    *   C. Else, close and set the end of the Exposure as the reach of the latest Drug Purchase and create a new
-    *   Exposure with the next DrugPurchase as the new Exposure.
-    * This strategy is suited for short term effects.
-    * !!! WARNING: THIS ONLY RETURNS EXPOSURES.
-    *
-    * @param minPurchases : Not used.
-    * @param  startDelay : period to be added to the start of each DrugPurchase.
-    * @param purchasesWindow : Not used.
-    * @param endThresholdGc : the period that defines the reach for Grand Condtionnement.
-    * @param endThresholdNgc : the period that defines the reach for Non Grand Condtionnement.
-    * @param endDelay : period added to the end of an exposure.
-    * @return: A DataFrame of Exposures.
-    */
+   * This strategy works as the following:
+   * 1. Each DrugPurchase will have a corresponding Exposure.
+   * 2. Each Exposure has one or multiple DrugPurchases.
+   * 3. An Exposure is defined recursively as follows:
+   *   A. The first DrugPurchase defines a new Exposure.
+   *   B. If there is a DrugPurchase within the defined window of the first DrugPurchase, then expand the current
+   * Exposure with the DrugPurchase.
+   *   C. Else, close and set the end of the Exposure as the reach of the latest Drug Purchase and create a new
+   * Exposure with the next DrugPurchase as the new Exposure.
+   * This strategy is suited for short term effects.
+   * !!! WARNING: THIS ONLY RETURNS EXPOSURES.
+   *
+   * @param minPurchases : Not used.
+   * @param startDelay : period to be added to delay the start of each DrugPurchase.
+   * @param purchasesWindow : Not used.
+   * @param endThresholdGc : the period that defines the reach for Grand Conditionnement.
+   * @param endThresholdNgc : the period that defines the reach for Non Grand Conditionnement.
+   * @param endDelay : period added to the end of an exposure.
+   * @return: A DataFrame of Exposures.
+   */
   def withStartEnd(
     minPurchases: Int = 2,
     startDelay: Period = 5.days,
@@ -48,7 +48,12 @@ private class LimitedExposurePeriodAdder(data: DataFrame) extends ExposurePeriod
 
     val delayedDrugPurchases = delayStart(data, startDelay)
 
-    val firstLastPurchase = getFirstAndLastPurchase(delayedDrugPurchases, endThresholdGc.get, endThresholdNgc.get, endDelay.get)
+    val firstLastPurchase = getFirstAndLastPurchase(
+      delayedDrugPurchases,
+      endThresholdGc.get,
+      endThresholdNgc.get,
+      endDelay.get
+    )
 
     toExposure(firstLastPurchase).select(outputColumns: _*)
   }

src/test/scala/fr/polytechnique/cmap/cnam/etl/transformers/exposures/LimitedExposurePeriodAdderSuite.scala

@@ -55,26 +55,26 @@ class LimitedExposurePeriodAdderSuite extends SharedContext {
 
     // Given
     val input = Seq(
-      ("A", "P", makeTS(2008, 1, 1), 1.0),
+      ("A", "P", makeTS(2008, 1, 1), 2.0),
       ("A", "P", makeTS(2008, 2, 1), 1.0),
       ("A", "P", makeTS(2008, 5, 1), 1.0),
       ("A", "P", makeTS(2008, 6, 1), 1.0),
       ("A", "P", makeTS(2008, 7, 1), 1.0),
       ("A", "P", makeTS(2009, 1, 1), 1.0),
       ("A", "P", makeTS(2009, 7, 1), 1.0),
       ("A", "P", makeTS(2009, 8, 1), 1.0),
-      ("A", "S", makeTS(2008, 2, 1), 1.0)
+      ("A", "S", makeTS(2008, 2, 1), 2.0)
     ).toDF(PatientID, Value, Start, Weight)
 
     val expected = Seq(
-      ("A", "P", makeTS(2008, 1, 1), 1.0, "first", makeTS(2008, 2, 11)),
+      ("A", "P", makeTS(2008, 1, 1), 2.0, "first", makeTS(2008, 4, 11)),
       ("A", "P", makeTS(2008, 2, 1), 1.0, "last", makeTS(2008, 3, 11)),
       ("A", "P", makeTS(2008, 5, 1), 1.0, "first", makeTS(2008, 6, 11)),
       ("A", "P", makeTS(2008, 7, 1), 1.0, "last", makeTS(2008, 8, 11)),
       ("A", "P", makeTS(2009, 1, 1), 1.0, "first", makeTS(2009, 2, 11)),
       ("A", "P", makeTS(2009, 7, 1), 1.0, "first", makeTS(2009, 8, 11)),
       ("A", "P", makeTS(2009, 8, 1), 1.0, "last", makeTS(2009, 9, 11)),
-      ("A", "S", makeTS(2008, 2, 1), 1.0, "first", makeTS(2008, 3, 11))
+      ("A", "S", makeTS(2008, 2, 1), 2.0, "first", makeTS(2008, 5, 11))
     ).toDF(PatientID, Value, Start, Weight, "Status", "purchaseReach")
 
     val result = mockInstance.getFirstAndLastPurchase(input, 1.months, 3.month, 10.days)