dotgit: add CheckAndSetReference tests

2 files changed, 27 insertions, 11 deletions

diff --git a/storage/filesystem/internal/dotgit/dotgit_test.go b/storage/filesystem/internal/dotgit/dotgit_test.go
index d1b04a7..446a204 100644
--- a/storage/filesystem/internal/dotgit/dotgit_test.go
+++ b/storage/filesystem/internal/dotgit/dotgit_test.go
@@ -55,10 +55,11 @@ func (s *SuiteDotGit) TestSetRefs(c *C) {
 	fs := osfs.New(tmp)
 	dir := New(fs)
 
-	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+	firstFoo := plumbing.NewReferenceFromStrings(
 		"refs/heads/foo",
 		"e8d3ffab552895c19b9fcf7aa264d277cde33881",
-	), nil)
+	)
+	err = dir.SetRef(firstFoo, nil)
 
 	c.Assert(err, IsNil)
 
@@ -105,6 +106,20 @@ func (s *SuiteDotGit) TestSetRefs(c *C) {
 	c.Assert(ref, NotNil)
 	c.Assert(ref.Hash().String(), Equals, "e8d3ffab552895c19b9fcf7aa264d277cde33881")
 
+	// Check that SetRef with a non-nil `old` works.
+	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+		"refs/heads/foo",
+		"6ecf0ef2c2dffb796033e5a02219af86ec6584e5",
+	), firstFoo)
+	c.Assert(err, IsNil)
+
+	// `firstFoo` is no longer the right `old` reference, so this
+	// should fail.
+	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+		"refs/heads/foo",
+		"6ecf0ef2c2dffb796033e5a02219af86ec6584e5",
+	), firstFoo)
+	c.Assert(err, NotNil)
 }
 
 func (s *SuiteDotGit) TestRefsFromPackedRefs(c *C) {
@@ -192,7 +207,7 @@ func (s *SuiteDotGit) TestRemoveRefFromReferenceFileAndPackedRefs(c *C) {
 	err := dir.SetRef(plumbing.NewReferenceFromStrings(
 		"refs/remotes/origin/branch",
 		"e8d3ffab552895c19b9fcf7aa264d277cde33881",
-	))
+	), nil)
 
 	// Make sure it only appears once in the refs list.
 	refs, err := dir.Refs()
diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index 69394af..c9306ee 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -204,16 +204,17 @@ func (r ReferenceStorage) SetReference(ref *plumbing.Reference) error {
 }
 
 func (r ReferenceStorage) CheckAndSetReference(ref, old *plumbing.Reference) error {
-	if ref != nil {
-		if old != nil {
-			tmp := r[ref.Name()]
-			if tmp != nil && tmp.Hash() != old.Hash() {
-				return ErrRefHasChanged
-			}
-		}
-		r[ref.Name()] = ref
+	if ref == nil {
+		return nil
 	}
 
+	if old != nil {
+		tmp := r[ref.Name()]
+		if tmp != nil && tmp.Hash() != old.Hash() {
+			return ErrRefHasChanged
+		}
+	}
+	r[ref.Name()] = ref
 	return nil
 }