[02/10] kafka git commit: KAFKA-5059: Implement Transactional Coordinator

[gu...@apache.org]

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
new file mode 100644
index 0000000..ab893cd
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
@@ -0,0 +1,362 @@
+/**
+ * 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 kafka.coordinator.group
+
+import kafka.common.OffsetAndMetadata
+
+import org.apache.kafka.common.TopicPartition
+
+import org.junit.Assert._
+import org.junit.{Before, Test}
+import org.scalatest.junit.JUnitSuite
+
+/**
+ * Test group state transitions and other GroupMetadata functionality
+ */
+class GroupMetadataTest extends JUnitSuite {
+  private val protocolType = "consumer"
+  private val groupId = "groupId"
+  private val clientId = "clientId"
+  private val clientHost = "clientHost"
+  private val rebalanceTimeoutMs = 60000
+  private val sessionTimeoutMs = 10000
+
+  private var group: GroupMetadata = null
+
+  @Before
+  def setUp() {
+    group = new GroupMetadata("groupId")
+  }
+
+  @Test
+  def testCanRebalanceWhenStable() {
+    assertTrue(group.canRebalance)
+  }
+
+  @Test
+  def testCanRebalanceWhenAwaitingSync() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(AwaitingSync)
+    assertTrue(group.canRebalance)
+  }
+
+  @Test
+  def testCannotRebalanceWhenPreparingRebalance() {
+    group.transitionTo(PreparingRebalance)
+    assertFalse(group.canRebalance)
+  }
+
+  @Test
+  def testCannotRebalanceWhenDead() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Empty)
+    group.transitionTo(Dead)
+    assertFalse(group.canRebalance)
+  }
+
+  @Test
+  def testStableToPreparingRebalanceTransition() {
+    group.transitionTo(PreparingRebalance)
+    assertState(group, PreparingRebalance)
+  }
+
+  @Test
+  def testStableToDeadTransition() {
+    group.transitionTo(Dead)
+    assertState(group, Dead)
+  }
+
+  @Test
+  def testAwaitingSyncToPreparingRebalanceTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(AwaitingSync)
+    group.transitionTo(PreparingRebalance)
+    assertState(group, PreparingRebalance)
+  }
+
+  @Test
+  def testPreparingRebalanceToDeadTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Dead)
+    assertState(group, Dead)
+  }
+
+  @Test
+  def testPreparingRebalanceToEmptyTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Empty)
+    assertState(group, Empty)
+  }
+
+  @Test
+  def testEmptyToDeadTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Empty)
+    group.transitionTo(Dead)
+    assertState(group, Dead)
+  }
+
+  @Test
+  def testAwaitingSyncToStableTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(AwaitingSync)
+    group.transitionTo(Stable)
+    assertState(group, Stable)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testStableToStableIllegalTransition() {
+    group.transitionTo(Stable)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testStableToAwaitingSyncIllegalTransition() {
+    group.transitionTo(AwaitingSync)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testPreparingRebalanceToPreparingRebalanceIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(PreparingRebalance)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testPreparingRebalanceToStableIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Stable)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testAwaitingSyncToAwaitingSyncIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(AwaitingSync)
+    group.transitionTo(AwaitingSync)
+  }
+
+  def testDeadToDeadIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Dead)
+    group.transitionTo(Dead)
+    assertState(group, Dead)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testDeadToStableIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Dead)
+    group.transitionTo(Stable)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testDeadToPreparingRebalanceIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Dead)
+    group.transitionTo(PreparingRebalance)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testDeadToAwaitingSyncIllegalTransition() {
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Dead)
+    group.transitionTo(AwaitingSync)
+  }
+
+  @Test
+  def testSelectProtocol() {
+    val memberId = "memberId"
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
+
+    group.add(member)
+    assertEquals("range", group.selectProtocol)
+
+    val otherMemberId = "otherMemberId"
+    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("range", Array.empty[Byte])))
+
+    group.add(otherMember)
+    // now could be either range or robin since there is no majority preference
+    assertTrue(Set("range", "roundrobin")(group.selectProtocol))
+
+    val lastMemberId = "lastMemberId"
+    val lastMember = new MemberMetadata(lastMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("range", Array.empty[Byte])))
+
+    group.add(lastMember)
+    // now we should prefer 'roundrobin'
+    assertEquals("roundrobin", group.selectProtocol)
+  }
+
+  @Test(expected = classOf[IllegalStateException])
+  def testSelectProtocolRaisesIfNoMembers() {
+    group.selectProtocol
+    fail()
+  }
+
+  @Test
+  def testSelectProtocolChoosesCompatibleProtocol() {
+    val memberId = "memberId"
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
+
+    val otherMemberId = "otherMemberId"
+    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("blah", Array.empty[Byte])))
+
+    group.add(member)
+    group.add(otherMember)
+    assertEquals("roundrobin", group.selectProtocol)
+  }
+
+  @Test
+  def testSupportsProtocols() {
+    // by default, the group supports everything
+    assertTrue(group.supportsProtocols(Set("roundrobin", "range")))
+
+    val memberId = "memberId"
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
+
+    group.add(member)
+    assertTrue(group.supportsProtocols(Set("roundrobin", "foo")))
+    assertTrue(group.supportsProtocols(Set("range", "foo")))
+    assertFalse(group.supportsProtocols(Set("foo", "bar")))
+
+    val otherMemberId = "otherMemberId"
+    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("blah", Array.empty[Byte])))
+
+    group.add(otherMember)
+
+    assertTrue(group.supportsProtocols(Set("roundrobin", "foo")))
+    assertFalse(group.supportsProtocols(Set("range", "foo")))
+  }
+
+  @Test
+  def testInitNextGeneration() {
+    val memberId = "memberId"
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, List(("roundrobin", Array.empty[Byte])))
+
+    group.transitionTo(PreparingRebalance)
+    member.awaitingJoinCallback = _ => ()
+    group.add(member)
+
+    assertEquals(0, group.generationId)
+    assertNull(group.protocol)
+
+    group.initNextGeneration()
+
+    assertEquals(1, group.generationId)
+    assertEquals("roundrobin", group.protocol)
+  }
+
+  @Test
+  def testInitNextGenerationEmptyGroup() {
+    assertEquals(Empty, group.currentState)
+    assertEquals(0, group.generationId)
+    assertNull(group.protocol)
+
+    group.transitionTo(PreparingRebalance)
+    group.initNextGeneration()
+
+    assertEquals(1, group.generationId)
+    assertNull(group.protocol)
+  }
+
+  @Test
+  def testOffsetCommit(): Unit = {
+    val partition = new TopicPartition("foo", 0)
+    val offset = OffsetAndMetadata(37)
+
+    group.prepareOffsetCommit(Map(partition -> offset))
+    assertTrue(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+
+    group.completePendingOffsetWrite(partition, offset)
+    assertTrue(group.hasOffsets)
+    assertEquals(Some(offset), group.offset(partition))
+  }
+
+  @Test
+  def testOffsetCommitFailure(): Unit = {
+    val partition = new TopicPartition("foo", 0)
+    val offset = OffsetAndMetadata(37)
+
+    group.prepareOffsetCommit(Map(partition -> offset))
+    assertTrue(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+
+    group.failPendingOffsetWrite(partition, offset)
+    assertFalse(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+  }
+
+  @Test
+  def testOffsetCommitFailureWithAnotherPending(): Unit = {
+    val partition = new TopicPartition("foo", 0)
+    val firstOffset = OffsetAndMetadata(37)
+    val secondOffset = OffsetAndMetadata(57)
+
+    group.prepareOffsetCommit(Map(partition -> firstOffset))
+    assertTrue(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+
+    group.prepareOffsetCommit(Map(partition -> secondOffset))
+    assertTrue(group.hasOffsets)
+
+    group.failPendingOffsetWrite(partition, firstOffset)
+    assertTrue(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+
+    group.completePendingOffsetWrite(partition, secondOffset)
+    assertTrue(group.hasOffsets)
+    assertEquals(Some(secondOffset), group.offset(partition))
+  }
+
+  @Test
+  def testOffsetCommitWithAnotherPending(): Unit = {
+    val partition = new TopicPartition("foo", 0)
+    val firstOffset = OffsetAndMetadata(37)
+    val secondOffset = OffsetAndMetadata(57)
+
+    group.prepareOffsetCommit(Map(partition -> firstOffset))
+    assertTrue(group.hasOffsets)
+    assertEquals(None, group.offset(partition))
+
+    group.prepareOffsetCommit(Map(partition -> secondOffset))
+    assertTrue(group.hasOffsets)
+
+    group.completePendingOffsetWrite(partition, firstOffset)
+    assertTrue(group.hasOffsets)
+    assertEquals(Some(firstOffset), group.offset(partition))
+
+    group.completePendingOffsetWrite(partition, secondOffset)
+    assertTrue(group.hasOffsets)
+    assertEquals(Some(secondOffset), group.offset(partition))
+  }
+
+  private def assertState(group: GroupMetadata, targetState: GroupState) {
+    val states: Set[GroupState] = Set(Stable, PreparingRebalance, AwaitingSync, Dead)
+    val otherStates = states - targetState
+    otherStates.foreach { otherState =>
+      assertFalse(group.is(otherState))
+    }
+    assertTrue(group.is(targetState))
+  }
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
new file mode 100644
index 0000000..48c9270
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
@@ -0,0 +1,88 @@
+/**
+ * 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 kafka.coordinator.group
+
+import java.util.Arrays
+
+import org.junit.Assert._
+import org.junit.Test
+import org.scalatest.junit.JUnitSuite
+
+class MemberMetadataTest extends JUnitSuite {
+  val groupId = "groupId"
+  val clientId = "clientId"
+  val clientHost = "clientHost"
+  val memberId = "memberId"
+  val protocolType = "consumer"
+  val rebalanceTimeoutMs = 60000
+  val sessionTimeoutMs = 10000
+
+
+  @Test
+  def testMatchesSupportedProtocols() {
+    val protocols = List(("range", Array.empty[Byte]))
+
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    assertTrue(member.matches(protocols))
+    assertFalse(member.matches(List(("range", Array[Byte](0)))))
+    assertFalse(member.matches(List(("roundrobin", Array.empty[Byte]))))
+    assertFalse(member.matches(List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))))
+  }
+
+  @Test
+  def testVoteForPreferredProtocol() {
+    val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
+
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    assertEquals("range", member.vote(Set("range", "roundrobin")))
+    assertEquals("roundrobin", member.vote(Set("blah", "roundrobin")))
+  }
+
+  @Test
+  def testMetadata() {
+    val protocols = List(("range", Array[Byte](0)), ("roundrobin", Array[Byte](1)))
+
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    assertTrue(Arrays.equals(Array[Byte](0), member.metadata("range")))
+    assertTrue(Arrays.equals(Array[Byte](1), member.metadata("roundrobin")))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testMetadataRaisesOnUnsupportedProtocol() {
+    val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
+
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    member.metadata("blah")
+    fail()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testVoteRaisesOnNoSupportedProtocols() {
+    val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
+
+    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    member.vote(Set("blah"))
+    fail()
+  }
+
+
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
new file mode 100644
index 0000000..85c631c
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
@@ -0,0 +1,104 @@
+/*
+ * 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 kafka.coordinator.transaction
+
+import kafka.common.KafkaException
+import kafka.utils.ZkUtils
+import org.easymock.{Capture, EasyMock, IAnswer}
+import org.junit.{After, Test}
+import org.junit.Assert._
+
+class ProducerIdManagerTest {
+
+  val zkUtils: ZkUtils = EasyMock.createNiceMock(classOf[ZkUtils])
+
+  @After
+  def tearDown(): Unit = {
+    EasyMock.reset(zkUtils)
+  }
+
+  @Test
+  def testGetPID() {
+    var zkVersion: Int = -1
+    var data: String = null
+    EasyMock.expect(zkUtils.readDataAndVersionMaybeNull(EasyMock.anyString()))
+      .andAnswer(new IAnswer[(Option[String], Int)] {
+        override def answer(): (Option[String], Int) = {
+          if (zkVersion == -1) {
+            (None.asInstanceOf[Option[String]], 0)
+          } else {
+            (Some(data), zkVersion)
+          }
+        }
+      })
+      .anyTimes()
+
+    val capturedVersion: Capture[Int] = EasyMock.newCapture()
+    val capturedData: Capture[String] = EasyMock.newCapture()
+    EasyMock.expect(zkUtils.conditionalUpdatePersistentPath(EasyMock.anyString(),
+      EasyMock.capture(capturedData),
+      EasyMock.capture(capturedVersion),
+      EasyMock.anyObject().asInstanceOf[Option[(ZkUtils, String, String) => (Boolean,Int)]]))
+      .andAnswer(new IAnswer[(Boolean, Int)] {
+        override def answer(): (Boolean, Int) = {
+          zkVersion = capturedVersion.getValue + 1
+          data = capturedData.getValue
+
+          (true, zkVersion)
+        }
+      })
+      .anyTimes()
+
+    EasyMock.replay(zkUtils)
+
+    val manager1: ProducerIdManager = new ProducerIdManager(0, zkUtils)
+    val manager2: ProducerIdManager = new ProducerIdManager(1, zkUtils)
+
+    val pid1 = manager1.nextPid()
+    val pid2 = manager2.nextPid()
+
+    assertEquals(0, pid1)
+    assertEquals(ProducerIdManager.PidBlockSize, pid2)
+
+    for (i <- 1 until ProducerIdManager.PidBlockSize.asInstanceOf[Int]) {
+      assertEquals(pid1 + i, manager1.nextPid())
+    }
+
+    for (i <- 1 until ProducerIdManager.PidBlockSize.asInstanceOf[Int]) {
+      assertEquals(pid2 + i, manager2.nextPid())
+    }
+
+    assertEquals(pid2 + ProducerIdManager.PidBlockSize, manager1.nextPid())
+    assertEquals(pid2 + ProducerIdManager.PidBlockSize * 2, manager2.nextPid())
+  }
+
+  @Test(expected = classOf[KafkaException])
+  def testExceedPIDLimit() {
+    EasyMock.expect(zkUtils.readDataAndVersionMaybeNull(EasyMock.anyString()))
+      .andAnswer(new IAnswer[(Option[String], Int)] {
+        override def answer(): (Option[String], Int) = {
+          (Some(ProducerIdManager.generatePidBlockJson(ProducerIdBlock(0,
+            Long.MaxValue - ProducerIdManager.PidBlockSize,
+            Long.MaxValue))), 0)
+        }
+      })
+      .anyTimes()
+    EasyMock.replay(zkUtils)
+    new ProducerIdManager(0, zkUtils)
+  }
+}
+

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorIntegrationTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorIntegrationTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorIntegrationTest.scala
new file mode 100644
index 0000000..80ef687
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorIntegrationTest.scala
@@ -0,0 +1,98 @@
+/*
+ * 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 unit.kafka.coordinator.transaction
+
+import java.util.Properties
+
+import kafka.common.Topic
+import kafka.coordinator.transaction.InitPidResult
+import kafka.integration.KafkaServerTestHarness
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.CompressionType
+import org.apache.kafka.common.requests.TransactionResult
+import org.apache.kafka.common.utils.Utils
+import org.junit.{Assert, Test}
+
+class TransactionCoordinatorIntegrationTest extends KafkaServerTestHarness {
+  val offsetsTopicCompressionCodec = CompressionType.GZIP
+  val overridingProps = new Properties()
+  overridingProps.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+  overridingProps.put(KafkaConfig.TransactionsTopicPartitionsProp, "1")
+  overridingProps.put(KafkaConfig.RequestTimeoutMsProp, "100")
+
+  override def generateConfigs = TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false).map {
+    KafkaConfig.fromProps(_, overridingProps)
+  }
+
+  @Test
+  def shouldCommitTransaction(): Unit = {
+    TestUtils.createTopic(zkUtils, Topic.TransactionStateTopicName, 1, 1, servers, servers.head.groupCoordinator.offsetsTopicConfigs)
+    val topic = "foo"
+    TestUtils.createTopic(this.zkUtils, topic, 1, 1, servers)
+
+    val tc = servers.head.transactionCoordinator
+
+    var initPidResult: InitPidResult = null
+    def callback(result: InitPidResult): Unit = {
+      initPidResult = result
+    }
+
+    val txnId = "txn"
+    tc.handleInitPid(txnId, 900000, callback)
+
+    @volatile var addPartitionErrors: Errors = null
+    def addPartitionsCallback(errors: Errors): Unit = {
+        addPartitionErrors = errors
+    }
+
+    tc.handleAddPartitionsToTransaction(txnId,
+      initPidResult.pid,
+      initPidResult.epoch,
+      Set[TopicPartition](new TopicPartition(topic, 0)),
+      addPartitionsCallback
+    )
+
+    while(addPartitionErrors == null) {
+      Utils.sleep(1)
+    }
+
+    Assert.assertEquals(Errors.NONE, addPartitionErrors)
+
+    /**
+     * TODO: Can't do this until the Broker side changes are done
+    @volatile var commitErrors: Errors = null
+    def commitCallback(errors: Errors): Unit ={
+      commitErrors = errors
+    }
+
+    tc.handleEndTransaction(txnId,
+      initPidResult.pid,
+      initPidResult.epoch,
+      TransactionResult.COMMIT,
+      commitCallback)
+
+    while(commitErrors == null) {
+      Utils.sleep(1)
+    }
+
+    Assert.assertEquals(Errors.NONE, commitErrors)
+    */
+  }
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
new file mode 100644
index 0000000..4d8d1c9
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
@@ -0,0 +1,787 @@
+/*
+ * 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 kafka.coordinator.transaction
+
+import kafka.server.DelayedOperationPurgatory
+import kafka.utils.timer.MockTimer
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.TransactionResult
+import org.apache.kafka.common.utils.MockTime
+import org.easymock.{Capture, EasyMock, IAnswer}
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.collection.mutable
+
+class TransactionCoordinatorTest {
+
+  val time = new MockTime()
+
+  var nextPid: Long = 0L
+  val pidManager: ProducerIdManager = EasyMock.createNiceMock(classOf[ProducerIdManager])
+  val transactionManager: TransactionStateManager = EasyMock.createNiceMock(classOf[TransactionStateManager])
+  val transactionMarkerChannelManager: TransactionMarkerChannelManager = EasyMock.createNiceMock(classOf[TransactionMarkerChannelManager])
+  val capturedTxn: Capture[TransactionMetadata] = EasyMock.newCapture()
+  val capturedErrorsCallback: Capture[Errors => Unit] = EasyMock.newCapture()
+  val brokerId = 0
+  private val transactionalId = "known"
+  private val pid = 10
+  private val epoch:Short = 1
+  private val txnTimeoutMs = 1
+
+  private val txnMarkerPurgatory = new DelayedOperationPurgatory[DelayedTxnMarker]("test", new MockTimer, reaperEnabled = false)
+  private val partitions = mutable.Set[TopicPartition](new TopicPartition("topic1", 0))
+
+  val coordinator: TransactionCoordinator = new TransactionCoordinator(brokerId,
+    pidManager,
+    transactionManager,
+    transactionMarkerChannelManager,
+    txnMarkerPurgatory,
+    time)
+
+  var result: InitPidResult = _
+  var error: Errors = Errors.NONE
+
+  private def mockPidManager(): Unit = {
+    EasyMock.expect(pidManager.nextPid())
+      .andAnswer(new IAnswer[Long] {
+        override def answer(): Long = {
+          nextPid += 1
+          nextPid - 1
+        }
+      })
+      .anyTimes()
+  }
+
+  private def initPidGenericMocks(transactionalId: String): Unit = {
+    mockPidManager()
+    EasyMock.expect(transactionManager.isCoordinatorFor(EasyMock.eq(transactionalId)))
+      .andReturn(true)
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.isCoordinatorLoadingInProgress(EasyMock.anyString()))
+      .andReturn(false)
+      .anyTimes()
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true)
+      .anyTimes()
+  }
+
+
+  @Test
+  def shouldAcceptInitPidAndReturnNextPidWhenTransactionalIdIsEmpty(): Unit = {
+    mockPidManager()
+    EasyMock.replay(pidManager)
+
+    coordinator.handleInitPid("", txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(0L, 0, Errors.NONE), result)
+    coordinator.handleInitPid("", txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(1L, 0, Errors.NONE), result)
+  }
+
+  @Test
+  def shouldAcceptInitPidAndReturnNextPidWhenTransactionalIdIsNull(): Unit = {
+    mockPidManager()
+    EasyMock.replay(pidManager)
+
+    coordinator.handleInitPid(null, txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(0L, 0, Errors.NONE), result)
+    coordinator.handleInitPid(null, txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(1L, 0, Errors.NONE), result)
+  }
+
+  @Test
+  def shouldInitPidWithEpochZeroForNewTransactionalId(): Unit = {
+    initPidGenericMocks(transactionalId)
+    EasyMock.expect(transactionManager.addTransaction(EasyMock.eq(transactionalId), EasyMock.capture(capturedTxn)))
+      .andAnswer(new IAnswer[TransactionMetadata] {
+        override def answer(): TransactionMetadata = {
+          capturedTxn.getValue
+        }
+      })
+      .once()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andAnswer(new IAnswer[Option[TransactionMetadata]] {
+        override def answer(): Option[TransactionMetadata] = {
+          if (capturedTxn.hasCaptured) {
+            Some(capturedTxn.getValue)
+          } else {
+            None
+          }
+        }
+      })
+      .once()
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.capture(capturedTxn),
+      EasyMock.capture(capturedErrorsCallback)))
+      .andAnswer(new IAnswer[Unit] {
+        override def answer(): Unit = {
+          capturedErrorsCallback.getValue.apply(Errors.NONE)
+        }
+      })
+      .anyTimes()
+    EasyMock.replay(pidManager, transactionManager)
+
+    coordinator.handleInitPid(transactionalId, txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(0L, 0, Errors.NONE), result)
+  }
+
+  @Test
+  def shouldRespondWithNotCoordinatorOnInitPidWhenNotCoordinatorForId(): Unit = {
+    mockPidManager()
+    EasyMock.replay(pidManager)
+    coordinator.handleInitPid("some-pid", txnTimeoutMs, initPidMockCallback)
+    assertEquals(InitPidResult(-1, -1, Errors.NOT_COORDINATOR), result)
+  }
+
+  @Test
+  def shouldRespondWithInvalidPidMappingOnAddPartitionsToTransactionWhenTransactionalIdNotPresent(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(None)
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 1, partitions, errorsCallback)
+    assertEquals(Errors.INVALID_PID_MAPPING, error)
+  }
+
+  @Test
+  def shouldRespondWithInvalidRequestAddPartitionsToTransactionWhenTransactionalIdIsEmpty(): Unit = {
+    coordinator.handleAddPartitionsToTransaction("", 0L, 1, partitions, errorsCallback)
+    assertEquals(Errors.INVALID_REQUEST, error)
+  }
+
+  @Test
+  def shouldRespondWithInvalidRequestAddPartitionsToTransactionWhenTransactionalIdIsNull(): Unit = {
+    coordinator.handleAddPartitionsToTransaction(null, 0L, 1, partitions, errorsCallback)
+    assertEquals(Errors.INVALID_REQUEST, error)
+  }
+
+  @Test
+  def shouldRespondWithNotCoordinatorOnAddPartitionsWhenNotCoordinator(): Unit = {
+    coordinator.handleAddPartitionsToTransaction("txn", 0L, 1, partitions, errorsCallback)
+    assertEquals(Errors.NOT_COORDINATOR, error)
+  }
+
+  @Test
+  def shouldRespondWithCoordinatorLoadInProgressOnAddPartitionsWhenCoordintorLoading(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.isCoordinatorLoadingInProgress(transactionalId))
+    .andReturn(true)
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 1, partitions, errorsCallback)
+    assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, error)
+  }
+
+  @Test
+  def shouldRespondWithInvalidTnxStateOnAddPartitionsWhenStateIsPrepareCommit(): Unit = {
+    validateInvalidTxnState(PrepareCommit)
+  }
+
+  @Test
+  def shouldRespondWithInvalidTnxStateOnAddPartitionsWhenStateIsPrepareAbort(): Unit = {
+    validateInvalidTxnState(PrepareAbort)
+  }
+
+  @Test
+  def shouldRespondWithInvalidTnxStateOnAddPartitionsWhenStateIsCompleteCommit(): Unit = {
+    validateInvalidTxnState(CompleteCommit)
+  }
+
+  @Test
+  def shouldRespondWithInvalidTnxStateOnAddPartitionsWhenStateIsCompleteAbort(): Unit = {
+    validateInvalidTxnState(CompleteAbort)
+  }
+
+  def validateInvalidTxnState(state: TransactionState): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(0, 0, 0, state, mutable.Set.empty, 0, 0)))
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 0, partitions, errorsCallback)
+    assertEquals(Errors.INVALID_TXN_STATE, error)
+  }
+
+  @Test
+  def shouldRespondWithInvalidTnxProduceEpochOnAddPartitionsWhenEpochsAreDifferent(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(0, 10, 0, PrepareCommit, mutable.Set.empty, 0, 0)))
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 0, partitions, errorsCallback)
+    assertEquals(Errors.INVALID_PRODUCER_EPOCH, error)
+  }
+
+  @Test
+  def shouldAppendNewMetadataToLogOnAddPartitionsWhenPartitionsAdded(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(0, 0, 0, Empty, mutable.Set.empty, 0, 0)))
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(new TransactionMetadata(0, 0, 0, Ongoing, partitions, time.milliseconds(), time.milliseconds())),
+      EasyMock.capture(capturedErrorsCallback)
+    ))
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 0, partitions, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldRespondWithErrorsNoneOnAddPartitionWhenNoErrorsAndPartitionsTheSame(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(0, 0, 0, Empty, partitions, 0, 0)))
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 0, partitions, errorsCallback)
+    assertEquals(Errors.NONE, error)
+    EasyMock.verify(transactionManager)
+
+  }
+
+  @Test
+  def shouldReplyWithInvalidPidMappingOnEndTxnWhenTxnIdDoesntExist(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId)).andReturn(None)
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_PID_MAPPING, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReplyWithInvalidPidMappingOnEndTxnWhenPidDosentMatchMapped(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(10, 0, 0, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_PID_MAPPING, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReplyWithProducerFencedOnEndTxnWhenEpochIsNotSameAsTransaction(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_PRODUCER_EPOCH, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnOkOnEndTxnWhenStatusIsCompleteCommitAndResultIsCommit(): Unit ={
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.NONE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnOkOnEndTxnWhenStatusIsCompleteAbortAndResultIsAbort(): Unit ={
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.ABORT, errorsCallback)
+    assertEquals(Errors.NONE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsCompleteAbortAndResultIsNotAbort(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_TXN_STATE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsCompleteCommitAndResultIsNotCommit(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.ABORT, errorsCallback)
+    assertEquals(Errors.INVALID_TXN_STATE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsPrepareCommit(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, PrepareCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_TXN_STATE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsPrepareAbort(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(new TransactionMetadata(pid, 1, 1, PrepareAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, 1, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_TXN_STATE, error)
+    EasyMock.verify(transactionManager)
+  }
+
+
+  @Test
+  def shouldAppendPrepareCommitToLogOnEndTxnWhenStatusIsOngoingAndResultIsCommit(): Unit = {
+    mockPrepare(PrepareCommit)
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldAppendPrepareAbortToLogOnEndTxnWhenStatusIsOngoingAndResultIsAbort(): Unit = {
+    mockPrepare(PrepareAbort)
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.ABORT, errorsCallback)
+    EasyMock.verify(transactionManager)
+  }
+
+
+  @Test
+  def shouldAppendCompleteAbortToLogOnEndTxnWhenStatusIsOngoingAndResultIsAbort(): Unit = {
+    mockComplete(PrepareAbort)
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.ABORT, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldAppendCompleteCommitToLogOnEndTxnWhenStatusIsOngoingAndResultIsCommit(): Unit = {
+    mockComplete(PrepareCommit)
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldRespondWithInvalidRequestOnEndTxnWhenTransactionalIdIsNull(): Unit = {
+    coordinator.handleEndTransaction(null, 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_REQUEST, error)
+  }
+
+  @Test
+  def shouldRespondWithInvalidRequestOnEndTxnWhenTransactionalIdIsEmpty(): Unit = {
+    coordinator.handleEndTransaction("", 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.INVALID_REQUEST, error)
+  }
+
+  @Test
+  def shouldRespondWithNotCoordinatorOnEndTxnWhenIsNotCoordinatorForId(): Unit = {
+    coordinator.handleEndTransaction("id", 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.NOT_COORDINATOR, error)
+  }
+
+  @Test
+  def shouldRespondWithCoordinatorLoadInProgressOnEndTxnWhenCoordinatorIsLoading(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(EasyMock.anyString()))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.isCoordinatorLoadingInProgress(EasyMock.anyString()))
+      .andReturn(true)
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction("id", 0, 0, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, error)
+  }
+
+  @Test
+  def shouldIncrementEpochAndUpdateMetadataOnHandleInitPidWhenExistingEmptyTransaction(): Unit = {
+    validateIncrementEpochAndUpdateMetadata(Empty)
+  }
+
+  @Test
+  def shouldIncrementEpochAndUpdateMetadataOnHandleInitPidWhenExistingCompleteTransaction(): Unit = {
+    validateIncrementEpochAndUpdateMetadata(CompleteAbort)
+  }
+
+  @Test
+  def shouldIncrementEpochAndUpdateMetadataOnHandleInitPidWhenExistingCompleteCommitTransaction(): Unit = {
+    validateIncrementEpochAndUpdateMetadata(CompleteCommit)
+  }
+
+  @Test
+  def shouldWaitForCommitToCompleteOnHandleInitPidAndExistingTransactionInPrepareCommitState(): Unit ={
+    validateWaitsForCompletionBeforeRespondingWithIncrementedEpoch(PrepareCommit)
+  }
+
+  @Test
+  def shouldWaitForCommitToCompleteOnHandleInitPidAndExistingTransactionInPrepareAbortState(): Unit ={
+    validateWaitsForCompletionBeforeRespondingWithIncrementedEpoch(PrepareAbort)
+  }
+
+  @Test
+  def shouldAbortTransactionOnHandleInitPidWhenExistingTransactionInOngoingState(): Unit = {
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true)
+
+    val metadata = new TransactionMetadata(pid, epoch, txnTimeoutMs, Ongoing, mutable.Set[TopicPartition](new TopicPartition("topic", 1)), 0, 0)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(metadata))
+      .once()
+
+    mockComplete(PrepareAbort)
+
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+
+    val completedMetadata = new TransactionMetadata(pid, epoch, txnTimeoutMs, CompleteAbort, mutable.Set.empty[TopicPartition], 0, 0)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(completedMetadata))
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.anyObject(classOf[TransactionMetadata]),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }
+    })
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleInitPid(transactionalId, txnTimeoutMs, initPidMockCallback)
+
+    assertEquals(InitPidResult(10, 2, Errors.NONE), result)
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldRemoveTransactionsForPartitionOnEmigration(): Unit = {
+    EasyMock.expect(transactionManager.removeTransactionsForPartition(0))
+    EasyMock.expect(transactionMarkerChannelManager.removeStateForPartition(0))
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleTxnEmigration(0)
+
+    EasyMock.verify(transactionManager, transactionMarkerChannelManager)
+  }
+
+  @Test
+  def shouldRetryOnCommitWhenTxnMarkerRequestFailsWithErrorOtherThanNotCoordinator(): Unit = {
+    val prepareMetadata = mockPrepare(PrepareCommit, runCallback = true)
+
+    EasyMock.expect(transactionManager.coordinatorEpochFor(transactionalId))
+      .andReturn(Some(0))
+
+    EasyMock.expect(transactionMarkerChannelManager.addTxnMarkerRequest(
+      EasyMock.eq(0),
+      EasyMock.anyObject(),
+      EasyMock.anyInt(),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NETWORK_EXCEPTION)
+      }
+    }).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }
+    })
+
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(prepareMetadata))
+      .once()
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionMarkerChannelManager)
+  }
+
+  @Test
+  def shouldNotRetryOnCommitWhenTxnMarkerRequestFailsWithNotCoordinator(): Unit = {
+    val prepareMetadata = mockPrepare(PrepareCommit, runCallback = true)
+
+    EasyMock.expect(transactionManager.coordinatorEpochFor(transactionalId))
+      .andReturn(Some(0))
+
+    EasyMock.expect(transactionMarkerChannelManager.addTxnMarkerRequest(
+      EasyMock.eq(0),
+      EasyMock.anyObject(),
+      EasyMock.anyInt(),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NOT_COORDINATOR)
+      }
+    })
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionMarkerChannelManager)
+  }
+
+  @Test
+  def shouldNotRetryOnCommitWhenAppendToLogFailsWithNotCoordinator(): Unit = {
+    mockComplete(PrepareCommit, Errors.NOT_COORDINATOR)
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldRetryOnCommitWhenAppendToLogFailsErrorsOtherThanNotCoordinator(): Unit = {
+    mockComplete(PrepareCommit, Errors.ILLEGAL_GENERATION)
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.handleEndTransaction(transactionalId, pid, epoch, TransactionResult.COMMIT, errorsCallback)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  private def validateWaitsForCompletionBeforeRespondingWithIncrementedEpoch(state: TransactionState) = {
+    val transactionId = "tid"
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionId))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+
+    val metadata = new TransactionMetadata(0, 0, 0, state, mutable.Set[TopicPartition](new TopicPartition("topic", 1)), 0, 0)
+    EasyMock.expect(transactionManager.getTransactionState(transactionId))
+      .andReturn(Some(metadata)).anyTimes()
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionId),
+      EasyMock.anyObject(classOf[TransactionMetadata]),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }
+    })
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleInitPid(transactionId, 10, initPidMockCallback)
+    // no result yet as hasn't completed
+    assertNull(result)
+    // complete the transaction
+    metadata.topicPartitions.clear()
+    metadata.state = if (state == PrepareCommit) CompleteCommit else CompleteAbort
+    txnMarkerPurgatory.checkAndComplete(0L)
+
+    assertEquals(InitPidResult(0, 1, Errors.NONE), result)
+    assertEquals(new TransactionMetadata(0, 1, 10, Empty, mutable.Set.empty, 0, time.milliseconds()), metadata)
+  }
+
+  private def validateIncrementEpochAndUpdateMetadata(state: TransactionState) = {
+    val transactionId = "tid"
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true)
+
+    val metadata = new TransactionMetadata(0, 0, 0, state, mutable.Set.empty[TopicPartition], 0, 0)
+    EasyMock.expect(transactionManager.getTransactionState(transactionId))
+      .andReturn(Some(metadata))
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionId),
+      EasyMock.anyObject(classOf[TransactionMetadata]),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }
+    })
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleInitPid(transactionId, 10, initPidMockCallback)
+
+    assertEquals(InitPidResult(0, 1, Errors.NONE), result)
+    assertEquals(10, metadata.txnTimeoutMs)
+    assertEquals(time.milliseconds(), metadata.lastUpdateTimestamp)
+    assertEquals(1, metadata.producerEpoch)
+    assertEquals(0, metadata.pid)
+  }
+
+  private def mockPrepare(transactionState: TransactionState, runCallback: Boolean = false) = {
+    val originalMetadata = new TransactionMetadata(pid,
+      epoch,
+      txnTimeoutMs,
+      Ongoing,
+      collection.mutable.Set.empty[TopicPartition],
+      0,
+      time.milliseconds())
+
+    val prepareCommitMetadata = new TransactionMetadata(pid,
+      epoch,
+      txnTimeoutMs,
+      transactionState,
+      collection.mutable.Set.empty[TopicPartition],
+      0,
+      time.milliseconds())
+
+    EasyMock.expect(transactionManager.isCoordinatorFor(transactionalId))
+      .andReturn(true)
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(originalMetadata))
+      .once()
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(EasyMock.eq(transactionalId),
+      EasyMock.eq(prepareCommitMetadata),
+      EasyMock.capture(capturedErrorsCallback)))
+      .andAnswer(new IAnswer[Unit] {
+        override def answer(): Unit = {
+          if (runCallback) capturedErrorsCallback.getValue.apply(Errors.NONE)
+        }
+      }).once()
+    prepareCommitMetadata
+  }
+
+  private def mockComplete(transactionState: TransactionState, appendError: Errors = Errors.NONE) = {
+
+
+    val prepareMetadata: TransactionMetadata = mockPrepare(transactionState, true)
+    val finalState = if (transactionState == PrepareAbort) CompleteAbort else CompleteCommit
+
+    EasyMock.expect(transactionManager.coordinatorEpochFor(transactionalId))
+      .andReturn(Some(0))
+
+    EasyMock.expect(transactionMarkerChannelManager.addTxnMarkerRequest(
+      EasyMock.eq(0),
+      EasyMock.anyObject(),
+      EasyMock.anyInt(),
+      EasyMock.capture(capturedErrorsCallback)
+    )).andAnswer(new IAnswer[Unit] {
+      override def answer(): Unit = {
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }
+    })
+
+    EasyMock.expect(transactionManager.getTransactionState(transactionalId))
+      .andReturn(Some(prepareMetadata))
+      .once()
+
+    val completedMetadata = new TransactionMetadata(pid,
+      epoch,
+      txnTimeoutMs,
+      finalState,
+      prepareMetadata.topicPartitions,
+      prepareMetadata.transactionStartTime,
+      prepareMetadata.lastUpdateTimestamp)
+
+    val firstAnswer = EasyMock.expect(transactionManager.appendTransactionToLog(EasyMock.eq(transactionalId),
+      EasyMock.eq(completedMetadata),
+      EasyMock.capture(capturedErrorsCallback)))
+      .andAnswer(new IAnswer[Unit] {
+        override def answer(): Unit = {
+          capturedErrorsCallback.getValue.apply(appendError)
+        }
+      })
+
+     if(appendError != Errors.NONE && appendError != Errors.NOT_COORDINATOR) {
+        firstAnswer.andAnswer(new IAnswer[Unit] {
+          override def answer(): Unit = {
+            capturedErrorsCallback.getValue.apply(Errors.NONE)
+          }
+        })
+     }
+
+
+  }
+
+
+  def initPidMockCallback(ret: InitPidResult): Unit = {
+    result = ret
+  }
+
+  def errorsCallback(ret: Errors): Unit = {
+    error = ret
+  }
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
new file mode 100644
index 0000000..cfb4a99
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
@@ -0,0 +1,109 @@
+/*
+ * 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 kafka.coordinator.transaction
+
+
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record.{CompressionType, SimpleRecord, MemoryRecords}
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import org.scalatest.junit.JUnitSuite
+
+import scala.collection.JavaConverters._
+
+class TransactionLogTest extends JUnitSuite {
+
+  val epoch: Short = 0
+  val transactionTimeoutMs: Int = 1000
+
+  val topicPartitions: Set[TopicPartition] = Set[TopicPartition](new TopicPartition("topic1", 0),
+    new TopicPartition("topic1", 1),
+    new TopicPartition("topic2", 0),
+    new TopicPartition("topic2", 1),
+    new TopicPartition("topic2", 2))
+
+  @Test
+  def shouldThrowExceptionWriteInvalidTxn() {
+    val txnMetadata = TransactionMetadata(0L, epoch, transactionTimeoutMs, 0)
+    txnMetadata.addPartitions(topicPartitions)
+
+    intercept[IllegalStateException] {
+      TransactionLog.valueToBytes(txnMetadata)
+    }
+  }
+
+  @Test
+  def shouldReadWriteMessages() {
+    val pidMappings = Map[String, Long]("zero" -> 0L,
+      "one" -> 1L,
+      "two" -> 2L,
+      "three" -> 3L,
+      "four" -> 4L,
+      "five" -> 5L)
+
+    val transactionStates = Map[Long, TransactionState](0L -> Empty,
+      1L -> Ongoing,
+      2L -> PrepareCommit,
+      3L -> CompleteCommit,
+      4L -> PrepareAbort,
+      5L -> CompleteAbort)
+
+    // generate transaction log messages
+    val txnRecords = pidMappings.map { case (transactionalId, pid) =>
+      val txnMetadata = TransactionMetadata(pid, epoch, transactionTimeoutMs, transactionStates(pid), 0)
+
+      if (!txnMetadata.state.equals(Empty))
+        txnMetadata.addPartitions(topicPartitions)
+
+      val keyBytes = TransactionLog.keyToBytes(transactionalId)
+      val valueBytes = TransactionLog.valueToBytes(txnMetadata)
+
+      new SimpleRecord(keyBytes, valueBytes)
+    }.toSeq
+
+    val records = MemoryRecords.withRecords(0, CompressionType.NONE, txnRecords: _*)
+
+    var count = 0
+    for (record <- records.records.asScala) {
+      val key = TransactionLog.readMessageKey(record.key())
+
+      key match {
+        case pidKey: TxnKey =>
+          val transactionalId = pidKey.transactionalId
+          val txnMetadata = TransactionLog.readMessageValue(record.value())
+
+          assertEquals(pidMappings(transactionalId), txnMetadata.pid)
+          assertEquals(epoch, txnMetadata.producerEpoch)
+          assertEquals(transactionTimeoutMs, txnMetadata.txnTimeoutMs)
+          assertEquals(transactionStates(txnMetadata.pid), txnMetadata.state)
+
+          if (txnMetadata.state.equals(Empty))
+            assertEquals(Set.empty[TopicPartition], txnMetadata.topicPartitions)
+          else
+            assertEquals(topicPartitions, txnMetadata.topicPartitions)
+
+          count = count + 1
+
+        case _ => fail(s"Unexpected transaction topic message key $key")
+      }
+    }
+
+    assertEquals(pidMappings.size, count)
+  }
+
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
new file mode 100644
index 0000000..352daa2
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
@@ -0,0 +1,283 @@
+/*
+ * 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 kafka.coordinator.transaction
+
+import kafka.api.{LeaderAndIsr, PartitionStateInfo}
+import kafka.common.{BrokerEndPointNotAvailableException, BrokerNotAvailableException, InterBrokerSendThread}
+import kafka.controller.LeaderIsrAndControllerEpoch
+import kafka.server.{DelayedOperationPurgatory, KafkaConfig, MetadataCache}
+import kafka.utils.{MockTime, TestUtils}
+import kafka.utils.timer.MockTimer
+import org.apache.kafka.clients.NetworkClient
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.{Errors, SecurityProtocol}
+import org.apache.kafka.common.requests.{TransactionResult, WriteTxnMarkersRequest}
+import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.{Node, TopicPartition}
+import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.collection.mutable
+
+class TransactionMarkerChannelManagerTest {
+  private val metadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+  private val interBrokerSendThread = EasyMock.createNiceMock(classOf[InterBrokerSendThread])
+  private val networkClient = EasyMock.createNiceMock(classOf[NetworkClient])
+  private val channel = new TransactionMarkerChannel(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+    metadataCache,
+    networkClient,
+    new MockTime())
+  private val purgatory = new DelayedOperationPurgatory[DelayedTxnMarker]("txn-purgatory-name",
+    new MockTimer,
+    reaperEnabled = false)
+  private val requestGenerator = TransactionMarkerChannelManager.requestGenerator(channel, purgatory)
+  private val partition1 = new TopicPartition("topic1", 0)
+  private val partition2 = new TopicPartition("topic1", 1)
+  private val broker1 = new Node(1, "host", 10)
+  private val broker2 = new Node(2, "otherhost", 10)
+  private val metadataPartition = 0
+  private val channelManager = new TransactionMarkerChannelManager(
+    KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:2181")),
+    metadataCache,
+    purgatory,
+    interBrokerSendThread,
+    channel)
+
+  @Test
+  def shouldGenerateEmptyMapWhenNoRequestsOutstanding(): Unit = {
+    assertTrue(requestGenerator().isEmpty)
+  }
+
+  @Test
+  def shouldGenerateRequestPerBroker(): Unit ={
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition2.topic(), partition2.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(2, 0, List.empty, 0), 0), Set.empty)))
+
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(2), EasyMock.anyObject())).andReturn(Some(broker2))
+
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1, partition2))
+
+
+    val expectedBroker1Request = new WriteTxnMarkersRequest.Builder(0,
+      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList(partition1)))).build()
+    val expectedBroker2Request = new WriteTxnMarkersRequest.Builder(0,
+      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList(partition2)))).build()
+
+    val requests: Map[Node, WriteTxnMarkersRequest] = requestGenerator().map{ result =>
+      (result.destination, result.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build())
+    }.toMap
+
+    val broker1Request = requests(broker1)
+    val broker2Request = requests(broker2)
+
+    assertEquals(2, requests.size)
+    assertEquals(expectedBroker1Request, broker1Request)
+    assertEquals(expectedBroker2Request, broker2Request)
+
+  }
+
+  @Test
+  def shouldGenerateRequestPerPartitionPerBroker(): Unit ={
+    val partitionOneEpoch = 0
+    val partitionTwoEpoch = 1
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, partitionOneEpoch, List.empty, 0), 0), Set.empty)))
+
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition2.topic(), partition2.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, partitionTwoEpoch, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1)).anyTimes()
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, partitionOneEpoch, Set[TopicPartition](partition1))
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, partitionTwoEpoch, Set[TopicPartition](partition2))
+
+    val expectedPartition1Request = new WriteTxnMarkersRequest.Builder(0,
+      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList(partition1)))).build()
+    val expectedPartition2Request = new WriteTxnMarkersRequest.Builder(1,
+      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList(partition2)))).build()
+
+    val requests = requestGenerator().map { result =>
+      val markersRequest = result.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build()
+      (markersRequest.coordinatorEpoch(), (result.destination, markersRequest))
+    }.toMap
+
+    val request1 = requests(partitionOneEpoch)
+    val request2 = requests(partitionTwoEpoch)
+    assertEquals(broker1, request1._1)
+    assertEquals(broker1, request2._1)
+    assertEquals(2, requests.size)
+    assertEquals(expectedPartition1Request, request1._2)
+    assertEquals(expectedPartition2Request, request2._2)
+  }
+
+  @Test
+  def shouldDrainBrokerQueueWhenGeneratingRequests(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    assertTrue(requestGenerator().nonEmpty)
+    assertTrue(requestGenerator().isEmpty)
+  }
+
+  @Test
+  def shouldRetryGettingLeaderWhenNotFound(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(None)
+      .andReturn(None)
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    EasyMock.verify(metadataCache)
+  }
+
+  @Test
+  def shouldRetryGettingLeaderWhenBrokerEndPointNotAvailableException(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+      .times(2)
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject()))
+      .andThrow(new BrokerEndPointNotAvailableException())
+      .andReturn(Some(broker1))
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    EasyMock.verify(metadataCache)
+  }
+
+  @Test
+  def shouldRetryGettingLeaderWhenLeaderDoesntExist(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+      .times(2)
+
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject()))
+      .andReturn(None)
+      .andReturn(Some(broker1))
+
+    EasyMock.replay(metadataCache)
+
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    EasyMock.verify(metadataCache)
+  }
+
+  @Test
+  def shouldAddPendingTxnRequest(): Unit = {
+    val metadata = new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1, partition2), 0, 0L)
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getPartitionInfo(partition2.topic(), partition2.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(2, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(2), EasyMock.anyObject())).andReturn(Some(broker2))
+
+    EasyMock.replay(metadataCache)
+
+    channelManager.addTxnMarkerRequest(metadataPartition, metadata, 0, completionCallback)
+
+    assertEquals(Some(metadata), channel.pendingTxnMetadata(metadataPartition, 1))
+
+  }
+
+  @Test
+  def shouldAddRequestToBrokerQueue(): Unit = {
+    val metadata = new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1), 0, 0L)
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+    EasyMock.replay(metadataCache)
+
+    channelManager.addTxnMarkerRequest(metadataPartition, metadata, 0, completionCallback)
+    assertEquals(1, requestGenerator().size)
+  }
+
+  @Test
+  def shouldAddDelayedTxnMarkerToPurgatory(): Unit = {
+    val metadata = new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1), 0, 0L)
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(EasyMock.eq(1), EasyMock.anyObject())).andReturn(Some(broker1))
+
+    EasyMock.replay(metadataCache)
+
+    channelManager.addTxnMarkerRequest(metadataPartition, metadata, 0, completionCallback)
+    assertEquals(1,purgatory.watched)
+  }
+
+  @Test
+  def shouldStartInterBrokerThreadOnStartup(): Unit = {
+    EasyMock.expect(interBrokerSendThread.start())
+    EasyMock.replay(interBrokerSendThread)
+    channelManager.start()
+    EasyMock.verify(interBrokerSendThread)
+  }
+
+
+  @Test
+  def shouldStopInterBrokerThreadOnShutdown(): Unit = {
+    EasyMock.expect(interBrokerSendThread.shutdown())
+    EasyMock.replay(interBrokerSendThread)
+    channelManager.shutdown()
+    EasyMock.verify(interBrokerSendThread)
+  }
+
+  @Test
+  def shouldClearPurgatoryForPartitionWhenPartitionEmigrated(): Unit = {
+    val metadata1 = new TransactionMetadata(0, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1), 0, 0)
+    purgatory.tryCompleteElseWatch(new DelayedTxnMarker(metadata1, (error:Errors) => {}),Seq(0L))
+    channel.maybeAddPendingRequest(0, metadata1)
+
+    val metadata2 = new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1), 0, 0)
+    purgatory.tryCompleteElseWatch(new DelayedTxnMarker(metadata2, (error:Errors) => {}),Seq(1L))
+    channel.maybeAddPendingRequest(0, metadata2)
+
+    val metadata3 = new TransactionMetadata(2, 0, 0, PrepareCommit, mutable.Set[TopicPartition](partition1), 0, 0)
+    purgatory.tryCompleteElseWatch(new DelayedTxnMarker(metadata3, (error:Errors) => {}),Seq(2L))
+    channel.maybeAddPendingRequest(1, metadata3)
+
+    channelManager.removeStateForPartition(0)
+
+    assertEquals(1, purgatory.watched)
+    // should not complete as they've been removed
+    purgatory.checkAndComplete(0L)
+    purgatory.checkAndComplete(1L)
+    
+    assertEquals(1, purgatory.watched)
+  }
+
+  def completionCallback(errors: Errors): Unit = {
+  }
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/f69d9415/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelTest.scala
----------------------------------------------------------------------
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelTest.scala
new file mode 100644
index 0000000..6bfeb9b
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelTest.scala
@@ -0,0 +1,179 @@
+/*
+ * 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 kafka.coordinator.transaction
+
+import kafka.api.{LeaderAndIsr, PartitionStateInfo}
+import kafka.controller.LeaderIsrAndControllerEpoch
+import kafka.server.{DelayedOperationPurgatory, MetadataCache}
+import kafka.utils.MockTime
+import kafka.utils.timer.MockTimer
+import org.apache.kafka.clients.NetworkClient
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.{Errors, SecurityProtocol}
+import org.apache.kafka.common.requests.{TransactionResult, WriteTxnMarkersRequest}
+import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.{Node, TopicPartition}
+import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.collection.mutable
+
+class TransactionMarkerChannelTest {
+
+  private val metadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+  private val networkClient = EasyMock.createNiceMock(classOf[NetworkClient])
+  private val purgatory = new DelayedOperationPurgatory[DelayedTxnMarker]("name", new MockTimer, reaperEnabled = false)
+  private val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+  private val channel = new TransactionMarkerChannel(listenerName, metadataCache, networkClient, new MockTime())
+  private val partition1 = new TopicPartition("topic1", 0)
+
+
+  @Test
+  def shouldAddEmptyBrokerQueueWhenAddingNewBroker(): Unit = {
+    channel.addOrUpdateBroker(new Node(1, "host", 10))
+    channel.addOrUpdateBroker(new Node(2, "host", 10))
+    assertEquals(0, channel.queueForBroker(1).get.markersQueue.size())
+    assertEquals(0, channel.queueForBroker(2).get.markersQueue.size())
+  }
+
+  @Test
+  def shouldUpdateDestinationBrokerNodeWhenUpdatingBroker(): Unit = {
+    val newDestination = new Node(1, "otherhost", 100)
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+
+    // getAliveEndpoint returns an updated node
+    EasyMock.expect(metadataCache.getAliveEndpoint(1, listenerName)).andReturn(Some(newDestination))
+    EasyMock.replay(metadataCache)
+
+    channel.addOrUpdateBroker(new Node(1, "host", 10))
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    val destinationAndQueue = channel.queueForBroker(1).get
+    assertEquals(newDestination, destinationAndQueue.destBrokerNode)
+    assertEquals(1, destinationAndQueue.markersQueue.size())
+  }
+
+  @Test
+  def shouldQueueRequestsByBrokerId(): Unit = {
+    channel.addOrUpdateBroker(new Node(1, "host", 10))
+    channel.addOrUpdateBroker(new Node(2, "otherhost", 10))
+    channel.addRequestForBroker(1, CoordinatorEpochAndMarkers(0, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+    channel.addRequestForBroker(1, CoordinatorEpochAndMarkers(0, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+    channel.addRequestForBroker(2, CoordinatorEpochAndMarkers(0, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+
+    assertEquals(2, channel.queueForBroker(1).get.markersQueue.size())
+    assertEquals(1, channel.queueForBroker(2).get.markersQueue.size())
+  }
+
+  @Test
+  def shouldNotAddPendingTxnIfOneAlreadyExistsForPid(): Unit = {
+    channel.maybeAddPendingRequest(0, new TransactionMetadata(0, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0))
+    assertFalse(channel.maybeAddPendingRequest(0, new TransactionMetadata(0, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0)))
+  }
+
+  @Test
+  def shouldAddRequestsToCorrectBrokerQueues(): Unit = {
+    val partition2 = new TopicPartition("topic1", 1)
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+
+    EasyMock.expect(metadataCache.getPartitionInfo(partition2.topic(), partition2.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(2, 0, List.empty, 0), 0), Set.empty)))
+
+    EasyMock.expect(metadataCache.getAliveEndpoint(1, listenerName)).andReturn(Some(new Node(1, "host", 10)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(2, listenerName)).andReturn(Some(new Node(2, "otherhost", 10)))
+
+    EasyMock.replay(metadataCache)
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1, partition2))
+
+    assertEquals(1, channel.queueForBroker(1).get.markersQueue.size)
+    assertEquals(1, channel.queueForBroker(2).get.markersQueue.size)
+  }
+  @Test
+  def shouldWakeupNetworkClientWhenRequestsQueued(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(1, listenerName)).andReturn(Some(new Node(1, "host", 10)))
+
+    EasyMock.expect(networkClient.wakeup())
+
+    EasyMock.replay(metadataCache, networkClient)
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    EasyMock.verify(networkClient)
+  }
+
+  @Test
+  def shouldAddNewBrokerQueueIfDoesntAlreadyExistWhenAddingRequest(): Unit = {
+    EasyMock.expect(metadataCache.getPartitionInfo(partition1.topic(), partition1.partition()))
+      .andReturn(Some(PartitionStateInfo(LeaderIsrAndControllerEpoch(LeaderAndIsr(1, 0, List.empty, 0), 0), Set.empty)))
+    EasyMock.expect(metadataCache.getAliveEndpoint(1, listenerName)).andReturn(Some(new Node(1, "host", 10)))
+
+    EasyMock.replay(metadataCache)
+    channel.addRequestToSend(0, 0, 0, TransactionResult.COMMIT, 0, Set[TopicPartition](partition1))
+
+    assertEquals(1, channel.queueForBroker(1).get.markersQueue.size)
+    EasyMock.verify(metadataCache)
+  }
+
+  @Test
+  def shouldGetPendingTxnMetadataByPid(): Unit = {
+    val metadataPartition = 0
+    val transaction = new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0)
+    channel.maybeAddPendingRequest(metadataPartition, transaction)
+    channel.maybeAddPendingRequest(metadataPartition, new TransactionMetadata(2, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0))
+    assertEquals(Some(transaction), channel.pendingTxnMetadata(metadataPartition, 1))
+  }
+
+  @Test
+  def shouldRemovePendingRequestsForPartitionWhenPartitionEmigrated(): Unit = {
+    channel.maybeAddPendingRequest(0, new TransactionMetadata(0, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0))
+    channel.maybeAddPendingRequest(0, new TransactionMetadata(1, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0))
+    val metadata = new TransactionMetadata(2, 0, 0, PrepareCommit, mutable.Set.empty, 0, 0)
+    channel.maybeAddPendingRequest(1, metadata)
+
+    channel.removeStateForPartition(0)
+
+    assertEquals(None, channel.pendingTxnMetadata(0, 0))
+    assertEquals(None, channel.pendingTxnMetadata(0, 1))
+    assertEquals(Some(metadata), channel.pendingTxnMetadata(1, 2))
+  }
+
+  @Test
+  def shouldRemoveBrokerRequestsForPartitionWhenPartitionEmigrated(): Unit = {
+    channel.addOrUpdateBroker(new Node(1, "host", 10))
+    channel.addRequestForBroker(1, CoordinatorEpochAndMarkers(0, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+    channel.addRequestForBroker(1, CoordinatorEpochAndMarkers(1, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+    channel.addRequestForBroker(1, CoordinatorEpochAndMarkers(1, 0, Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(0, 0, TransactionResult.COMMIT, Utils.mkList()))))
+
+    channel.removeStateForPartition(1)
+
+
+    val markersQueue = channel.queueForBroker(1).get.markersQueue
+    assertEquals(1, markersQueue.size())
+    assertEquals(0, markersQueue.peek().metadataPartition)
+  }
+
+
+
+  def errorCallback(error: Errors): Unit = {}
+}