package com.microsoft.cdm.test import com.microsoft.cdm.utils.{AADProvider, ADLGen2Provider, CDMModel} import org.apache.spark.sql.{Row, SparkSession} import org.scalatest.{BeforeAndAfter, FunSuite} import scala.util.{Random, Try} class ReadWriteTests extends FunSuite with BeforeAndAfter { private val appId = sys.env("APP_ID") private val appKey = sys.env("APP_KEY") private val tenantId = sys.env("TENANT_ID") private val inputModel = sys.env("DEMO_INPUT_MODEL") private val outputTestDir = sys.env("OUTPUT_TEST_DIR") private val spark = SparkSession.builder().master("local").appName("DemoTest").getOrCreate() spark.sparkContext.setLogLevel("ERROR") test("read and write basic CDM folders") { val entities = getEntities(inputModel) val outputDir = outputTestDir + "output" + Random.alphanumeric.take(5).mkString("") + "/" val outputModelName = Random.alphanumeric.take(5).mkString("") val collections: Map[String, Int] = entities map (t => t -> readWrite(inputModel, t, write=true, outputDir, outputModelName).length) toMap val outputModel = outputDir + "model.json" val outputEntities = getEntities(outputModel) assert(outputEntities.size == entities.size) collections.foreach{case (entity, size) => assert(size == readWrite(outputModel, entity, write=false).length) } println("Done!") } private def getEntities(modelUri: String): Iterable[String] = { val aadProvider = new AADProvider(appId, appKey, tenantId) val adls = new ADLGen2Provider(aadProvider) val modelJson = new CDMModel(adls.getFullFile(modelUri)) modelJson.listEntities() } private def readWrite(modelUri: String, entity: String, write: Boolean = true, outputDirectory: String = "", outputModelName: String = ""): Array[Row] = { println("%s : %s".format(entity, modelUri)) val df = spark.read.format("com.microsoft.cdm") .option("cdmModel", modelUri) .option("entity", entity) .option("appId", appId) .option("appKey", appKey) .option("tenantId", tenantId) .load() val collection = Try(df.collect()).getOrElse(null) if(write) { df.write.format("com.microsoft.cdm") .option("entity", entity) .option("appId", appId) .option("appKey", appKey) .option("tenantId", tenantId) .option("cdmFolder", outputDirectory) .option("cdmModelName", outputModelName) .save() } collection } }