Skip to content

Commit dc2116b

Browse files
CoderYellowCoderYellow
committed
refactor: extract reusable Spark test logic into SparkBigDataScenario
- Moved common Spark test setup and verification logic into an abstract `SparkBigDataScenario`. - Simplified `SparkBigDataTestExample` by extending `SparkBigDataScenario` and overriding key properties. - Improved modularity, reusability, and maintainability of Spark test infrastructure.
1 parent 16f1135 commit dc2116b

2 files changed

Lines changed: 337 additions & 249 deletions

File tree

‎example/spark/src/test/kotlin/org/openprojectx/bigdata/test/example/spark/SparkBigDataScenario.kt‎

Lines changed: 327 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,327 @@
1+
package org.openprojectx.bigdata.test.example.spark
2+
3+
import org.apache.hadoop.hive.conf.HiveConf
4+
import org.apache.hadoop.hive.metastore.HiveMetaStoreClient
5+
import org.apache.spark.sql.SparkSession
6+
import org.junit.jupiter.api.Test
7+
import org.openprojectx.bigdata.test.core.BigDataService
8+
import org.openprojectx.bigdata.test.core.BigDataTestKit
9+
import org.openprojectx.bigdata.test.extensions.core.BigDataExtensionResult
10+
import java.net.URI
11+
import java.nio.file.Files
12+
13+
abstract class SparkBigDataScenario {
14+
protected abstract val runId: String
15+
protected abstract val s3BucketExtensionId: String
16+
protected abstract val gcsBucketExtensionId: String
17+
18+
@Test
19+
fun runsSparkScenario(kit: BigDataTestKit, extensions: BigDataExtensionResult) {
20+
val environment = createEnvironment(kit, extensions)
21+
createSparkSession(environment).use { spark ->
22+
val context = SparkScenarioContext(environment, spark)
23+
infrastructureChecks().forEach { it.verify(context) }
24+
sourceChecks().forEach { it.verify(context) }
25+
targetChecks().forEach { it.verify(context) }
26+
}
27+
}
28+
29+
protected open fun infrastructureChecks(): List<SparkScenarioCheck> =
30+
listOf(
31+
SparkScenarioCheck { context ->
32+
assertHdfsConfigStore(
33+
spark = context.spark,
34+
hdfsUri = context.environment.hdfsUri,
35+
hdfsPath = context.environment.s3CredentialProviderHdfsPath,
36+
)
37+
},
38+
)
39+
40+
protected open fun sourceChecks(): List<SparkScenarioCheck> =
41+
listOf(
42+
SparkScenarioCheck { context ->
43+
assertKafkaAvroInput(
44+
spark = context.spark,
45+
bootstrapServers = context.environment.kafkaBootstrapServers,
46+
topic = kafkaAvroTopic(),
47+
)
48+
},
49+
)
50+
51+
protected open fun targetChecks(): List<SparkScenarioCheck> =
52+
listOf(
53+
SparkScenarioCheck { context ->
54+
assertIcebergTable(
55+
spark = context.spark,
56+
catalog = "s3",
57+
namespace = "demo_$runId",
58+
table = "events_s3",
59+
storageName = "s3",
60+
)
61+
},
62+
SparkScenarioCheck { context ->
63+
assertIcebergTable(
64+
spark = context.spark,
65+
catalog = "gcs_local",
66+
namespace = "demo_$runId",
67+
table = "events_gcs",
68+
storageName = "gcs",
69+
dataPath = context.environment.gcsIcebergDataPath,
70+
)
71+
},
72+
SparkScenarioCheck { context ->
73+
assertIcebergTable(
74+
spark = context.spark,
75+
catalog = "hms",
76+
namespace = "hms_demo_$runId",
77+
table = "events_hms",
78+
storageName = "hms",
79+
)
80+
assertHiveMetastoreTable(
81+
hiveMetastoreUri = context.environment.hiveMetastoreUri,
82+
database = "hms_demo_$runId",
83+
table = "events_hms",
84+
)
85+
},
86+
SparkScenarioCheck { context ->
87+
assertHiveExternalParquetTable(
88+
spark = context.spark,
89+
hiveMetastoreUri = context.environment.hiveMetastoreUri,
90+
database = "hive_s3_demo_$runId",
91+
table = "events_parquet_s3",
92+
location = "s3a://${context.environment.s3Bucket}/hive-parquet/events_parquet_s3",
93+
storageName = "hive-s3-parquet",
94+
)
95+
},
96+
SparkScenarioCheck { context ->
97+
assertHiveExternalParquetTable(
98+
spark = context.spark,
99+
hiveMetastoreUri = context.environment.hiveMetastoreUri,
100+
database = "hive_gcs_demo_$runId",
101+
table = "events_parquet_gcs",
102+
location = context.environment.gcsIcebergDataPath,
103+
storageName = "gcs",
104+
seedData = false,
105+
)
106+
},
107+
)
108+
109+
protected open fun kafkaAvroTopic(): String = "spark-avro-events"
110+
111+
protected open fun configureSpark(
112+
builder: SparkSession.Builder,
113+
environment: SparkScenarioEnvironment,
114+
): SparkSession.Builder = builder
115+
116+
private fun createEnvironment(kit: BigDataTestKit, extensions: BigDataExtensionResult): SparkScenarioEnvironment {
117+
val hdfs = kit.endpoint(BigDataService.HDFS)
118+
val hiveMetastore = kit.endpoint(BigDataService.HIVE_METASTORE)
119+
val kafka = kit.endpoint(BigDataService.KAFKA)
120+
val s3 = kit.endpoint(BigDataService.LOCALSTACK_S3)
121+
val gcs = kit.endpoint(BigDataService.FAKE_GCS)
122+
val gcsBucket = extensions.required("$gcsBucketExtensionId.bucket")
123+
return SparkScenarioEnvironment(
124+
runId = runId,
125+
hdfsUri = hdfs.property("fs.defaultFS"),
126+
hiveMetastoreUri = hiveMetastore.property("hive.metastore.uris"),
127+
kafkaBootstrapServers = kafka.property("bootstrap.servers"),
128+
s3Endpoint = s3.property("aws.endpoint-url.s3"),
129+
s3Bucket = extensions.required("$s3BucketExtensionId.bucket"),
130+
gcsEndpoint = gcs.property("google.cloud.storage.host"),
131+
gcsBucket = gcsBucket,
132+
gcsIcebergDataPath = "${extensions.required("$gcsBucketExtensionId.gs.uri")}/data/demo_$runId/events_gcs",
133+
s3CredentialProviderPath = extensions.required("s3-jceks.credential-provider.path"),
134+
s3CredentialProviderHdfsPath = extensions.required("s3-jceks.hdfs.path"),
135+
)
136+
}
137+
138+
private fun createSparkSession(environment: SparkScenarioEnvironment): SparkSession =
139+
configureSpark(
140+
SparkSession.builder()
141+
.appName("bigdata-test-spark-example")
142+
.master("local[2]")
143+
.config("spark.ui.enabled", "false")
144+
.config("spark.driver.extraJavaOptions", "--add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED")
145+
.config("spark.executor.extraJavaOptions", "--add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED")
146+
.config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
147+
.config("spark.sql.catalog.s3", "org.apache.iceberg.spark.SparkCatalog")
148+
.config("spark.sql.catalog.s3.type", "hadoop")
149+
.config("spark.sql.catalog.s3.warehouse", "s3a://${environment.s3Bucket}/warehouse")
150+
.config("spark.sql.catalog.gcs_local", "org.apache.iceberg.spark.SparkCatalog")
151+
.config("spark.sql.catalog.gcs_local.type", "hadoop")
152+
.config("spark.sql.catalog.gcs_local.warehouse", "file:${Files.createTempDirectory("bigdata-test-gcs-iceberg-warehouse-")}")
153+
.config("spark.sql.catalog.hms", "org.apache.iceberg.spark.SparkCatalog")
154+
.config("spark.sql.catalog.hms.type", "hive")
155+
.config("spark.sql.catalog.hms.uri", environment.hiveMetastoreUri)
156+
.config("spark.sql.catalog.hms.warehouse", "file:${Files.createTempDirectory("bigdata-test-hms-warehouse-")}")
157+
.config("spark.sql.warehouse.dir", "file:${Files.createTempDirectory("bigdata-test-spark-warehouse-")}")
158+
.config("hive.metastore.uris", environment.hiveMetastoreUri)
159+
.config("spark.hadoop.fs.defaultFS", environment.hdfsUri)
160+
.config("spark.hadoop.hadoop.security.credential.provider.path", environment.s3CredentialProviderPath)
161+
.config("spark.hadoop.fs.s3a.endpoint", environment.s3Endpoint)
162+
.config("spark.hadoop.fs.s3a.aws.credentials.provider", "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider")
163+
.config("spark.hadoop.fs.s3a.access.key", "test")
164+
.config("spark.hadoop.fs.s3a.secret.key", "test")
165+
.config("spark.hadoop.fs.s3a.path.style.access", "true")
166+
.config("spark.hadoop.fs.s3a.connection.ssl.enabled", "false")
167+
.config("spark.hadoop.fs.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem")
168+
.config("spark.hadoop.fs.AbstractFileSystem.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFS")
169+
.config("spark.hadoop.fs.gs.project.id", "bigdata-test")
170+
.config("spark.hadoop.fs.gs.storage.root.url", "${environment.gcsEndpoint.trimEnd('/')}/")
171+
.config("spark.hadoop.fs.gs.storage.service.path", "storage/v1/")
172+
.config("spark.hadoop.fs.gs.client.type", "HTTP_API_CLIENT")
173+
.config("spark.hadoop.fs.gs.http.connect-timeout", "4000")
174+
.config("spark.hadoop.fs.gs.auth.type", "UNAUTHENTICATED")
175+
.config("spark.hadoop.fs.gs.create.items.conflict.check.enable", "false")
176+
.config("spark.hadoop.fs.gs.implicit.dir.repair.enable", "false")
177+
.config("spark.hadoop.fs.gs.hierarchical.namespace.folders.enable", "false"),
178+
environment,
179+
).enableHiveSupport()
180+
.getOrCreate()
181+
182+
protected fun assertHdfsConfigStore(spark: SparkSession, hdfsUri: String, hdfsPath: String) {
183+
val exists = org.apache.hadoop.fs.FileSystem.get(URI.create(hdfsUri), spark.sparkContext().hadoopConfiguration())
184+
.use { fs -> fs.exists(org.apache.hadoop.fs.Path(hdfsPath)) }
185+
check(exists) { "Expected S3 JCEKS file in HDFS for ${spark.sparkContext().appName()}" }
186+
}
187+
188+
protected fun assertKafkaAvroInput(spark: SparkSession, bootstrapServers: String, topic: String) {
189+
val rows = spark.read()
190+
.format("kafka")
191+
.option("kafka.bootstrap.servers", bootstrapServers)
192+
.option("subscribe", topic)
193+
.option("startingOffsets", "earliest")
194+
.option("endingOffsets", "latest")
195+
.load()
196+
check(rows.count() == 2L) { "Expected two Avro Kafka records in $topic" }
197+
}
198+
199+
protected fun assertIcebergTable(
200+
spark: SparkSession,
201+
catalog: String,
202+
namespace: String,
203+
table: String,
204+
storageName: String,
205+
dataPath: String? = null,
206+
) {
207+
val identifier = "$catalog.$namespace.$table"
208+
spark.sql("CREATE NAMESPACE IF NOT EXISTS $catalog.$namespace")
209+
val tableProperties = dataPath?.let { " TBLPROPERTIES ('write.data.path'='$it')" }.orEmpty()
210+
spark.sql(
211+
"""
212+
CREATE TABLE $identifier (
213+
id INT,
214+
name STRING,
215+
storage STRING
216+
) USING iceberg$tableProperties
217+
""".trimIndent(),
218+
)
219+
spark.sql("INSERT INTO $identifier VALUES (1, 'alpha', '$storageName'), (2, 'beta', '$storageName')")
220+
val count = spark.table(identifier).where("storage = '$storageName'").count()
221+
check(count == 2L) { "Expected two Iceberg rows in $identifier" }
222+
}
223+
224+
protected fun assertHiveMetastoreTable(hiveMetastoreUri: String, database: String, table: String) {
225+
val conf = HiveConf()
226+
conf.setVar(HiveConf.ConfVars.METASTOREURIS, hiveMetastoreUri)
227+
val client = HiveMetaStoreClient(conf)
228+
try {
229+
check(client.getAllDatabases().contains(database)) { "Expected HMS database $database" }
230+
val hmsTable = client.getTable(database, table)
231+
check(hmsTable.dbName == database) { "Expected HMS table $database.$table, got ${hmsTable.dbName}.${hmsTable.tableName}" }
232+
check(hmsTable.tableName == table) { "Expected HMS table $database.$table, got ${hmsTable.dbName}.${hmsTable.tableName}" }
233+
check(hmsTable.parameters["table_type"] == "ICEBERG") {
234+
"Expected HMS table $database.$table to be an Iceberg table, parameters=${hmsTable.parameters}"
235+
}
236+
check(hmsTable.parameters.containsKey("metadata_location")) {
237+
"Expected HMS table $database.$table to contain Iceberg metadata_location"
238+
}
239+
} finally {
240+
client.close()
241+
}
242+
}
243+
244+
protected fun assertHiveExternalParquetTable(
245+
spark: SparkSession,
246+
hiveMetastoreUri: String,
247+
database: String,
248+
table: String,
249+
location: String,
250+
storageName: String,
251+
seedData: Boolean = true,
252+
) {
253+
val identifier = "$database.$table"
254+
spark.sql("CREATE DATABASE IF NOT EXISTS $database")
255+
if (seedData) {
256+
spark.sql(
257+
"""
258+
SELECT 1 AS id, 'alpha' AS name, '$storageName' AS storage
259+
UNION ALL
260+
SELECT 2 AS id, 'beta' AS name, '$storageName' AS storage
261+
""".trimIndent(),
262+
).write().mode("overwrite").parquet(location)
263+
}
264+
spark.sql(
265+
"""
266+
CREATE EXTERNAL TABLE $identifier (
267+
id INT,
268+
name STRING,
269+
storage STRING
270+
) STORED AS PARQUET
271+
LOCATION '$location'
272+
""".trimIndent(),
273+
)
274+
val count = spark.table(identifier).where("storage = '$storageName'").count()
275+
check(count == 2L) { "Expected two Hive external Parquet rows in $identifier" }
276+
assertParquetFiles(spark, location)
277+
assertHiveMetastoreExternalParquetTable(hiveMetastoreUri, database, table, location)
278+
}
279+
280+
private fun assertParquetFiles(spark: SparkSession, location: String) {
281+
val files = org.apache.hadoop.fs.FileSystem.get(URI.create(location), spark.sparkContext().hadoopConfiguration())
282+
.use { fs -> fs.listStatus(org.apache.hadoop.fs.Path(location)).map { it.path.name } }
283+
check(files.any { it.endsWith(".parquet") }) { "Expected Parquet files under $location, got $files" }
284+
}
285+
286+
private fun assertHiveMetastoreExternalParquetTable(hiveMetastoreUri: String, database: String, table: String, location: String) {
287+
val conf = HiveConf()
288+
conf.setVar(HiveConf.ConfVars.METASTOREURIS, hiveMetastoreUri)
289+
val client = HiveMetaStoreClient(conf)
290+
try {
291+
val hmsTable = client.getTable(database, table)
292+
check(hmsTable.dbName == database) { "Expected HMS table $database.$table, got ${hmsTable.dbName}.${hmsTable.tableName}" }
293+
check(hmsTable.tableName == table) { "Expected HMS table $database.$table, got ${hmsTable.dbName}.${hmsTable.tableName}" }
294+
check(hmsTable.sd.location == location) {
295+
"Expected HMS table $database.$table location $location, got ${hmsTable.sd.location}"
296+
}
297+
check(hmsTable.sd.inputFormat.contains("Parquet")) {
298+
"Expected HMS table $database.$table to be Parquet, inputFormat=${hmsTable.sd.inputFormat}"
299+
}
300+
} finally {
301+
client.close()
302+
}
303+
}
304+
}
305+
306+
fun interface SparkScenarioCheck {
307+
fun verify(context: SparkScenarioContext)
308+
}
309+
310+
data class SparkScenarioContext(
311+
val environment: SparkScenarioEnvironment,
312+
val spark: SparkSession,
313+
)
314+
315+
data class SparkScenarioEnvironment(
316+
val runId: String,
317+
val hdfsUri: String,
318+
val hiveMetastoreUri: String,
319+
val kafkaBootstrapServers: String,
320+
val s3Endpoint: String,
321+
val s3Bucket: String,
322+
val gcsEndpoint: String,
323+
val gcsBucket: String,
324+
val gcsIcebergDataPath: String,
325+
val s3CredentialProviderPath: String,
326+
val s3CredentialProviderHdfsPath: String,
327+
)

0 commit comments

Comments
 (0)