vicennial commented on code in PR #40256: URL: https://github.com/apache/spark/pull/40256#discussion_r1123897616
########## connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/ArtifactSuite.scala: ########## @@ -0,0 +1,241 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.connect.client + +import java.io.InputStream +import java.nio.file.{Files, Path, Paths} +import java.util.concurrent.TimeUnit +import java.util.zip.{CheckedInputStream, CRC32} + +import com.google.protobuf.ByteString +import io.grpc.{ManagedChannel, Server} +import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} +import org.scalatest.BeforeAndAfterEach + +import org.apache.spark.connect.proto +import org.apache.spark.connect.proto.AddArtifactsRequest +import org.apache.spark.sql.connect.client.util.ConnectFunSuite + +class ArtifactSuite extends ConnectFunSuite with BeforeAndAfterEach { + + private var client: SparkConnectClient = _ + private var service: DummySparkConnectService = _ + private var server: Server = _ + private var artifactManager: ArtifactManager = _ + private var channel: ManagedChannel = _ + + private def startDummyServer(): Unit = { + service = new DummySparkConnectService() + server = InProcessServerBuilder + .forName(getClass.getName) + .addService(service) + .build() + server.start() + } + + private def createArtifactManager(): Unit = { + channel = InProcessChannelBuilder.forName(getClass.getName).directExecutor().build() + artifactManager = new ArtifactManager(proto.UserContext.newBuilder().build(), channel) + } + + override def beforeEach(): Unit = { + super.beforeEach() + startDummyServer() + createArtifactManager() + client = null + } + + override def afterEach(): Unit = { + if (server != null) { + server.shutdownNow() + assert(server.awaitTermination(5, TimeUnit.SECONDS), "server failed to shutdown") + } + + if (channel != null) { + channel.shutdownNow() + } + + if (client != null) { + client.shutdown() + } + } + + private val CHUNK_SIZE: Int = 32 * 1024 + protected def artifactFilePath: Path = baseResourcePath.resolve("artifact-tests") + + /** + * Check if the data sent to the server (stored in `artifactChunk`) is equivalent to the local + * data at `localPath`. + * @param artifactChunk + * @param localPath + */ + private def assertFileDataEquality( + artifactChunk: AddArtifactsRequest.ArtifactChunk, + localPath: Path): Unit = { + val in = new CheckedInputStream(Files.newInputStream(localPath), new CRC32) + val localData = ByteString.readFrom(in) + assert(artifactChunk.getData == localData) + assert(artifactChunk.getCrc == in.getChecksum.getValue) + } + + private def singleChunkArtifactTest(path: String): Unit = { + test(s"Single Chunk Artifact - $path") { + val artifactPath = artifactFilePath.resolve(path) + artifactManager.addArtifact(artifactPath.toString) + + val receivedRequests = service.getAndClearLatestAddArtifactRequests() + // Single `AddArtifactRequest` + assert(receivedRequests.size == 1) + + val request = receivedRequests.head + assert(request.hasBatch) + + val batch = request.getBatch + // Single artifact in batch + assert(batch.getArtifactsList.size() == 1) + + val singleChunkArtifact = batch.getArtifacts(0) + val namePrefix = artifactPath.getFileName.toString match { + case jar if jar.endsWith(".jar") => "jars" + case cf if cf.endsWith(".class") => "classes" + } + assert(singleChunkArtifact.getName.equals(namePrefix + "/" + path)) + assertFileDataEquality(singleChunkArtifact.getData, artifactPath) + } + } + + singleChunkArtifactTest("smallClassFile.class") + + singleChunkArtifactTest("smallJar.jar") + + private def readNextChunk(in: InputStream): ByteString = { + val buf = new Array[Byte](CHUNK_SIZE) + var bytesRead = 0 + var count = 0 + while (count != -1 && bytesRead < CHUNK_SIZE) { + count = in.read(buf, bytesRead, CHUNK_SIZE - bytesRead) + if (count != -1) { + bytesRead += count + } + } + if (bytesRead == 0) ByteString.empty() + else ByteString.copyFrom(buf, 0, bytesRead) + } + + /** + * Read data in a chunk of `CHUNK_SIZE` bytes from `in` and verify equality with server-side + * data stored in `chunk`. + * @param in + * @param chunk + * @return + */ + private def checkChunkDataAndCrc( Review Comment: Added truth/golden files 👍 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
