Merge pull request #158 from dolmen-go/fix-detection-of-missing-Done-…

…in-Defer-from-Defer

Fix detection of missing done in Defer from Defer

cleanup_test.go

@@ -37,6 +37,26 @@ func TestCDeferWithoutDone(t *testing.T) {
 	c.Assert(tc.cleanup, qt.PanicMatches, `Done not called after Defer`)
 }
 
+func TestCDeferFromDefer(t *testing.T) {
+	c := qt.New(t)
+	tc := &testingTWithCleanup{
+		TB:      t,
+		cleanup: func() {},
+	}
+	c1 := qt.New(tc)
+	c1.Defer(func() {
+		c1.Log("defer 1")
+		// This defer is triggered from the first Done().
+		// It should have its own Done() call too.
+		c1.Defer(func() {
+			c1.Log("defer 2")
+		})
+	})
+	c1.Done()
+	// Check that we report the missing second Done().
+	c.Assert(tc.cleanup, qt.PanicMatches, `Done not called after Defer`)
+}
+
 func TestCDeferVsCleanupOrder(t *testing.T) {
 	c := qt.New(t)
 	var defers []int

quicktest.go

@@ -109,22 +109,19 @@ func (c *C) Defer(f func()) {
 		// Use TB.Cleanup when available, but add a check
 		// that Done has been called so that we don't run
 		// into unexpected Go version incompatibilities.
-		if c.doneNeeded {
-			// We've already installed the wrapper func that checks for Done
-			// so we can avoid doing it again.
-			cleaner.Cleanup(f)
-			return
+		if !c.doneNeeded {
+			c.doneNeeded = true
+			cleaner.Cleanup(func() {
+				c.mu.Lock()
+				defer c.mu.Unlock()
+				if c.doneNeeded {
+					panic("Done not called after Defer")
+				}
+			})
 		}
-		c.doneNeeded = true
-		cleaner.Cleanup(func() {
-			c.mu.Lock()
-			doneNeeded := c.doneNeeded
-			c.mu.Unlock()
-			if doneNeeded {
-				panic("Done not called after Defer")
-			}
-			f()
-		})
+
+		cleaner.Cleanup(f)
+
 		return
 	}