@Test
  public void testDrainOnContainerStart() {
    int numPageViews = 200;
    int numPartitions = 4;
    long delayBetweenMessagesInMillis = 500L;
    String runId = "DrainTestId";

    InMemorySystemDescriptor isd = new InMemorySystemDescriptor("test");
    InMemoryInputDescriptor<TestTableData.PageView> inputDescriptor = isd.getInputDescriptor("PageView", new NoOpSerde<>());
    InMemoryMetadataStoreFactory metadataStoreFactory = new InMemoryMetadataStoreFactory();

    Map<String, String> customConfig = ImmutableMap.of(
        ApplicationConfig.APP_RUN_ID, runId,
        JobConfig.DRAIN_MONITOR_POLL_INTERVAL_MILLIS, "100",
        JobConfig.DRAIN_MONITOR_ENABLED, "true");

    // Create a TestRunner
    // Set a InMemoryMetadataFactory.This factory is shared between TestRunner and DrainUtils's write drain method
    TestRunner testRunner = TestRunner.of(new PageViewEventCountLowLevelApplication())
        .setInMemoryMetadataFactory(metadataStoreFactory)
        .addConfig(customConfig)
        .addInputStream(inputDescriptor, TestTableData.generatePartitionedPageViews(numPageViews, numPartitions), delayBetweenMessagesInMillis);

    Config configFromRunner = testRunner.getConfig();
    MetadataStore
        metadataStore = metadataStoreFactory.getMetadataStore("NoOp", configFromRunner, new MetricsRegistryMap());

    // Write configs to the coordinator stream here as neither the passthrough JC nor the StreamProcessor is writing
    // configs to coordinator stream. RemoteApplicationRunner typically write the configs to the metadata store
    // before starting the JC.
    // We are doing this so that DrainUtils.writeDrainNotification can read app.run.id from the config
    CoordinatorStreamUtil.writeConfigToCoordinatorStream(metadataStore, configFromRunner);

    // write on the test thread to ensure that drain notification is available on container start
    DrainUtils.writeDrainNotification(metadataStore);

    testRunner.run(Duration.ofSeconds(40));

    assertEquals(RECEIVED.size(), 0);
    RECEIVED.clear();
    clearMetadataStore(metadataStore);
  }
This test will test drain and consumption of some messages from the in-memory topic.
In order to simulate the real-world behaviour of drain, the test adds messages to the in-memory topic buffer in
a delayed fashion instead of all at once. The test then writes the drain notification message to the in-memory
metadata store to drain and stop the pipeline. This write is done shortly after the pipeline starts and before all
the messages are written to the topic's buffer. As a result, the total count of the processed messages will be less
than the expected count of messages.
samza-test/src/test/java/org/apache/samza/test/drain/DrainLowLevelApiIntegrationTest.java