diff --git a/tools/build_mocks b/tools/build_mocks
index e23ae58af5..131a2c80b0 100755
--- a/tools/build_mocks
+++ b/tools/build_mocks
@@ -6,7 +6,7 @@ COMMENT_PREFIX = '/// '
 current_method = None
 current_package = None
 
-def split_to_parts(line):
+def split_to_parts(line, mod_file=None):
     global current_method
     global current_package
     # Line is beginning of method
@@ -17,6 +17,8 @@ def split_to_parts(line):
         
         *current_package, method_name = current_method.split('.')
 
+        yield (current_package, "\n")
+        yield (current_package, '# ' + mod_file + "\n")
         line = line.replace(current_method, method_name)
         
     yield (current_package, line)
@@ -50,7 +52,7 @@ def build_module(mod_file, dest):
             continue
         
         l = l[len(COMMENT_PREFIX):]#.strip()
-        store_to_file(dest, split_to_parts(l))
+        store_to_file(dest, split_to_parts(l, mod_file))
     
 def build_directory(dir, dest):
     print("Building mocks for", dir, "to", dest)
@@ -69,7 +71,12 @@ def clear_directory(top_dir):
             os.remove(os.path.join(root, name))
         for name in dirs:
             #print('Deleting directory', os.path.join(root, name))
-            os.rmdir(os.path.join(root, name))
+            try:
+                os.rmdir(os.path.join(root, name))
+            except FileNotFoundError:
+                pass
+
+        os.rmdir(root)
         
 if __name__ == '__main__':
     clear_directory('../mocks')