Add Cleanup and Only

Author: evanphx

neko.go

@@ -18,9 +18,11 @@ type test struct {
 type Organizer struct {
 	t *testing.T
 
-	mocks []*mock.Mock
-	setup []func()
-	tests []test
+	only    *test
+	mocks   []*mock.Mock
+	setup   []func()
+	cleanup []func()
+	tests   []test
 }
 
 // Create a new Organizer against testing's T interface
@@ -38,11 +40,20 @@ func (o *Organizer) Setup(f func()) {
 	o.setup = append(o.setup, f)
 }
 
+// Add some work to be done after each test
+func (o *Organizer) Cleanup(f func()) {
+	o.cleanup = append(o.cleanup, f)
+}
+
 // Add a test.
 func (o *Organizer) It(name string, f func()) {
 	o.tests = append(o.tests, test{name, f})
 }
 
+func (o *Organizer) Only(name string, f func()) {
+	o.only = &test{name, f}
+}
+
 // Useful by allowing the developer to simply add 'N' before
 // It to disable a block.
 func (o *Organizer) NIt(name string, f func()) {
@@ -51,28 +62,45 @@ func (o *Organizer) NIt(name string, f func()) {
 
 // Coordinate running the tests with the setups and mocks
 func (o *Organizer) Run() {
+	if o.only != nil {
+		o.runTest(o.only)
+		return
+	}
+
 	for _, test := range o.tests {
-		if test.Func == nil {
-			o.t.Logf("==== DISABLED: %s ====", test.Name)
-			continue
-		}
+		o.runTest(&test)
+	}
+}
 
-		o.t.Logf("==== %s ====", test.Name)
+func (o *Organizer) runTest(test *test) {
+	if test.Func == nil {
+		o.t.Logf("==== DISABLED: %s ====", test.Name)
+		return
+	}
+
+	o.t.Logf("==== %s ====", test.Name)
+
+	for _, mock := range o.mocks {
+		mock.ExpectedCalls = nil
+		mock.Calls = nil
+	}
 
-		for _, mock := range o.mocks {
-			mock.ExpectedCalls = nil
-			mock.Calls = nil
-		}
+	for _, setup := range o.setup {
+		setup()
+	}
+
+	defer o.runCleanup()
 
-		for _, setup := range o.setup {
-			setup()
-		}
+	test.Func()
 
-		test.Func()
+	for _, mock := range o.mocks {
+		mock.AssertExpectations(o.t)
+	}
+}
 
-		for _, mock := range o.mocks {
-			mock.AssertExpectations(o.t)
-		}
+func (o *Organizer) runCleanup() {
+	for _, cleanup := range o.cleanup {
+		cleanup()
 	}
 }