From 03ecd7055d314164250f49763782538ddb3389b0 Mon Sep 17 00:00:00 2001
From: "John R. Lenton" <jlenton@gmail.com>
Date: Mon, 22 Mar 2021 21:56:38 +0000
Subject: [PATCH] 'CleanupOrFail': Cleanup but returns error if it would
 cleanup with -u

I'd rather know early if there are snapshots that need cleaning up,
rather than finding out while adding something unrelated. Keeps things
tidy.

I thought of giving Cleanup() this behaviour directly, but suspected a
breaking change like that would not be welcome. So instead I wrote its
wacky sister function. If you see the name, cringe, and think of a
better one, speak up -- I'm terrible at names.

HTH, HAND!
---
 README.md            | 15 ++++++++++++-
 abide.go             | 26 ++++++++++++++++++++++
 abide_test.go        | 53 ++++++++++++++++++++++++++++++++++++++++++++
 example/main_test.go | 10 ++++++---
 4 files changed, 100 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 877ae32..d414f75 100644
--- a/README.md
+++ b/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"
 }
-``` 
\ No newline at end of file
+``` 
diff --git a/abide.go b/abide.go
index 41df490..572f1dd 100644
--- a/abide.go
+++ b/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
 
diff --git a/abide_test.go b/abide_test.go
index 4211283..42fd3f8 100644
--- a/abide_test.go
+++ b/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()
 
diff --git a/example/main_test.go b/example/main_test.go
index 7d1a24d..0fd8dc2 100644
--- a/example/main_test.go
+++ b/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) {