Merge pull request #52 from alexhornbake/alex/fix_update_cleanup_bug

fix cleanup after update bug and concurrent writes to allSnapshots

abide.go

@@ -14,6 +14,7 @@ import (
 var (
 	args         *arguments
 	allSnapshots snapshots
+	allSnapMutex sync.Mutex
 )
 
 var (
@@ -200,8 +201,18 @@ func getSnapshot(id snapshotID) *snapshot {
 	return allSnapshots[id]
 }
 
-// createSnapshot creates or updates a Snapshot.
+// createSnapshot creates a Snapshot.
 func createSnapshot(id snapshotID, value string) (*snapshot, error) {
+	return writeSnapshot(id, value, false)
+}
+
+// updateSnapshot creates a Snapshot.
+func updateSnapshot(id snapshotID, value string) (*snapshot, error) {
+	return writeSnapshot(id, value, true)
+}
+
+// writeSnapshot creates or updates a Snapshot.
+func writeSnapshot(id snapshotID, value string, isUpdate bool) (*snapshot, error) {
 	if !id.isValid() {
 		return nil, errInvalidSnapshotID
 	}
@@ -223,11 +234,15 @@ func createSnapshot(id snapshotID, value string) (*snapshot, error) {
 	path := filepath.Join(dir, fmt.Sprintf("%s%s", pkg, snapshotExt))
 
 	s := &snapshot{
-		id:    id,
-		value: value,
-		path:  path,
+		id:        id,
+		value:     value,
+		path:      path,
+		evaluated: isUpdate,
 	}
+
+	allSnapMutex.Lock()
 	allSnapshots[id] = s
+	allSnapMutex.Unlock()
 
 	err = allSnapshots.save()
 	if err != nil {

abide_test.go

@@ -69,6 +69,46 @@ func TestCleanup(t *testing.T) {
 	}
 }
 
+func TestCleanupUpdate(t *testing.T) {
+	defer testingCleanup()
+
+	// this snapshot is updated, should be evaluated, and not removed
+	_ = testingSnapshot("1", "A")
+	t2 := &testing.T{}
+	createOrUpdateSnapshot(t2, "1", "B")
+
+	// this snapshot is never evaluated, and should be removed
+	_ = testingSnapshot("2", "B")
+
+	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 := Cleanup()
+	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 exist.")
+	}
+	snapshot = getSnapshot("2")
+	if snapshot != nil {
+		t.Fatal("Expected snapshot[2] to be removed.")
+	}
+}
+
 func TestSnapshotIDIsValid(t *testing.T) {
 	id := snapshotID("1")
 	if !id.isValid() {

assert.go

@@ -160,7 +160,7 @@ func createOrUpdateSnapshot(t *testing.T, id, data string) {
 	if diff != "" {
 		if snapshot != nil && args.shouldUpdate {
 			fmt.Printf("Updating snapshot `%s`\n", id)
-			_, err = createSnapshot(snapshotID(id), data)
+			_, err = updateSnapshot(snapshotID(id), data)
 			if err != nil {
 				t.Fatal(err)
 			}