wip: Added unit test

Author: kr-igor

dd-java-agent/instrumentation/spark-executor/build.gradle

@@ -33,13 +33,24 @@ ext {
 dependencies {
   compileOnly group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
   compileOnly group: 'org.apache.spark', name: 'spark-sql_2.12', version: '2.4.0'
+  compileOnly group: 'org.apache.spark', name:'spark-sql-kafka-0-10_2.12', version: "2.4.0"
 
   baseTestImplementation group: 'org.apache.spark', name: "spark-core_2.12", version: "2.4.0"
   baseTestImplementation group: 'org.apache.spark', name: "spark-sql_2.12", version: "2.4.0"
+  baseTestImplementation group: 'org.apache.spark', name: "spark-sql_2.12", version: "2.4.0"
+  baseTestImplementation group: 'org.apache.spark', name:'spark-sql-kafka-0-10_2.12', version: "2.4.0"
+  testImplementation group: 'org.apache.kafka', name: 'kafka_2.12', version: '2.4.0'
+  testImplementation group: 'org.apache.kafka', name: 'kafka-clients', version: '2.4.0'
+  testImplementation group: 'org.springframework.kafka', name: 'spring-kafka', version: '2.4.0.RELEASE'
+  testImplementation group: 'org.springframework.kafka', name: 'spring-kafka-test', version: '2.4.0.RELEASE'
 
   latest212DepTestImplementation group: 'org.apache.spark', name: "spark-core_2.12", version: '3.+'
   latest212DepTestImplementation group: 'org.apache.spark', name: "spark-sql_2.12", version: '3.+'
+  latest212DepTestImplementation group: 'org.apache.spark', name: "spark-sql_2.12", version: "3.+"
+  latest212DepTestImplementation group: 'org.apache.spark', name:'spark-sql-kafka-0-10_2.12', version: "2.4.0"
 
   latest213DepTestImplementation group: 'org.apache.spark', name: "spark-core_2.13", version: '3.+'
   latest213DepTestImplementation group: 'org.apache.spark', name: "spark-sql_2.13", version: '3.+'
+  latest212DepTestImplementation group: 'org.apache.spark', name: "spark-sql_2.13", version: "3.+"
+  latest212DepTestImplementation group: 'org.apache.spark', name:'spark-sql-kafka-0-10_2.13', version: "3.+"
 }

dd-java-agent/instrumentation/spark-executor/src/baseTest/groovy/SparkExecutorTest.groovy

@@ -1,17 +1,38 @@
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.spark.api.java.function.VoidFunction2
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.RowFactory
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.streaming.Trigger
 import org.apache.spark.sql.types.StructType
+import org.junit.ClassRule
+import org.springframework.kafka.core.DefaultKafkaProducerFactory
+import org.springframework.kafka.test.EmbeddedKafkaBroker
+import org.springframework.kafka.test.rule.EmbeddedKafkaRule
+import org.springframework.kafka.test.utils.KafkaTestUtils
+import spock.lang.Shared
+
 
 class SparkExecutorTest extends AgentTestRunner {
+  static final SOURCE_TOPIC = "source"
+  static final SINK_TOPIC = "sink"
+
+  @Shared
+  @ClassRule
+  EmbeddedKafkaRule kafkaRule = new EmbeddedKafkaRule(1, false, 1, SOURCE_TOPIC, SINK_TOPIC)
+  EmbeddedKafkaBroker embeddedKafka = kafkaRule.embeddedKafka
 
   @Override
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark-executor.enabled", "true")
+    injectSysConfig("dd.integration.spark.enabled", "true")
+    injectSysConfig("dd.integration.kafka.enabled", "true")
+    injectSysConfig("dd.data.streams.enabled", "true")
+    injectSysConfig("dd.trace.debug", "true")
   }
 
   private Dataset<Row> generateSampleDataframe(SparkSession spark) {
@@ -23,6 +44,57 @@ class SparkExecutorTest extends AgentTestRunner {
     spark.createDataFrame(rows, structType)
   }
 
+  def "test dsm service name override"() {
+    setup:
+    def sparkSession = SparkSession.builder()
+      .config("spark.master", "local[2]")
+      .config("spark.driver.bindAddress", "localhost")
+      //      .config("spark.sql.shuffle.partitions", "2")
+      .appName("test-app")
+      .getOrCreate()
+
+    def producerProps = KafkaTestUtils.producerProps(embeddedKafka.getBrokersAsString())
+    def producer = new DefaultKafkaProducerFactory<Integer, String>(producerProps).createProducer()
+
+    when:
+    for (int i = 0; i < 100; i++) {
+      producer.send(new ProducerRecord<>(SOURCE_TOPIC, i, i.toString()))
+    }
+    producer.flush()
+
+    def df = sparkSession
+      .readStream()
+      .format("kafka")
+      .option("kafka.bootstrap.servers", embeddedKafka.getBrokersAsString())
+      .option("startingOffsets", "earliest")
+      .option("failOnDataLoss", "false")
+      .option("subscribe", SOURCE_TOPIC)
+      .load()
+
+    def query = df
+      .selectExpr("CAST(key AS STRING) as key", "CAST(value AS STRING) as value")
+      .writeStream()
+      .format("kafka")
+      .option("kafka.bootstrap.servers", embeddedKafka.getBrokersAsString())
+      .option("checkpointLocation", "/tmp/" + System.currentTimeMillis().toString())
+      .option("topic", SINK_TOPIC)
+      .trigger(Trigger.Once())
+      .foreachBatch(new VoidFunction2<Dataset<Row>, Long>() {
+        @Override
+        void call(Dataset<Row> rowDataset, Long aLong) throws Exception {
+          rowDataset.show()
+          rowDataset.write()
+        }
+      })
+      .start()
+
+    query.processAllAvailable()
+
+    then:
+    query.stop()
+    producer.close()
+  }
+
   def "generate spark task run spans"() {
     setup:
     def sparkSession = SparkSession.builder()