Apache Beam，在 DoFn 的 @Setup Lifecycle 方法中初始化的模拟外部客户端

Question

我已经创建了以下 DoFn 类，但我不确定如何模拟 APIClient，因为在 Junit 中使用模拟的 APICall 对象创建 JsonToGenericRecordMapper 的实例，被 @setup 方法覆盖，该方法正在创建 APICall 的实际实例？

public class JsonToGenericRecordMapper extends DoFn<String, GenericRecord> {
    
    private APICall apiCall;
    
     @setup
     public void setup () 
    {
        api = new APICall();
    }
    
    @ProcessElement
    public void processElement(ProcessContext processContext) {

        String enrichedItemCost = processContext.element();

        Schema schema = CommonUtils.schema;
        GenericRecord genericRecord = new GenericData.Record(schema);
        genericRecord.put(CommonUtils.SCHEMA_FIELD_KEY, enrichedItemCost);

        processContext.output(genericRecord);
    }
}




TestPipeline testPipeline = TestPipeline.create();
APICAll apiCall = Mockito.mock(APICall.class);
dataPCollection.apply(new JsonToGenericRecordMapper(apiCall));

Answer 1

您可以创建两个构造函数（一个空的，一个带有 APICall 参数的构造函数），并使用其中一个进行测试，类似的东西（我只是扩展了您的示例来说明这一点，所以这只是一个示例）：

public class JsonToGenericRecordMapper extends DoFn<String, GenericRecord> {

private APICall apiCall;

public JsonToGenericRecordMapper() {
}

public JsonToGenericRecordMapper(APICall apiCall) {
    this.apiCall = apiCall;
}


 @setup
 public void setup () 
 { 
    if(apiCall==null) {
        api = new APICall();
    }
 }

@ProcessElement
public void processElement(ProcessContext processContext) {

    String enrichedItemCost = processContext.element();

    Schema schema = CommonUtils.schema;
    GenericRecord genericRecord = new GenericData.Record(schema);
    genericRecord.put(CommonUtils.SCHEMA_FIELD_KEY, enrichedItemCost);

    processContext.output(genericRecord);
  }
}

TestPipeline testPipeline = TestPipeline.create();
APICAll apiCall = Mockito.mock(APICall.class);
dataPCollection.apply(new JsonToGenericRecordMapper(apiCall));

Answer 2

DoFn 的模拟组件不能很好地工作，因为模拟在管道执行期间无法在序列化过程中幸存下来。最好使用 fakes，而不是模拟，正如 Beam 文档的“使用 Fakes”部分中所建议的那样 (https://beam.apache.org/documentation/io/testing/)。以下是我如何使用 kotlin 中的工作实现对 DoFn 进行单元测试：

JsonToGenericRecordMapper

class JsonToGenericRecordMapper : DoFn<String?, GenericRecord?>() {
    @Transient
    private var apiClient: ApiClient? = null

    @Setup
    fun setup() {
        // Initial setup for non-test execution
        apiClient = ApiClient()
    }

    @ProcessElement
    fun processElement(processContext: ProcessContext) {
        var response = apiClient?.call()
        // Continue with your processElement logic
    }

    @VisibleForTesting
    fun setApiClientForTests(client: ApiClient) {
        this.apiClient = client
    }
}

ApiClient 与假的

import java.io.Serializable as JavaSerializable

open class ApiClient : JavaSerializable {
    open fun call(): String {
        // Actual API call
        return "API Response"
    }
}

class FakeApiClient : ApiClient() {
    override fun call(): String {
        // Mocked API call
        return "Mocked response"
    }
}

测试

class JsonToGenericRecordMapperTest {
        @get:Rule
        val pipeline = TestPipeline.create()
    
        @Test
        fun `test JsonToGenericRecordMapper processes elements correctly`() {
            val fakeApiClient = FakeApiClient()
            val mapper = JsonToGenericRecordMapper()
                .apply {
                    setApiClientForTests(fakeApiClient)
                }
    
            val expectedOutput = /* Define your expected output */
            val input = "Test input"
    
            // Prepare the input PCollection
            val inputPCollection = pipeline.apply(Create.of(input))
    
            // Apply the DoFn
            val output: PCollection = inputPCollection.apply(
                ParDo.of(
                    mapper
                )
            )
    
            PAssert.that(output).containsInAnyOrder(expectedOutput)
            pipeline.run().waitUntilFinish()
        }
    }

但是，您仍然可以在假对象中有效地使用模拟。当您需要在测试环境中模拟复杂的行为或外部依赖项（例如 OAuth 令牌生成）时，此策略特别有用。为了方便在 Kotlin 中进行模拟，尤其是在使用 Gradle 时，请确保在项目依赖项中包含 Mockito Kotlin 库。这是所需的依赖项：

添加到您的 build.gradle.kts (Kotlin DSL)：

testImplementation group: 'org.mockito.kotlin', name: 'mockito-kotlin', version: '5.2.1'

JsonToGenericRecordMapper

import com.google.auth.oauth2.IdToken
    
class JsonToGenericRecordMapper : DoFn<String?, GenericRecord?>() {

    @Transient
    private var apiClient: ApiClient? = null

    @Setup
    fun setup() {
        apiClient = ApiClient()
    }

    @ProcessElement
    fun processElement(processContext: ProcessContext) {
        // Example of using the IdToken within the processElement method
        val idToken = apiClient?.createIdToken()
        val accessToken = idToken.tokenValue
    }
}

客户和假客户

open class ApiClient : JavaSerializable {
    open fun createIdToken(): IdToken {
        /* actual generation logic here */
        return IdToken.create("actual token")
    }
}

import org.mockito.kotlin.mock
import org.mockito.kotlin.whenever

class FakeApiClient : ApiClient() {
    override fun createIdToken(): IdToken {
        val fakeTokenIdToken: IdToken = mock()  
        whenever(fakeTokenIdToken.tokenValue).thenReturn("fake token")
        return fakeTokenIdToken
    }
}