Fix tests

locationtech · Jul 31, 2017 · 70e4634 · 70e4634
1 parent 9bedc00
commit 70e4634
Showing 1 changed file with 43 additions and 21 deletions.
diff --git a/s3-test/src/test/scala/geotrellis/spark/io/s3/S3GeoTiffRDDSpec.scala b/s3-test/src/test/scala/geotrellis/spark/io/s3/S3GeoTiffRDDSpec.scala
@@ -35,7 +35,23 @@ class S3GeoTiffRDDSpec
   extends FunSpec
     with Matchers
     with RasterMatchers
-    with TestEnvironment {
+    with TestEnvironment
+    with BeforeAndAfterEach {
+
+  override def afterEach() {
+    try super.afterEach()
+    finally setDefaultWindowSize
+  }
+
+  implicit def toOption[T](t: T): Option[T] = Option(t)
+
+  val defaultWindowSize: Option[Int] = S3GeoTiffRDD.windowSize
+  def setDefaultWindowSize: Unit = setWindowSize(defaultWindowSize)
+  def setWindowSize(size: Option[Int]): Unit = {
+    val field = S3GeoTiffRDD.getClass.getDeclaredField("windowSize")
+    field.setAccessible(true)
+    field.set(S3GeoTiffRDD, size)
+  }
 
   describe("S3GeoTiffRDD") {
     implicit val mockClient = new MockS3Client()
@@ -48,9 +64,11 @@ class S3GeoTiffRDDSpec
       mockClient.putObject(bucket, key, geoTiffBytes)
 
       val source1 =
-        S3GeoTiffRDD.spatial(bucket, key, S3GeoTiffRDD.Options(maxTileSize = None, bySegments = false, partitionBytes = None, getS3Client = () => new MockS3Client))
-      val source2 =
-        S3GeoTiffRDD.spatial(bucket, key, S3GeoTiffRDD.Options(maxTileSize = Some(128), getS3Client = () => new MockS3Client))
+        S3GeoTiffRDD.spatial(bucket, key, S3GeoTiffRDD.Options(partitionBytes = None, getS3Client = () => new MockS3Client))
+      val source2 = {
+        setWindowSize(128)
+        S3GeoTiffRDD.spatial(bucket, key, S3GeoTiffRDD.Options(getS3Client = () => new MockS3Client))
+      }
 
       source1.count should be < (source2.count)
 
@@ -69,9 +87,11 @@ class S3GeoTiffRDDSpec
       mockClient.putObject(bucket, key, geoTiffBytes)
 
       val source1 =
-        S3GeoTiffRDD.spatialMultiband(bucket, key, S3GeoTiffRDD.Options(maxTileSize = None, bySegments = false, getS3Client = () => new MockS3Client))
-      val source2 =
-        S3GeoTiffRDD.spatialMultiband(bucket, key, S3GeoTiffRDD.Options(maxTileSize = Some(20), getS3Client = () => new MockS3Client))
+        S3GeoTiffRDD.spatialMultiband(bucket, key, S3GeoTiffRDD.Options(partitionBytes = None, getS3Client = () => new MockS3Client))
+      val source2 = {
+        setWindowSize(20)
+        S3GeoTiffRDD.spatialMultiband(bucket, key, S3GeoTiffRDD.Options(getS3Client = () => new MockS3Client))
+      }
 
       //source1.count should be < (source2.count)
       val (_, md) = source1.collectMetadata[SpatialKey](FloatingLayoutScheme(20, 40))
@@ -89,17 +109,18 @@ class S3GeoTiffRDDSpec
       mockClient.putObject(bucket, key, geoTiffBytes)
 
       val source1 = S3GeoTiffRDD.temporal(bucket, key, S3GeoTiffRDD.Options(
-        maxTileSize = None,
-        bySegments = false,
+        partitionBytes = None,
         timeTag = "ISO_TIME",
         timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
         getS3Client = () => new MockS3Client))
 
-      val source2 = S3GeoTiffRDD.temporal(bucket, key, S3GeoTiffRDD.Options(
-        maxTileSize = Some(128),
-        timeTag = "ISO_TIME",
-        timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
-        getS3Client = () => new MockS3Client))
+      val source2 = {
+        setWindowSize(128)
+        S3GeoTiffRDD.temporal(bucket, key, S3GeoTiffRDD.Options(
+          timeTag = "ISO_TIME",
+          timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
+          getS3Client = () => new MockS3Client))
+      }
 
       source1.count should be < (source2.count)
 
@@ -128,17 +149,18 @@ class S3GeoTiffRDDSpec
       val geoTiffBytes = multiband.toByteArray
       mockClient.putObject(bucket, key, geoTiffBytes)
       val source1 = S3GeoTiffRDD.temporalMultiband(bucket, key, S3GeoTiffRDD.Options(
-        maxTileSize = None,
-        bySegments = false,
+        partitionBytes = None,
         timeTag = "ISO_TIME",
         timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
         getS3Client = () => new MockS3Client))
 
-      val source2 = S3GeoTiffRDD.temporalMultiband(bucket, key, S3GeoTiffRDD.Options(
-        maxTileSize = Some(256),
-        timeTag = "ISO_TIME",
-        timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
-        getS3Client = () => new MockS3Client))
+      val source2 = {
+        setWindowSize(256)
+        S3GeoTiffRDD.temporalMultiband(bucket, key, S3GeoTiffRDD.Options(
+          timeTag = "ISO_TIME",
+          timeFormat = "yyyy-MM-dd'T'HH:mm:ss",
+          getS3Client = () => new MockS3Client))
+      }
 
       source1.count should be < (source2.count)