'CleanupOrFail': Cleanup but returns error if it would cleanup with -u

README.md

@@ -51,6 +51,19 @@ func TestMain(m *testing.M) {
 
 Once included, if the update `-u` flag is used when running tests, any snapshot that is no longer in use will be removed. Note: if a single test is run, pruning _will not occur_.
 
+Alternatively `CleanupOrFail` can be used to fail a test run if a snapshot needs cleaning up but the `-u` flag wasn't given (and it's not a single-test run):
+
+```go
+func TestMain(m *testing.M) {
+  if m.Run() == 0 {
+    if err := abide.CleanupOrFail(); err != nil {
+      fmt.Fprintln(os.Stderr, err.Error())
+      os.Exit(1)
+    }
+  }
+}
+```
+
 ## Snapshots
 
 A snapshot is essentially a lock file for an http response. Instead of having to manually compare every aspect of an http response to it's expected value, it can be automatically generated and used for matching in subsequent testing.
@@ -105,4 +118,4 @@ To write snapshots to a directory other than the default `__snapshot__`, adjust
 func init() {
   abide.SnapshotDir = "testdata"
 }
-``` 
+``` 

abide.go

@@ -50,6 +50,32 @@ func Cleanup() error {
 	return allSnapshots.save()
 }
 
+// CleanupOrFail is an optional method which will behave like
+// Cleanup() if the `-u` flag was given, but which returns an error if
+// `-u` was not given and there were things to clean up.
+func CleanupOrFail() error {
+	if args.singleRun {
+		return nil
+	}
+	if args.shouldUpdate {
+		return Cleanup()
+	}
+
+	failed := 0
+	for _, s := range allSnapshots {
+		if !s.evaluated {
+			failed++
+			fmt.Fprintf(os.Stderr, "Unused snapshot `%s`\n", s.id)
+		}
+	}
+
+	if failed > 0 {
+		return fmt.Errorf("%d unused snapshots", failed)
+	}
+
+	return nil
+}
+
 // snapshotID represents the unique identifier for a snapshot.
 type snapshotID string
 

abide_test.go

@@ -1,6 +1,7 @@
 package abide
 
 import (
+	"fmt"
 	"os"
 	"reflect"
 	"testing"
@@ -34,6 +35,8 @@ func TestCleanup(t *testing.T) {
 	_ = testingSnapshot("1", "A")
 
 	// If shouldUpdate = false, the snapshot must remain.
+	args.shouldUpdate = false
+	args.singleRun = false
 	err := Cleanup()
 	if err != nil {
 		t.Fatal(err)
@@ -69,6 +72,56 @@ func TestCleanup(t *testing.T) {
 	}
 }
 
+func TestCleanupOrFail(t *testing.T) {
+	defer testingCleanup()
+
+	_ = testingSnapshot("1", "A")
+
+	args.shouldUpdate = false
+	args.singleRun = true
+	// singleRun means no cleanup
+	err := CleanupOrFail()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// shouldUpdate=false and singleRun=false -> CleanupOrFail fails
+	args.singleRun = false
+	err = CleanupOrFail()
+	if fmt.Sprint(err) != "1 unused snapshots" {
+		t.Fatalf("expected `1 unused snapshots`, got %v", err)
+	}
+
+	err = loadSnapshots()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	snapshot := getSnapshot("1")
+	if snapshot == nil {
+		t.Fatal("Expected snapshot[1] to exist.")
+	}
+
+	// If shouldUpdate = true and singleRun = false, the snapshot must be removed.
+	args.shouldUpdate = true
+	args.singleRun = false
+	err = CleanupOrFail()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// call private reloadSnapshots to repeat once-executing function
+	err = reloadSnapshots()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	snapshot = getSnapshot("1")
+	if snapshot != nil {
+		t.Fatal("Expected snapshot[1] to be removed.")
+	}
+}
+
 func TestCleanupUpdate(t *testing.T) {
 	defer testingCleanup()
 

example/main_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -11,9 +12,12 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	exit := m.Run()
-	abide.Cleanup()
-	os.Exit(exit)
+	if m.Run() == 0 {
+		if err := abide.CleanupOrFail(); err != nil {
+			fmt.Fprintln(os.Stderr, err.Error())
+			os.Exit(1)
+		}
+	}
 }
 
 func TestRequests(t *testing.T) {