public static <K, V> WriteRecords<K, V> writeRecords() {
    return new AutoValue_KafkaIO_WriteRecords.Builder<K, V>()
        .setProducerConfig(WriteRecords.DEFAULT_PRODUCER_PROPERTIES)
        .setEOS(false)
        .setNumShards(0)
        .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
        .setBadRecordRouter(BadRecordRouter.THROWING_ROUTER)
        .setBadRecordErrorHandler(new DefaultErrorHandler<>())
        .build();
  }
@Test
  public void testSinkProducerRecordsWithCustomPartition() throws Exception {
    int numElements = 1000;

    try (MockProducerWrapper producerWrapper = new MockProducerWrapper(new LongSerializer())) {

      ProducerSendCompletionThread completionThread =
          new ProducerSendCompletionThread(producerWrapper.mockProducer).start();

      final String defaultTopic = "test";
      final Integer partition = 1;

      p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata())
          .apply(ParDo.of(new KV2ProducerRecord(defaultTopic, partition)))
          .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of()))
          .apply(
              KafkaIO.<Integer, Long>writeRecords()
                  .withBootstrapServers("none")
                  .withKeySerializer(IntegerSerializer.class)
                  .withValueSerializer(LongSerializer.class)
                  .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));

      p.run();

      completionThread.shutdown();

      // Verify that messages are written with user-defined timestamp
      List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history();

      for (int i = 0; i < numElements; i++) {
        ProducerRecord<Integer, Long> record = sent.get(i);
        assertEquals(defaultTopic, record.topic());
        assertEquals(partition, record.partition());
        assertEquals(i, record.key().intValue());
        assertEquals(i, record.value().longValue());
      }
    }
  }