diff --git a/src/main/java/info/batey/kafka/unit/KafkaUnit.java b/src/main/java/info/batey/kafka/unit/KafkaUnit.java
index 0fda327..a350cc5 100644
--- a/src/main/java/info/batey/kafka/unit/KafkaUnit.java
+++ b/src/main/java/info/batey/kafka/unit/KafkaUnit.java
@@ -46,7 +46,7 @@
 import java.util.*;
 import java.util.concurrent.*;
 
-public class KafkaUnit {
+public class KafkaUnit<K, V> {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(KafkaUnit.class);
 
@@ -57,7 +57,7 @@ public class KafkaUnit {
     private final String brokerString;
     private int zkPort;
     private int brokerPort;
-    private Producer<String, String> producer = null;
+    private Producer<K, V> producer = null;
     private Properties kafkaBrokerConfig = new Properties();
 
     public KafkaUnit() throws IOException {
@@ -252,7 +252,7 @@ public List<T> call() throws Exception {
     }
 
     @SafeVarargs
-    public final void sendMessages(KeyedMessage<String, String> message, KeyedMessage<String, String>... messages) {
+    public final void sendMessages(KeyedMessage<K, V> message, KeyedMessage<K, V>... messages) {
         if (producer == null) {
             Properties props = new Properties();
             props.put("serializer.class", StringEncoder.class.getName());
diff --git a/src/test/java/info/batey/kafka/unit/KafkaIntegrationTest.java b/src/test/java/info/batey/kafka/unit/KafkaIntegrationTest.java
index 6a32695..8c3242c 100644
--- a/src/test/java/info/batey/kafka/unit/KafkaIntegrationTest.java
+++ b/src/test/java/info/batey/kafka/unit/KafkaIntegrationTest.java
@@ -38,7 +38,7 @@
 
 public class KafkaIntegrationTest {
 
-    private KafkaUnit kafkaUnitServer;
+    private KafkaUnit<String, String> kafkaUnitServer;
 
     @Before
     public void setUp() throws Exception {