diff --git a/java_gen/codegen.py b/java_gen/codegen.py
index f3619d4..bc9b9a0 100644
--- a/java_gen/codegen.py
+++ b/java_gen/codegen.py
@@ -37,6 +37,7 @@
 from loxi_ir import *
 import lang_java
 import test_data
+from import_cleaner import ImportCleaner
 
 import loxi_utils.loxi_utils as loxi_utils
 
@@ -84,6 +85,14 @@
         print "filename: %s" % filename
         with open(filename, "w") as f:
             loxi_utils.render_template(f, template, [self.templates_dir], context, prefix=prefix)
+        
+        try:
+            cleaner = ImportCleaner(filename)
+            cleaner.find_used_imports()
+            cleaner.rewrite_file(filename)
+        except:
+            print 'Cannot clean imports from file %s' % filename
+        
 
     def create_of_const_enums(self):
         for enum in self.java_model.enums:
@@ -126,13 +135,16 @@
                 else:
                     print "Class %s ignored by generate_class" % java_class.name
 
-    def create_unit_test(self, unit_test):
-        if unit_test.has_test_data:
-            self.render_class(clazz=unit_test,
-                    template='unit_test.java', src_dir="src/test/java",
-                    version=unit_test.java_class.version,
-                    test=unit_test, msg=unit_test.java_class,
-                    test_data=unit_test.test_data)
+    def create_unit_test(self, unit_tests):
+        if unit_tests.has_test_data:
+            for i in range(unit_tests.length):
+                unit_test = unit_tests.get_test_unit(i)
+                if unit_test.has_test_data:
+                    self.render_class(clazz=unit_test,
+                            template='unit_test.java', src_dir="src/test/java",
+                            version=unit_test.java_class.version,
+                            test=unit_test, msg=unit_test.java_class,
+                            test_data=unit_test.test_data)
 
     def create_of_factories(self):
         factory = self.java_model.of_factory
diff --git a/java_gen/import_cleaner.py b/java_gen/import_cleaner.py
new file mode 100755
index 0000000..83897d4
--- /dev/null
+++ b/java_gen/import_cleaner.py
@@ -0,0 +1,77 @@
+#!/usr/bin/python
+
+import sys
+import re
+
+class ImportLine:
+    def __init__(self, line):
+        self.line = line
+        class_name = None
+        if line[len(line) - 1] == '*':
+            class_name = '*'
+        else:
+            i = 7
+            while i < len(line) - 1:
+                if re.match('\.[A-Z][\..]*$', line[i - 1 : len(line) - 1]):
+                    class_name = line[i : len(line) - 1]
+                    break
+                i = i + 1
+            if class_name is None:
+                class_name = line[line.rfind('.') + 1 : len(line) - 1]
+        self.class_name = class_name
+
+
+class ImportCleaner:
+    def __init__(self, path):
+        f = open(path)
+        self.imp_lines = []
+        self.code_lines = []
+        self.imports_first_line = -1
+        i = 0
+        for line in f:
+            if len(line) > 6 and re.match('^[ \t]*import ', line):
+                self.imp_lines.append(ImportLine(line.rstrip()))
+                if self.imports_first_line == -1:
+                    self.imports_first_line = i
+            else:
+                self.code_lines.append(line.rstrip())
+            i = i + 1
+        f.close()
+
+    def find_used_imports(self):
+        self.used_imports = []
+        for line in self.code_lines:
+            temp = []
+            for imp in self.imp_lines:
+                if imp.class_name == '*' or line.find(imp.class_name) > -1:
+                    temp.append(imp)
+            for x in temp:
+                self.imp_lines.remove(x)
+                self.used_imports.append(x)
+
+    def rewrite_file(self, path):
+        f = open(path, 'w')
+        imports_written = False
+        for i in range(len(self.code_lines)):
+            if not imports_written and self.imports_first_line == i:
+                # Put all imports
+                for imp in self.used_imports:
+                    f.write(imp.line + '\n')
+                imports_written = True
+            # Put next code line
+            f.write(self.code_lines[i] + '\n')
+        f.close()
+
+def main(argv):
+    if len(argv) != 2:
+        print 'Usage: ImportCleaner <java file>'
+        return
+
+    filename = argv[1]
+    print 'Cleaning imports from file %s' % (filename)
+    cleaner = ImportCleaner(filename)
+    cleaner.find_used_imports()
+    cleaner.rewrite_file(filename)
+
+if __name__ == '__main__':
+    main(sys.argv)
diff --git a/java_gen/java_model.py b/java_gen/java_model.py
index 4ac497f..b93b330 100644
--- a/java_gen/java_model.py
+++ b/java_gen/java_model.py
@@ -200,7 +200,7 @@
             self.parent_interface = parent_interface
         else:
             self.parent_interface = None
-
+            
     def class_info(self):
         if re.match(r'OF.+StatsRequest$', self.name):
             return ("", "OFStatsRequest")
@@ -300,7 +300,7 @@
     @property
     @memoize
     def unit_test(self):
-        return JavaUnitTest(self)
+        return JavaUnitTestSet(self)
 
     @property
     def name(self):
@@ -563,18 +563,58 @@
 ### Unit Test
 #######################################################################
 
-class JavaUnitTest(object):
+class JavaUnitTestSet(object):
     def __init__(self, java_class):
         self.java_class = java_class
-        self.data_file_name = "of{version}/{name}.data".format(version=java_class.version.of_version,
+        first_data_file_name = "of{version}/{name}.data".format(version=java_class.version.of_version,
                                                      name=java_class.c_name[3:])
+        data_file_template = "of{version}/{name}.".format(version=java_class.version.of_version,
+                                                     name=java_class.c_name[3:]) + "{i}.data"
+        test_class_name = self.java_class.name + "Test"
+        self.test_units = []
+        if test_data.exists(first_data_file_name):
+            self.test_units.append(JavaUnitTest(java_class, first_data_file_name, test_class_name))
+        i = 1
+        while test_data.exists(data_file_template.format(i=i)):
+            self.test_units.append(JavaUnitTest(java_class, data_file_template.format(i=i), test_class_name + str(i)))
+            i = i + 1
+        
+    @property
+    def package(self):
+        return self.java_class.package
+
+    @property
+    def has_test_data(self):
+        return len(self.test_units) > 0
+
+    @property
+    def length(self):
+        return len(self.test_units)
+    
+    def get_test_unit(self, i):
+        return self.test_units[i]
+
+
+class JavaUnitTest(object):
+    def __init__(self, java_class, file_name=None, test_class_name=None):
+        self.java_class = java_class
+        if file_name is None:
+            self.data_file_name = "of{version}/{name}.data".format(version=java_class.version.of_version,
+                                                         name=java_class.c_name[3:])
+        else:
+            self.data_file_name = file_name
+        if test_class_name is None:
+            self.test_class_name = self.java_class.name + "Test"
+        else:
+            self.test_class_name = test_class_name
+        
     @property
     def package(self):
         return self.java_class.package
 
     @property
     def name(self):
-        return self.java_class.name + "Test"
+        return self.test_class_name
 
     @property
     def has_test_data(self):