2 files changed, 11 insertions, 100 deletions

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()