1 files changed, 31 insertions, 3 deletions

diff --git a/src/tools/mkunit.cpp b/src/tools/mkunit.cpp
index f6ea669..5711c42 100644
--- a/src/tools/mkunit.cpp
+++ b/src/tools/mkunit.cpp
@@ -42,6 +42,7 @@ enum TokType
     tokTest,
     tokChar,
     tokBlock,
+    tokCleanup,
     tokEof
 };
 
@@ -54,6 +55,7 @@ Bu::Formatter &operator<<( Bu::Formatter &f, TokType t )
         case tokTest:       return f << "tokTest";
         case tokChar:       return f << "tokChar";
         case tokBlock:      return f << "tokBlock";
+        case tokCleanup:    return f << "tokCleanup";
         case tokEof:        return f << "tokEof";
     }
 
@@ -193,6 +195,8 @@ public:
                     {
                         if( sTok == "test" )
                             return tokTest;
+                        else if( sTok == "cleanup" )
+                            return tokCleanup;
                         else
                         {
                             v = sTok;
@@ -361,6 +365,15 @@ public:
                                 s.lTest.append( t );
                             }
                             break;
+                        
+                        case tokCleanup:
+                            {
+                                if( nextToken( v, sWs, iL, iC ) != tokBlock )
+                                    throw Bu::ExceptionBase("%d:%d: Expected "
+                                        "{...} block.",
+                                        iL, iC );
+                            }
+                            break;
 
                         case tokChar:
                             if( v.get<char>() == '}' )
@@ -438,7 +451,7 @@ public:
                         for( TestList::iterator i = s.lTest.begin(); i; i++ )
                         {
                             f << "\t\tadd( static_cast<Bu::UnitSuite::Test>("
-                                "&" << sClass << "::" << (*i).sName << "), \""
+                                "&" << sClass << "::_test_" << (*i).sName << "), \""
                                                                 << (*i).sName << "\", Bu::UnitSuite::"
                                 "expectPass );" << f.nl;
                         }
@@ -481,9 +494,24 @@ public:
                                         "{...} block.",
                                         iL, iC );
                                 
-                                f << "\tvoid " << t.sName << "()"
+                                f << "void _test_" << t.sName << "()"
+                                    << f.nl << "#line " << iL
+                                    << " \"" << sIn << "\"" << f.nl << "    "
+                                    << v << f.nl;
+                            }
+                            break;
+
+                        case tokCleanup:
+                            {
+                                fOut.write( sWs );
+                                if( nextToken( v, sWs, iL, iC ) != tokBlock )
+                                    throw Bu::ExceptionBase("%d:%d: Expected "
+                                        "{...} block.",
+                                        iL, iC );
+                                
+                                f << "void cleanup()"
                                     << f.nl << "#line " << iL
-                                    << " \"" << sIn << "\"" << f.nl
+                                    << " \"" << sIn << "\"" << f.nl << "    "
                                     << v << f.nl;
                             }
                             break;