storage: Add object type hint parameter to ObjectStorage.Get. (#69)

Some storage backends can optimize object lookup if they get
the object type that is expected. So we the signature of the Get
method is now Get(Hash, ObjectType).

Added generic tests for storage backends.

5 files changed, 120 insertions, 108 deletions

diff --git a/storage/filesystem/object.go b/storage/filesystem/object.go
index f3a1dda..6024ae0 100644
--- a/storage/filesystem/object.go
+++ b/storage/filesystem/object.go
@@ -38,7 +38,7 @@ func (s *ObjectStorage) Set(core.Object) (core.Hash, error) {
 
 // Get returns the object with the given hash, by searching for it in
 // the packfile.
-func (s *ObjectStorage) Get(h core.Hash) (core.Object, error) {
+func (s *ObjectStorage) Get(h core.Hash, t core.ObjectType) (core.Object, error) {
 	offset, err := s.index.Get(h)
 	if err != nil {
 		return nil, err
@@ -71,7 +71,14 @@ func (s *ObjectStorage) Get(h core.Hash) (core.Object, error) {
 	p := packfile.NewParser(r)
 
 	obj := s.NewObject()
-	return obj, p.FillObject(obj)
+	err = p.FillObject(obj)
+	if err != nil {
+		return nil, err
+	}
+	if core.AnyObject != t && obj.Type() != t {
+		return nil, core.ErrObjectNotFound
+	}
+	return obj, nil
 }
 
 // Iter returns an iterator for all the objects in the packfile with the
@@ -80,11 +87,11 @@ func (s *ObjectStorage) Iter(t core.ObjectType) (core.ObjectIter, error) {
 	var objects []core.Object
 
 	for hash := range s.index {
-		object, err := s.Get(hash)
+		object, err := s.Get(hash, core.AnyObject)
 		if err != nil {
 			return nil, err
 		}
-		if object.Type() == t {
+		if t == core.AnyObject || object.Type() == t {
 			objects = append(objects, object)
 		}
 	}
diff --git a/storage/filesystem/object_test.go b/storage/filesystem/object_test.go
index 692a69b..956fdeb 100644
--- a/storage/filesystem/object_test.go
+++ b/storage/filesystem/object_test.go
@@ -64,9 +64,11 @@ func (s *FsSuite) TearDownSuite(c *C) {
 
 func (s *FsSuite) TestHashNotFound(c *C) {
 	sto := s.newObjectStorage(c, "binary-relations")
-
-	_, err := sto.Get(core.ZeroHash)
-	c.Assert(err, Equals, core.ErrObjectNotFound)
+	types := []core.ObjectType{core.AnyObject, core.TagObject, core.CommitObject, core.BlobObject, core.TreeObject}
+	for t := range types {
+		_, err := sto.Get(core.ZeroHash, core.ObjectType(t))
+		c.Assert(err, Equals, core.ErrObjectNotFound)
+	}
 }
 
 func (s *FsSuite) newObjectStorage(c *C, fixtureName string) core.ObjectStorage {
@@ -156,7 +158,7 @@ func equalsStorages(a, b core.ObjectStorage) (bool, string, error) {
 				break
 			}
 
-			bo, err := b.Get(ao.Hash())
+			bo, err := b.Get(ao.Hash(), core.AnyObject)
 			if err != nil {
 				return false, "", fmt.Errorf("getting object with hash %s: %s",
 					ao.Hash(), err)
diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index d67368f..8033541 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -130,12 +130,11 @@ func (o *ObjectStorage) Set(obj core.Object) (core.Hash, error) {
 }
 
 // Get returns a object with the given hash
-func (o *ObjectStorage) Get(h core.Hash) (core.Object, error) {
+func (o *ObjectStorage) Get(h core.Hash, t core.ObjectType) (core.Object, error) {
 	obj, ok := o.Objects[h]
-	if !ok {
+	if !ok || (core.AnyObject != t && obj.Type() != t) {
 		return nil, core.ErrObjectNotFound
 	}
-
 	return obj, nil
 }
 
@@ -143,6 +142,8 @@ func (o *ObjectStorage) Get(h core.Hash) (core.Object, error) {
 func (o *ObjectStorage) Iter(t core.ObjectType) (core.ObjectIter, error) {
 	var series []core.Object
 	switch t {
+	case core.AnyObject:
+		series = flattenObjectMap(o.Objects)
 	case core.CommitObject:
 		series = flattenObjectMap(o.Commits)
 	case core.TreeObject:
diff --git a/storage/memory/storage_test.go b/storage/memory/storage_test.go
index ac97584..e291609 100644
--- a/storage/memory/storage_test.go
+++ b/storage/memory/storage_test.go
@@ -6,14 +6,20 @@ import (
 
 	. "gopkg.in/check.v1"
 	"gopkg.in/src-d/go-git.v4/core"
+	. "gopkg.in/src-d/go-git.v4/storage/test"
 )
 
 func Test(t *testing.T) { TestingT(t) }
 
-type StorageSuite struct{}
+type StorageSuite struct { }
 
 var _ = Suite(&StorageSuite{})
 
+func (s *StorageSuite) TestObjectStorage(c *C) {
+	storage := NewStorage()
+	RunObjectStorageSuite(c, storage.ObjectStorage())
+}
+
 func (s *StorageSuite) TestStorageObjectStorage(c *C) {
 	storage := NewStorage()
 	o := storage.ObjectStorage()
@@ -30,102 +36,6 @@ func (s *StorageSuite) TestStorageReferenceStorage(c *C) {
 	c.Assert(o == e, Equals, true)
 }
 
-func (s *StorageSuite) TestObjectStorageSetAndGet(c *C) {
-	storage := NewStorage()
-	os := storage.ObjectStorage()
-
-	commit := &core.MemoryObject{}
-	commit.SetType(core.CommitObject)
-
-	h, err := os.Set(commit)
-	c.Assert(err, IsNil)
-	c.Assert(h.String(), Equals, "dcf5b16e76cce7425d0beaef62d79a7d10fce1f5")
-
-	e, err := os.Get(h)
-	c.Assert(commit == e, Equals, true)
-
-	tree := &core.MemoryObject{}
-	tree.SetType(core.TreeObject)
-
-	h, err = os.Set(tree)
-	c.Assert(err, IsNil)
-	c.Assert(h.String(), Equals, "4b825dc642cb6eb9a060e54bf8d69288fbee4904")
-
-	e, err = os.Get(h)
-	c.Assert(tree == e, Equals, true)
-
-	blob := &core.MemoryObject{}
-	blob.SetType(core.BlobObject)
-
-	h, err = os.Set(blob)
-	c.Assert(err, IsNil)
-	c.Assert(h.String(), Equals, "e69de29bb2d1d6434b8b29ae775ad8c2e48c5391")
-
-	e, err = os.Get(h)
-	c.Assert(blob == e, Equals, true)
-
-	tag := &core.MemoryObject{}
-	tag.SetType(core.TagObject)
-
-	h, err = os.Set(tag)
-	c.Assert(err, IsNil)
-	c.Assert(h.String(), Equals, "d994c6bb648123a17e8f70a966857c546b2a6f94")
-
-	e, err = os.Get(h)
-	c.Assert(tag == e, Equals, true)
-}
-
-func (s *StorageSuite) TestObjectStorageIter(c *C) {
-	commit := &core.MemoryObject{}
-	commit.SetType(core.CommitObject)
-	tree := &core.MemoryObject{}
-	tree.SetType(core.TreeObject)
-	blob := &core.MemoryObject{}
-	blob.SetType(core.BlobObject)
-	tag := &core.MemoryObject{}
-	tag.SetType(core.TagObject)
-
-	storage := NewStorage()
-	os := storage.ObjectStorage()
-
-	os.Set(commit)
-	os.Set(tree)
-	os.Set(blob)
-	os.Set(tag)
-
-	i, err := os.Iter(core.CommitObject)
-	c.Assert(err, IsNil)
-
-	e, err := i.Next()
-	c.Assert(err, IsNil)
-	c.Assert(commit == e, Equals, true)
-
-	i, err = os.Iter(core.TreeObject)
-	c.Assert(err, IsNil)
-
-	e, err = i.Next()
-	c.Assert(err, IsNil)
-	c.Assert(tree == e, Equals, true)
-
-	i, err = os.Iter(core.BlobObject)
-	c.Assert(err, IsNil)
-
-	e, err = i.Next()
-	c.Assert(err, IsNil)
-	c.Assert(blob == e, Equals, true)
-
-	i, err = os.Iter(core.TagObject)
-	c.Assert(err, IsNil)
-
-	e, err = i.Next()
-	c.Assert(err, IsNil)
-	c.Assert(tag == e, Equals, true)
-
-	e, err = i.Next()
-	c.Assert(e, IsNil)
-	c.Assert(err, Equals, io.EOF)
-}
-
 func (s *StorageSuite) TestReferenceStorageSetAndGet(c *C) {
 	storage := NewStorage()
 	rs := storage.ReferenceStorage()
diff --git a/storage/test/storage_suite.go b/storage/test/storage_suite.go
new file mode 100644
index 0000000..2463d9d
--- /dev/null
+++ b/storage/test/storage_suite.go
@@ -0,0 +1,92 @@
+package test
+
+import (
+	. "gopkg.in/check.v1"
+	"gopkg.in/src-d/go-git.v4/core"
+	"io"
+)
+
+type TestObject struct {
+	Object core.Object
+	Hash   string
+	Type   core.ObjectType
+}
+
+func RunObjectStorageSuite(c *C, os core.ObjectStorage) {
+	commit := &core.MemoryObject{}
+	commit.SetType(core.CommitObject)
+	tree := &core.MemoryObject{}
+	tree.SetType(core.TreeObject)
+	blob := &core.MemoryObject{}
+	blob.SetType(core.BlobObject)
+	tag := &core.MemoryObject{}
+	tag.SetType(core.TagObject)
+
+	testObjects := map[core.ObjectType]TestObject{
+		core.CommitObject: TestObject{commit, "dcf5b16e76cce7425d0beaef62d79a7d10fce1f5", core.CommitObject},
+		core.TreeObject:   TestObject{tree, "4b825dc642cb6eb9a060e54bf8d69288fbee4904", core.TreeObject},
+		core.BlobObject:   TestObject{blob, "e69de29bb2d1d6434b8b29ae775ad8c2e48c5391", core.BlobObject},
+		core.TagObject:    TestObject{tag, "d994c6bb648123a17e8f70a966857c546b2a6f94", core.TagObject},
+	}
+
+	validTypes := []core.ObjectType{core.CommitObject, core.BlobObject, core.TagObject, core.TreeObject}
+
+	for _, to := range testObjects {
+		comment := Commentf("failed for type %s", to.Type.String())
+
+		h, err := os.Set(to.Object)
+		c.Assert(err, IsNil)
+		c.Assert(h.String(), Equals, to.Hash, comment)
+
+		o, err := os.Get(h, to.Type)
+		c.Assert(err, IsNil)
+		c.Assert(o, Equals, to.Object)
+
+		o, err = os.Get(h, core.AnyObject)
+		c.Assert(err, IsNil)
+		c.Assert(o, Equals, to.Object)
+
+		for _, validType := range validTypes {
+			if validType == to.Type {
+				continue
+			}
+			o, err = os.Get(h, validType)
+			c.Assert(o, IsNil)
+			c.Assert(err, Equals, core.ErrObjectNotFound)
+		}
+	}
+
+	for _, validType := range validTypes {
+		comment := Commentf("failed for type %s)", validType.String())
+		i, err := os.Iter(validType)
+		c.Assert(err, IsNil, comment)
+
+		o, err := i.Next()
+		c.Assert(err, IsNil)
+		c.Assert(o, Equals, testObjects[validType].Object, comment)
+
+		o, err = i.Next()
+		c.Assert(o, IsNil)
+		c.Assert(err, Equals, io.EOF, comment)
+	}
+
+	i, err := os.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+
+	foundObjects := []core.Object{}
+	i.ForEach(func(o core.Object) error {
+		foundObjects = append(foundObjects, o)
+		return nil
+	})
+	c.Assert(foundObjects, HasLen, len(testObjects))
+	for _, to := range testObjects {
+		found := false
+		for _, o := range foundObjects {
+			if to.Object == o {
+				found = true
+				break
+			}
+		}
+		c.Assert(found, Equals, true, Commentf("Object of type %s not found", to.Type.String()))
+	}
+}