improved OPCollectionHashingVectorizerTest test (#286)

* improved OPCollectionHashingVectorizerTest test
salesforce · Apr 16, 2019 · 699f978 · 699f978
1 parent 2c06124
commit 699f978
Showing 1 changed file with 20 additions and 6 deletions.
diff --git a/.../test/scala/com/salesforce/op/stages/impl/feature/OPCollectionHashingVectorizerTest.scala b/.../test/scala/com/salesforce/op/stages/impl/feature/OPCollectionHashingVectorizerTest.scala
@@ -32,16 +32,31 @@ package com.salesforce.op.stages.impl.feature
 
 import com.salesforce.op.features.types._
 import com.salesforce.op.test.TestOpVectorColumnType.{IndCol, PivotColNoInd}
-import com.salesforce.op.test.{TestFeatureBuilder, TestOpVectorMetadataBuilder, TestSparkContext}
+import com.salesforce.op.test.{OpTransformerSpec, TestFeatureBuilder, TestOpVectorMetadataBuilder}
 import com.salesforce.op.utils.spark.OpVectorMetadata
 import com.salesforce.op.utils.spark.RichDataset._
-import com.salesforce.op.utils.spark.RichMetadata._
+import org.apache.spark.sql.Dataset
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
-import org.scalatest.{Assertions, FlatSpec, Matchers}
+import org.apache.spark.ml.linalg.Vectors
 
 @RunWith(classOf[JUnitRunner])
-class OPCollectionHashingVectorizerTest extends FlatSpec with TestSparkContext with AttributeAsserts {
+class OPCollectionHashingVectorizerTest extends OpTransformerSpec[OPVector, OPCollectionHashingVectorizer[TextList]]
+ with AttributeAsserts {
+
+ val (inputData, f1) =
+ TestFeatureBuilder("textList1", Seq[TextList](TextList(Seq("x", "y"))))
+
+ /**
+ * [[OpTransformer]] instance to be tested
+ */
+ override val transformer: OPCollectionHashingVectorizer[TextList] = new OPCollectionHashingVectorizer().setInput(f1)
+
+ /**
+ * Expected result of the transformer applied on the Input Dataset
+ */
+ override val expectedResult: Seq[OPVector] = Seq(OPVector(Vectors.sparse(512, Array(107, 224), Array(1.0, 1.0))))
+
 
  val (catData, top, bot) = TestFeatureBuilder("top", "bot",
  Seq[(MultiPickList, MultiPickList)](
@@ -68,7 +83,7 @@ class OPCollectionHashingVectorizerTest extends FlatSpec with TestSparkContext w
  )
  )
 
- Spec[OPCollectionHashingVectorizer[_]] should "have the correct number of features" in {
+ it should "have the correct number of features" in {
  val vectorizer = new OPCollectionHashingVectorizer[MultiPickList].setInput(top, bot)
  vectorizer.setNumFeatures(128).getNumFeatures() shouldBe 128
  }
@@ -446,5 +461,4 @@ class OPCollectionHashingVectorizerTest extends FlatSpec with TestSparkContext w
  )
  meta.getColumnHistory().length shouldBe 10
  }
-
 }