improve testing framework

.scalafmt.conf

@@ -39,6 +39,7 @@ rewrite {
 
 project.excludeFilters = [
   "MultiJoin.scala",
+  "SMBMultiJoin.scala",
   "TupleCoders.scala",
   "scalafix/input/*",
   "scalafix/output/*"

build.sbt

@@ -1285,7 +1285,7 @@ lazy val `scio-smb`: Project = project
   .in(file("scio-smb"))
   .dependsOn(
     `scio-core`,
-    `scio-test` % Test
+    `scio-test` % "test->test"
   )
   .settings(commonSettings)
   .settings(publishSettings)

integration/src/test/scala/com/spotify/scio/neo4j/Neo4jIOIT.scala

@@ -90,7 +90,6 @@ class Neo4jIOIT extends PipelineSpec with Eventually with ForAllTestContainer {
 
   val americanPresident: Movie = Movie("American President", 1995)
 
-  val options: PipelineOptions = PipelineOptionsFactory.create()
   lazy val neo4jOptions: Neo4jOptions = Neo4jOptions(
     Neo4jConnectionOptions(container.boltUrl, container.username, container.password)
   )
@@ -107,7 +106,7 @@ class Neo4jIOIT extends PipelineSpec with Eventually with ForAllTestContainer {
       Role(martin, Movie("American President", 1995), "A.J. MacInerney")
     )
 
-    runWithRealContext(options) { sc =>
+    runWithRealContext() { sc =>
       val resultQueryRoles = sc.neo4jCypher[Role](neo4jOptions, queryRoles)
       resultQueryRoles should containInAnyOrder(expectedRoles)
     }
@@ -132,7 +131,7 @@ class Neo4jIOIT extends PipelineSpec with Eventually with ForAllTestContainer {
       Role(morgan, Movie("The Shawshank Redemption", 1994), "Ellis Boyd 'Red' Redding")
     )
 
-    runWithRealContext(options) { sc =>
+    runWithRealContext() { sc =>
       val resultQueryMovieYear = sc
         .parallelize(queryParams)
         .neo4jCypher[Role](neo4jOptions, queryRoles)
@@ -154,7 +153,7 @@ class Neo4jIOIT extends PipelineSpec with Eventually with ForAllTestContainer {
                           |  (c:Country {name: origin.country})
                           |CREATE (m)-[:ORIGIN]->(c)
                           |""".stripMargin
-    runWithRealContext(options) { sc =>
+    runWithRealContext() { sc =>
       sc
         .parallelize(movieOrigins)
         .saveAsNeo4j(neo4jOptions, insertOrigins)

scio-core/src/main/scala/com/spotify/scio/ScioContext.scala

@@ -308,10 +308,12 @@ object ScioContext {
     new ScioContext(options, artifacts)
 
   /** Create a new [[ScioContext]] instance for testing. */
-  def forTest(): ScioContext = {
-    val opts = PipelineOptionsFactory
-      .fromArgs("--appName=" + TestUtil.newTestId())
-      .as(classOf[PipelineOptions])
+  def forTest(): ScioContext = forTest(TestUtil.newTestId())
+
+  /** Create a new [[ScioContext]] instance for testing. */
+  def forTest(testId: String): ScioContext = {
+    val opts = PipelineOptionsFactory.create()
+    opts.as(classOf[ApplicationNameOptions]).setAppName(testId)
     new ScioContext(opts, List[String]())
   }
 

scio-core/src/main/scala/com/spotify/scio/io/ScioIO.scala

@@ -18,8 +18,10 @@
 package com.spotify.scio.io
 
 import com.spotify.scio.ScioContext
+import com.spotify.scio.coders.Coder
 import com.spotify.scio.testing.TestDataManager
 import com.spotify.scio.values.SCollection
+
 import scala.annotation.unused
 
 sealed trait TapT[A] extends Serializable {
@@ -159,6 +161,12 @@ trait TestIO[T] extends ScioIO[T] {
     throw new UnsupportedOperationException(s"$this is for testing purpose only")
 }
 
+trait KeyedIO[T] { self: ScioIO[T] =>
+  type KeyT
+  def keyBy: T => KeyT
+  def keyCoder: Coder[KeyT]
+}
+
 /**
  * Special version of [[ScioIO]] for use with [[ScioContext.customInput]] and
  * [[SCollection.saveAsCustomOutput]].

scio-core/src/main/scala/com/spotify/scio/testing/TestDataManager.scala

@@ -64,8 +64,8 @@ private[scio] class TestInput(val m: Map[String, JobInputSource[_]]) {
   val s: MSet[String] =
     java.util.concurrent.ConcurrentHashMap.newKeySet[String]().asScala
 
-  def apply[T](io: ScioIO[T]): JobInputSource[T] = {
-    val key = io.testId
+  def apply[T](io: ScioIO[T]): JobInputSource[T] = apply(io.testId)
+  def apply[T](key: String): JobInputSource[T] = {
     require(
       m.contains(key),
       s"Missing test input: $key, available: ${m.keys.mkString("[", ", ", "]")}"
@@ -85,10 +85,9 @@ private[scio] class TestInput(val m: Map[String, JobInputSource[_]]) {
 private[scio] class TestOutput(val m: Map[String, SCollection[_] => Any]) {
   val s: MSet[String] =
     java.util.concurrent.ConcurrentHashMap.newKeySet[String]().asScala
-
-  def apply[T](io: ScioIO[T]): SCollection[T] => Any = {
+  def apply[T](io: ScioIO[T]): SCollection[T] => Any = apply(io.testId)
+  def apply[T](key: String): SCollection[T] => Any = {
     // TODO: support Materialize outputs, maybe Materialized[T]?
-    val key = io.testId
     require(
       m.contains(key),
       s"Missing test output: $key, available: ${m.keys.mkString("[", ", ", "]")}"

scio-core/src/main/scala/com/spotify/scio/testing/TestUtil.scala

@@ -20,11 +20,17 @@ package com.spotify.scio.testing
 import java.util.UUID
 
 private[scio] object TestUtil {
-  def newTestId(className: String = "TestClass"): String = {
+
+  def newTestId(clazz: Class[_]): String = {
+    val className = clazz.getSimpleName
     val uuid = UUID.randomUUID().toString.replaceAll("-", "")
     s"JobTest-$className-$uuid"
   }
+  def newTestId(): String = {
+    val uuid = UUID.randomUUID().toString.replaceAll("-", "")
+    s"JobTest-$uuid"
+  }
 
   def isTestId(appName: String): Boolean =
-    "JobTest-[^-]+-[a-z0-9]+".r.pattern.matcher(appName).matches()
+    "JobTest(-[^-]+)?-[a-z0-9]+".r.pattern.matcher(appName).matches()
 }

scio-core/src/main/scala/com/spotify/scio/values/SCollection.scala

@@ -1435,7 +1435,7 @@ sealed trait SCollection[T] extends PCollectionWrapper[T] {
   )(implicit ev: T <:< String): SCollection[A] =
     if (context.isTest) {
       val id = context.testId.get
-      this.flatMap(s => TestDataManager.getInput(id)(ReadIO(ev(s))).asIterable.get)
+      this.flatMap(s => TestDataManager.getInput(id)(ReadIO[A](ev(s))).asIterable.get)
     } else {
       this
         .covary_[String]

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroSortedBucketIO.java

@@ -36,7 +36,6 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 
 /** API for reading and writing Avro sorted-bucket files. */
 public class AvroSortedBucketIO {
@@ -192,11 +191,6 @@ public static <K1, K2, T extends SpecificRecord> TransformOutput<K1, K2, T> tran
   /** Reads from Avro sorted-bucket files, to be used with {@link SortedBucketIO.CoGbk}. */
   @AutoValue
   public abstract static class Read<T extends IndexedRecord> extends SortedBucketIO.Read<T> {
-    @Nullable
-    abstract ImmutableList<String> getInputDirectories();
-
-    abstract String getFilenameSuffix();
-
     @Nullable
     abstract Schema getSchema();
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/JsonSortedBucketIO.java

@@ -33,7 +33,6 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 
 /** API for reading and writing BigQuery {@link TableRow} JSON sorted-bucket files. */
 public class JsonSortedBucketIO {
@@ -108,10 +107,6 @@ public static <K1, K2> TransformOutput<K1, K2> transformOutput(
    */
   @AutoValue
   public abstract static class Read extends SortedBucketIO.Read<TableRow> {
-    @Nullable
-    abstract ImmutableList<String> getInputDirectories();
-
-    abstract String getFilenameSuffix();
 
     abstract Compression getCompression();
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetAvroSortedBucketIO.java

@@ -36,7 +36,6 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.parquet.filter2.predicate.FilterPredicate;
 import org.apache.parquet.hadoop.metadata.CompressionCodecName;
@@ -200,11 +199,6 @@ public static <K1, K2, T extends SpecificRecord> TransformOutput<K1, K2, T> tran
   /** Reads from Avro sorted-bucket files, to be used with {@link SortedBucketIO.CoGbk}. */
   @AutoValue
   public abstract static class Read<T extends IndexedRecord> extends SortedBucketIO.Read<T> {
-    @Nullable
-    abstract ImmutableList<String> getInputDirectories();
-
-    abstract String getFilenameSuffix();
-
     @Nullable
     abstract Schema getSchema();
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/SortedBucketIO.java

@@ -493,6 +493,12 @@ public abstract static class TransformOutput<K1, K2, V> implements Serializable
 
   /** Represents a single sorted-bucket source written using {@link SortedBucketSink}. */
   public abstract static class Read<V> implements Serializable {
+
+    @Nullable
+    abstract ImmutableList<String> getInputDirectories();
+
+    abstract String getFilenameSuffix();
+
     public abstract TupleTag<V> getTupleTag();
 
     protected abstract BucketedInput<V> toBucketedInput(SortedBucketSource.Keying keying);

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/TensorFlowBucketIO.java

@@ -30,7 +30,6 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 import org.tensorflow.proto.example.Example;
 
 /**
@@ -125,10 +124,6 @@ public static <K1, K2> TransformOutput<K1, K2> transformOutput(
    */
   @AutoValue
   public abstract static class Read extends SortedBucketIO.Read<Example> {
-    @Nullable
-    abstract ImmutableList<String> getInputDirectories();
-
-    abstract String getFilenameSuffix();
 
     abstract Compression getCompression();