diff --git a/Makefile b/Makefile
index fd2be473..5448f4bd 100644
--- a/Makefile
+++ b/Makefile
@@ -58,8 +58,8 @@ cbins:
 
 $(ALL_SRC): ;
 
-# TODO(pourchet): Remove this hack to make dep more reliable. For some reason `go mod download` fails
-# sometimes on TravisCI, so run it a few times if it fails.
+# TODO(pourchet): Remove this hack to make dep more reliable. For some reason `go mod download`
+# fails sometimes on TravisCI, so run it a few times if it fails.
 vendor: go.mod go.sum
 	go mod download || go mod download || go mod download
 	go mod vendor
@@ -75,6 +75,8 @@ mocks: ext-tools
 	@echo "Generating mocks"
 	mkdir -p mocks/net/http
 	$(EXT_TOOLS_DIR)/mockgen -destination=mocks/net/http/mockhttp.go -package=mockhttp net/http RoundTripper
+	mkdir -p mocks/lib/registry
+	$(EXT_TOOLS_DIR)/mockgen -destination=mocks/lib/registry/mockclient.go -package=mockregistry github.com/uber/makisu/lib/registry Client
 
 env: test/python/requirements.txt
 	[ -d env ] || virtualenv --setuptools env
diff --git a/bin/makisu/cmd/build.go b/bin/makisu/cmd/build.go
index 47112a46..ba55f4ae 100644
--- a/bin/makisu/cmd/build.go
+++ b/bin/makisu/cmd/build.go
@@ -72,9 +72,9 @@ type buildCmd struct {
 func getBuildCmd() *buildCmd {
 	buildCmd := &buildCmd{
 		Command: &cobra.Command{
-			Use:                   "build -t=<image_tag> [flags] <context_path>",
+			Use: "build -t=<image_tag> [flags] <context_path>",
 			DisableFlagsInUseLine: true,
-			Short:                 "Build docker image, optionally push to registries and/or load into docker daemon",
+			Short: "Build docker image, optionally push to registries and/or load into docker daemon",
 		},
 	}
 	buildCmd.Args = func(cmd *cobra.Command, args []string) error {
@@ -245,7 +245,6 @@ func (cmd *buildCmd) Build(contextDir string) error {
 			}
 			defer rootPreserver.RestoreRoot()
 		}
-		log.Debugf("build.Cmd.Build() first call")
 		buildContext.MemFS.Remove()
 		defer buildContext.MemFS.Remove()
 	}
diff --git a/lib/builder/build_node.go b/lib/builder/build_node.go
index d2c9a4a0..7db73c67 100644
--- a/lib/builder/build_node.go
+++ b/lib/builder/build_node.go
@@ -166,6 +166,7 @@ func (n *buildNode) pushCacheLayer(cacheMgr cache.Manager) error {
 func (n *buildNode) pullCacheLayer(cacheMgr cache.Manager) bool {
 	digestPair, err := cacheMgr.PullCache(n.CacheID())
 	if err != nil {
+		// TODO: distinguish cache not found and pull failure.
 		log.Errorf("Failed to fetch intermediate layer with cache ID %s: %s", n.CacheID(), err)
 		return false
 	} else if digestPair == nil {
diff --git a/lib/builder/build_plan.go b/lib/builder/build_plan.go
index feabfbf6..2d82d5e6 100644
--- a/lib/builder/build_plan.go
+++ b/lib/builder/build_plan.go
@@ -70,11 +70,10 @@ func NewBuildPlan(
 		return nil, fmt.Errorf("build alias list: %s", err)
 	}
 
-	digestPairs := make(image.DigestPairMap)
 	for i, parsedStage := range parsedStages {
 		// Add this stage to the plan.
 		stage, err := newBuildStage(
-			ctx, parsedStage.From.Alias, parsedStage, digestPairs, plan.opts)
+			ctx, parsedStage.From.Alias, parsedStage, plan.opts)
 		if err != nil {
 			return nil, fmt.Errorf("failed to convert parsed stage: %s", err)
 		}
@@ -86,7 +85,7 @@ func NewBuildPlan(
 		plan.stages[i] = stage
 	}
 
-	if err := plan.handleCopyFromDirs(aliases, digestPairs); err != nil {
+	if err := plan.handleCopyFromDirs(aliases); err != nil {
 		return nil, fmt.Errorf("handle cross refs: %s", err)
 	}
 	return plan, nil
@@ -95,9 +94,7 @@ func NewBuildPlan(
 // handleCopyFromDirs goes through all of the stages in the build plan and looks
 // at the `COPY --from` steps to make sure they are valid. If the --from source
 // is another image, we create a new image stage in the build plan.
-func (plan *BuildPlan) handleCopyFromDirs(
-	aliases map[string]bool, digestPairs image.DigestPairMap) error {
-
+func (plan *BuildPlan) handleCopyFromDirs(aliases map[string]bool) error {
 	for _, stage := range plan.stages {
 		for alias, dirs := range stage.copyFromDirs {
 			if _, ok := aliases[alias]; !ok {
@@ -110,7 +107,7 @@ func (plan *BuildPlan) handleCopyFromDirs(
 					return fmt.Errorf("copy from nonexistent stage %s", alias)
 				}
 				remoteImageStage, err := newRemoteImageStage(
-					plan.baseCtx, alias, digestPairs, plan.opts)
+					plan.baseCtx, alias, plan.opts)
 				if err != nil {
 					return fmt.Errorf("new image stage: %s", err)
 				}
@@ -135,7 +132,7 @@ func buildAliases(stages dockerfile.Stages) (map[string]bool, error) {
 			if _, ok := aliases[parsedStage.From.Alias]; ok {
 				return nil, fmt.Errorf("duplicate stage alias: %s", parsedStage.From.Alias)
 			} else if _, err := strconv.Atoi(parsedStage.From.Alias); err == nil {
-				// Docker would return `name can't start with a number or contain symbols`
+				// Docker would return `name can't start with a number or contain symbols`.
 				return nil, fmt.Errorf("stage alias cannot be a number: %s", parsedStage.From.Alias)
 			}
 		} else {
@@ -148,13 +145,6 @@ func buildAliases(stages dockerfile.Stages) (map[string]bool, error) {
 
 // Execute executes all build stages in order.
 func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
-	// Execute pre-build procedures. Try to pull some reusable layers from the
-	// registry.
-	// TODO: Pull in parallel
-	for _, stage := range plan.stages {
-		stage.pullCacheLayers(plan.cacheMgr)
-	}
-
 	for alias, stage := range plan.remoteImageStages {
 		// Building that pseudo stage will unpack the image directly into the
 		// stage's cross stage directory.
@@ -174,6 +164,9 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 		currStage = plan.stages[k]
 		log.Infof("* Stage %d/%d : %s", k+1, len(plan.stages), currStage.String())
 
+		// Try to pull reusable layers cached from previous builds.
+		currStage.pullCacheLayers(plan.cacheMgr)
+
 		lastStage := k == len(plan.stages)-1
 		_, copiedFrom := plan.copyFromDirs[currStage.alias]
 
diff --git a/lib/builder/build_stage.go b/lib/builder/build_stage.go
index eb9e5cee..bb40714a 100644
--- a/lib/builder/build_stage.go
+++ b/lib/builder/build_stage.go
@@ -46,12 +46,11 @@ type buildStageOptions struct {
 
 // buildStage represents a sequence of steps to build intermediate layers or a final image.
 type buildStage struct {
-	ctx               *context.BuildContext
-	alias             string
-	copyFromDirs      map[string][]string
-	nodes             []*buildNode
-	lastImageConfig   *image.Config
-	sharedDigestPairs image.DigestPairMap
+	ctx             *context.BuildContext
+	alias           string
+	copyFromDirs    map[string][]string
+	nodes           []*buildNode
+	lastImageConfig *image.Config
 
 	opts *buildStageOptions
 }
@@ -59,7 +58,7 @@ type buildStage struct {
 // newBuildStage initializes a buildStage.
 func newBuildStage(
 	baseCtx *context.BuildContext, alias string, parsedStage *dockerfile.Stage,
-	digestPairs image.DigestPairMap, planOpts *buildPlanOptions) (*buildStage, error) {
+	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
 	ctx, err := context.NewBuildContext(
@@ -74,13 +73,12 @@ func newBuildStage(
 		return nil, fmt.Errorf("new dockerfile steps: %s", err)
 	}
 
-	return newBuildStageHelper(ctx, alias, steps, digestPairs, planOpts)
+	return newBuildStageHelper(ctx, alias, steps, planOpts)
 }
 
 // newRemoteImageStage initializes a buildStage.
 func newRemoteImageStage(
-	baseCtx *context.BuildContext, alias string, digestPairs image.DigestPairMap,
-	planOpts *buildPlanOptions) (*buildStage, error) {
+	baseCtx *context.BuildContext, alias string, planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
 	ctx, err := context.NewBuildContext(
@@ -105,12 +103,12 @@ func newRemoteImageStage(
 		allowModifyFS: planOpts.allowModifyFS,
 	}
 
-	return newBuildStageHelper(ctx, alias, steps, digestPairs, opts)
+	return newBuildStageHelper(ctx, alias, steps, opts)
 }
 
 func newBuildStageHelper(
 	ctx *context.BuildContext, alias string, steps []step.BuildStep,
-	digestPairs image.DigestPairMap, planOpts *buildPlanOptions) (*buildStage, error) {
+	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Convert each step to a build node.
 	var requireOnDisk bool
@@ -134,11 +132,10 @@ func newBuildStageHelper(
 	}
 
 	stage := &buildStage{
-		ctx:               ctx,
-		copyFromDirs:      copyFromDirs,
-		alias:             alias,
-		nodes:             nodes,
-		sharedDigestPairs: digestPairs,
+		ctx:          ctx,
+		copyFromDirs: copyFromDirs,
+		alias:        alias,
+		nodes:        nodes,
 		opts: &buildStageOptions{
 			allowModifyFS: planOpts.allowModifyFS,
 			forceCommit:   planOpts.forceCommit,
@@ -172,14 +169,6 @@ func createDockerfileSteps(
 // build performs the build for that stage. There are side effects that should
 // be expected on each node within the stage.
 func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom bool) error {
-	// Reuse the digestpairs that other stages have populated.
-	for _, node := range stage.nodes {
-		if pairs, ok := stage.sharedDigestPairs[node.CacheID()]; ok {
-			log.Infof("* Reusing digest pairs computed from earlier step %s", node.CacheID())
-			node.digestPairs = pairs
-		}
-	}
-
 	var err error
 	diffIDs := make([]image.Digest, 0)
 	histories := make([]image.History, 0)
@@ -214,11 +203,6 @@ func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom boo
 				Author:    "makisu",
 			})
 		}
-
-		// Update the shared map of cacheID to digest pair.
-		if len(node.digestPairs) != 0 {
-			stage.sharedDigestPairs[node.CacheID()] = node.digestPairs
-		}
 	}
 	stage.lastImageConfig.Created = time.Now()
 	stage.lastImageConfig.History = histories
@@ -310,12 +294,12 @@ func (stage *buildStage) saveManifest(
 	return manifest, nil
 }
 
-// pullCacheLayers attempts to pull reusable layers from the distributed cache. Terminates once
-// a node that can be cached fails to pull its layer.
+// pullCacheLayers attempts to pull reusable layers from the distributed cache.
+// Terminates once a node that can be cached fails to pull its layer.
 func (stage *buildStage) pullCacheLayers(cacheMgr cache.Manager) {
-	// Skip the first node since it's a FROM step. We do not want to try
-	// to pull from cache because the step itself will pull the right layers when
-	// it gets executed.
+	// Skip the first node since it's a FROM step. We do not want to try to pull
+	// from cache because the step itself will pull the right layers when it
+	// gets executed.
 	for _, node := range stage.nodes[1:] {
 		// Stop once the cache chain is broken.
 		if node.HasCommit() || stage.opts.forceCommit {
diff --git a/lib/builder/build_stage_test.go b/lib/builder/build_stage_test.go
index f9e73402..bf988f71 100644
--- a/lib/builder/build_stage_test.go
+++ b/lib/builder/build_stage_test.go
@@ -123,10 +123,10 @@ func TestPullCacheLayers(t *testing.T) {
 				allowModifyFS: false,
 			}
 
-			stage, err := newBuildStage(ctx, alias, tc.stage, image.DigestPairMap{}, opts)
+			stage, err := newBuildStage(ctx, alias, tc.stage, opts)
 			require.NoError(err)
 
-			kvStore := keyvalue.MemStore{}
+			kvStore := keyvalue.MockStore{}
 			cacheMgr := cache.New(ctx.ImageStore, kvStore, registry.NoopClientFixture())
 
 			for i, node := range stage.nodes {
diff --git a/lib/builder/step/add_copy_step.go b/lib/builder/step/add_copy_step.go
index 84a76fe1..6ed3d30c 100644
--- a/lib/builder/step/add_copy_step.go
+++ b/lib/builder/step/add_copy_step.go
@@ -90,8 +90,9 @@ func (s *addCopyStep) ContextDirs() (string, []string) {
 	return s.fromStage, s.fromPaths
 }
 
-// SetCacheID sets the cache ID of the step given a seed SHA256 value. Calculates the
-// ID randomly if copying from another stage, else checksums the file contents.
+// SetCacheID sets the cache ID of the step given a seed SHA256 value.
+// Calculates the ID randomly if copying from another stage, else checksums the
+// file contents.
 func (s *addCopyStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	if s.fromStage != "" {
 		// It is copying from a previous stage, return random bytes.
@@ -120,7 +121,8 @@ func (s *addCopyStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	return nil
 }
 
-// Execute executes the add/copy step. If modifyFS is true, actually performs the on-disk copy.
+// Execute executes the add/copy step. If modifyFS is true, actually performs
+// the on-disk copy.
 func (s *addCopyStep) Execute(ctx *context.BuildContext, modifyFS bool) (err error) {
 	sourceRoot := s.contextRootDir(ctx)
 	sources := s.resolveFromPaths(ctx)
@@ -188,7 +190,9 @@ func (s *addCopyStep) contextRootDir(ctx *context.BuildContext) string {
 	return ctx.ContextDir
 }
 
-func checksumPathContents(ctx *context.BuildContext, path string, fi os.FileInfo, checksum io.Writer) error {
+func checksumPathContents(
+	ctx *context.BuildContext, path string, fi os.FileInfo, checksum io.Writer) error {
+
 	// Skip special files.
 	if utils.IsSpecialFile(fi) {
 		if fi.IsDir() {
diff --git a/lib/builder/step/base_step.go b/lib/builder/step/base_step.go
index 37c805a6..36a991cc 100644
--- a/lib/builder/step/base_step.go
+++ b/lib/builder/step/base_step.go
@@ -66,13 +66,13 @@ func (s *baseStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	return nil
 }
 
-// SetWorkingDir set the working dir of the current step
-// Exporting the logic to this method allows for an easier `ApplyCtxAndConfig` overwriting
+// SetWorkingDir set the working dir of the current step.
+// Exporting the logic to this method allows overwriting `ApplyCtxAndConfig`.
 func (s *baseStep) SetWorkingDir(
 	ctx *context.BuildContext, imageConfig *image.Config) error {
 	s.workingDir = ctx.RootDir // Default workingDir to root.
 
-	// Set working dir from imageConfig
+	// Set working dir from imageConfig.
 	if imageConfig != nil && imageConfig.Config.WorkingDir != "" {
 		s.workingDir = os.ExpandEnv(imageConfig.Config.WorkingDir)
 	}
diff --git a/lib/builder/step/step.go b/lib/builder/step/step.go
index 098035cb..f31190a0 100644
--- a/lib/builder/step/step.go
+++ b/lib/builder/step/step.go
@@ -116,7 +116,8 @@ func NewDockerfileStep(
 		step, err = NewFromStep(s.Args, s.Image, s.Alias)
 	case *dockerfile.HealthcheckDirective:
 		s, _ := d.(*dockerfile.HealthcheckDirective)
-		step, err = NewHealthcheckStep(s.Args, s.Interval, s.Timeout, s.StartPeriod, s.Retries, s.Test, s.Commit)
+		step, err = NewHealthcheckStep(
+			s.Args, s.Interval, s.Timeout, s.StartPeriod, s.Retries, s.Test, s.Commit)
 	case *dockerfile.LabelDirective:
 		s, _ := d.(*dockerfile.LabelDirective)
 		step = NewLabelStep(s.Args, s.Labels, s.Commit)
diff --git a/lib/cache/cache_manager.go b/lib/cache/cache_manager.go
index 9a2a00bc..ff9521d8 100644
--- a/lib/cache/cache_manager.go
+++ b/lib/cache/cache_manager.go
@@ -65,13 +65,24 @@ func (manager noopCacheManager) WaitForPush() error {
 // It needs an additional key-value store for cache key/layer name lookup.
 // It implements CacheManager interface.
 type registryCacheManager struct {
-	imageStore     *storage.ImageStore
-	kvStore        keyvalue.Store
-	registryClient registry.Client
-
 	sync.Mutex
 	wg         sync.WaitGroup
 	pushErrors utils.MultiErrors
+
+	// imageStore manages local files.
+	imageStore *storage.ImageStore
+
+	// kvStore stores cache key-value pairs in an external storage.
+	kvStore keyvalue.Store
+
+	// memKVStore stores cache key-value pairs in memory. This is specifically
+	// needed for the case where a layer is still being pushed, thus kvStore
+	// doesn't contain the key-value pairs yet, but the layer can be reused for
+	// the following stage.
+	memKVStore map[string]string
+
+	// registryClient is the client for docker registry.
+	registryClient registry.Client
 }
 
 var (
@@ -80,8 +91,8 @@ var (
 	ErrorLayerNotFound = errors.Errorf("layer not found in cache")
 )
 
-// New returns a new cache manager that interacts with the registry
-// passed in as well as the local filesystem through the image store.
+// New returns a new cache manager that interacts with the registry passed in as
+// well as the local filesystem through the image store.
 // By default the registry field is left blank.
 func New(
 	imageStore *storage.ImageStore, kvStore keyvalue.Store,
@@ -94,34 +105,41 @@ func New(
 	return &registryCacheManager{
 		imageStore:     imageStore,
 		kvStore:        kvStore,
+		memKVStore:     make(map[string]string),
 		registryClient: registryClient,
 	}
 }
 
 // PullCache tries to fetch the layer corresponding to the cache ID.
 // If the layer is not found, it returns ErrorLayerNotFound.
-// This function is blocking
+// This function is blocking.
 func (manager *registryCacheManager) PullCache(cacheID string) (*image.DigestPair, error) {
 	manager.Lock()
 	defer manager.Unlock()
 
 	var entry string
 	var err error
-	for i := 0; ; i++ {
-		entry, err = manager.kvStore.Get(_cachePrefix + cacheID)
-		if err == nil && entry != "" {
-			break
-		} else if entry == "" {
-			return nil, errors.Wrapf(ErrorLayerNotFound, "find layer %s", cacheID)
-		} else {
-			if i >= 2 {
-				return nil, fmt.Errorf("query cache id %s: %s", cacheID, err)
+	key := _cachePrefix + cacheID
+	entry, ok := manager.memKVStore[key]
+	if ok {
+		log.Infof("Found mapping in cacheID mem kv store: %s => %s", cacheID, entry)
+	} else {
+		for i := 0; ; i++ {
+			entry, err = manager.kvStore.Get(key)
+			if err == nil && entry != "" {
+				break
+			} else if entry == "" {
+				return nil, errors.Wrapf(ErrorLayerNotFound, "find layer %s", cacheID)
+			} else {
+				if i >= 2 {
+					return nil, fmt.Errorf("query cache id %s: %s", cacheID, err)
+				}
+				log.Info("Retrying query for cacheID %s", cacheID)
+				time.Sleep(time.Second)
 			}
-			log.Info("Retrying query for cacheID %s", cacheID)
-			time.Sleep(time.Second)
 		}
+		log.Infof("Found mapping in cacheID kv store: %s => %s", cacheID, entry)
 	}
-	log.Infof("Found mapping in cacheID KVStore: %s => %s", cacheID, entry)
 
 	if entry == _cacheEmptyEntry {
 		return nil, nil
@@ -140,6 +158,7 @@ func (manager *registryCacheManager) PullCache(cacheID string) (*image.DigestPai
 		if manager.registryClient == nil {
 			return nil, fmt.Errorf("registry client not configured to pull cache")
 		}
+
 		// Pull layer from docker registry.
 		info, err = manager.registryClient.PullLayer(gzipDigest)
 		if err != nil {
@@ -164,6 +183,13 @@ func (manager *registryCacheManager) PullCache(cacheID string) (*image.DigestPai
 
 // PushCache tries to push an image layer asynchronously.
 func (manager *registryCacheManager) PushCache(cacheID string, digestPair *image.DigestPair) error {
+	manager.Lock()
+	defer manager.Unlock()
+
+	key := _cachePrefix + cacheID
+	entry := createEntry(digestPair)
+	manager.memKVStore[key] = entry
+
 	if manager.registryClient == nil {
 		manager.pushErrors.Add(fmt.Errorf("registry client not configured to push cache"))
 		return nil
@@ -174,9 +200,6 @@ func (manager *registryCacheManager) PushCache(cacheID string, digestPair *image
 	go func() {
 		defer manager.wg.Done()
 
-		manager.Lock()
-		defer manager.Unlock()
-
 		if digestPair != nil {
 			if err := manager.registryClient.PushLayer(digestPair.GzipDescriptor.Digest); err != nil {
 				manager.pushErrors.Add(fmt.Errorf("push layer %s: %s", digestPair.GzipDescriptor.Digest, err))
@@ -184,7 +207,9 @@ func (manager *registryCacheManager) PushCache(cacheID string, digestPair *image
 			}
 		}
 
-		entry := createEntry(digestPair)
+		manager.Lock()
+		defer manager.Unlock()
+
 		if err := manager.kvStore.Put(_cachePrefix+cacheID, entry); err != nil {
 			manager.pushErrors.Add(fmt.Errorf("store tag mapping (%s,%s): %s", cacheID, entry, err))
 			return
diff --git a/lib/cache/cache_manager_test.go b/lib/cache/cache_manager_test.go
index fea19bcb..9d93c3e2 100644
--- a/lib/cache/cache_manager_test.go
+++ b/lib/cache/cache_manager_test.go
@@ -15,15 +15,19 @@
 package cache_test
 
 import (
+	"sync"
 	"testing"
 
-	"github.com/pkg/errors"
-	"github.com/stretchr/testify/require"
 	"github.com/uber/makisu/lib/cache"
 	"github.com/uber/makisu/lib/cache/keyvalue"
 	"github.com/uber/makisu/lib/context"
 	"github.com/uber/makisu/lib/docker/image"
 	"github.com/uber/makisu/lib/registry"
+	mockregistry "github.com/uber/makisu/mocks/lib/registry"
+
+	"github.com/golang/mock/gomock"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNoopCache(t *testing.T) {
@@ -49,13 +53,13 @@ func TestNoopCache(t *testing.T) {
 	require.NoError(err)
 }
 
-func TestMemStore(t *testing.T) {
+func TestCachePushAndPull(t *testing.T) {
 	require := require.New(t)
 
 	ctx, cleanup := context.BuildContextFixture()
 	defer cleanup()
 
-	kvStore := keyvalue.MemStore{}
+	kvStore := keyvalue.MockStore{}
 	cacheMgr := cache.New(ctx.ImageStore, kvStore, registry.NoopClientFixture())
 
 	_, err := cacheMgr.PullCache("cacheid1")
@@ -75,3 +79,46 @@ func TestMemStore(t *testing.T) {
 	_, err = cacheMgr.PullCache("cacheid2")
 	require.NoError(err)
 }
+
+func TestCachePullWithOngoingPushing(t *testing.T) {
+	require := require.New(t)
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ctx, cleanup := context.BuildContextFixture()
+	defer cleanup()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	mockClient := mockregistry.NewMockClient(ctrl)
+	mockClient.EXPECT().PullLayer(image.Digest("sha256:testgzip")).Return(nil, nil)
+	mockClient.EXPECT().PushLayer(image.Digest("sha256:testgzip")).
+		Do(func(d image.Digest) {
+			wg.Done()
+			// Push doesn't return until the end of the test.
+			select {}
+		}).Return(nil)
+
+	kvStore := keyvalue.MockStore{}
+	cacheMgr := cache.New(ctx.ImageStore, kvStore, mockClient)
+
+	_, err := cacheMgr.PullCache("cacheid1")
+	require.Equal(cache.ErrorLayerNotFound, errors.Cause(err))
+
+	go func() {
+		err = cacheMgr.PushCache(
+			"cacheid2",
+			&image.DigestPair{
+				TarDigest:      image.Digest("sha256:test"),
+				GzipDescriptor: image.Descriptor{Digest: image.Digest("sha256:testgzip")},
+			},
+		)
+	}()
+
+	// Wait until layer is being pushed to make sure memKVStore is populated.
+	wg.Wait()
+	_, err = cacheMgr.PullCache("cacheid2")
+	require.NoError(err)
+}
diff --git a/lib/cache/keyvalue/http_store_test.go b/lib/cache/keyvalue/http_store_test.go
index 43215270..ad47f4c2 100644
--- a/lib/cache/keyvalue/http_store_test.go
+++ b/lib/cache/keyvalue/http_store_test.go
@@ -20,10 +20,10 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/uber/makisu/mocks/net/http"
+
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
-
-	"github.com/uber/makisu/mocks/net/http"
 )
 
 const _testURL = "http://localhost:0/test"
diff --git a/lib/cache/keyvalue/mem_store.go b/lib/cache/keyvalue/mock_store.go
similarity index 76%
rename from lib/cache/keyvalue/mem_store.go
rename to lib/cache/keyvalue/mock_store.go
index 00f230b2..03f7d113 100644
--- a/lib/cache/keyvalue/mem_store.go
+++ b/lib/cache/keyvalue/mock_store.go
@@ -14,20 +14,20 @@
 
 package keyvalue
 
-// MemStore implements Client interface. It stores cache key-value mappings
+// MockStore implements Client interface. It stores cache key-value mappings
 // in memory.
-type MemStore map[string]string
+type MockStore map[string]string
 
 // Get returns the value of a key previously set in memory.
-func (m MemStore) Get(key string) (string, error) {
+func (m MockStore) Get(key string) (string, error) {
 	return m[key], nil
 }
 
 // Put stores a key and its value in memory.
-func (m MemStore) Put(key, value string) error {
+func (m MockStore) Put(key, value string) error {
 	m[key] = value
 	return nil
 }
 
 // Cleanup does nothing, but is implemented to comply with Client interface.
-func (m MemStore) Cleanup() error { return nil }
+func (m MockStore) Cleanup() error { return nil }
diff --git a/lib/docker/image/digest.go b/lib/docker/image/digest.go
index 93367f9f..a5c48381 100644
--- a/lib/docker/image/digest.go
+++ b/lib/docker/image/digest.go
@@ -28,9 +28,6 @@ type Digest string
 // DigestPairs is a list of DigestPair
 type DigestPairs []*DigestPair
 
-// DigestPairMap is a map from string to DigestPairs
-type DigestPairMap map[string]DigestPairs
-
 // Hex returns the hex part of the digest.
 // Example:
 //   e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
diff --git a/lib/mountutils/mountutils.go b/lib/mountutils/mountutils.go
index 7b218999..fc5e77bf 100644
--- a/lib/mountutils/mountutils.go
+++ b/lib/mountutils/mountutils.go
@@ -55,7 +55,7 @@ func newMountInfo() *mountInfo {
 func (info *mountInfo) initialize() error {
 	content, err := ioutil.ReadFile(info.mountsFile)
 	if os.IsNotExist(err) {
-		log.Debug("Cannot init mountmanager, /proc/mounts does not exist")
+		log.Debug("Skipping mountmanager init, /proc/mounts does not exist")
 		return nil
 	} else if err != nil {
 		return fmt.Errorf("mountmanager initialize: %s", err)
diff --git a/mocks/lib/registry/mockclient.go b/mocks/lib/registry/mockclient.go
new file mode 100644
index 00000000..c3ec8be8
--- /dev/null
+++ b/mocks/lib/registry/mockclient.go
@@ -0,0 +1,151 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/uber/makisu/lib/registry (interfaces: Client)
+
+// Package mockregistry is a generated GoMock package.
+package mockregistry
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	image "github.com/uber/makisu/lib/docker/image"
+	os "os"
+	reflect "reflect"
+)
+
+// MockClient is a mock of Client interface
+type MockClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockClientMockRecorder
+}
+
+// MockClientMockRecorder is the mock recorder for MockClient
+type MockClientMockRecorder struct {
+	mock *MockClient
+}
+
+// NewMockClient creates a new mock instance
+func NewMockClient(ctrl *gomock.Controller) *MockClient {
+	mock := &MockClient{ctrl: ctrl}
+	mock.recorder = &MockClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockClient) EXPECT() *MockClientMockRecorder {
+	return m.recorder
+}
+
+// Pull mocks base method
+func (m *MockClient) Pull(arg0 string) (*image.DistributionManifest, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Pull", arg0)
+	ret0, _ := ret[0].(*image.DistributionManifest)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Pull indicates an expected call of Pull
+func (mr *MockClientMockRecorder) Pull(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Pull", reflect.TypeOf((*MockClient)(nil).Pull), arg0)
+}
+
+// PullImageConfig mocks base method
+func (m *MockClient) PullImageConfig(arg0 image.Digest) (os.FileInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PullImageConfig", arg0)
+	ret0, _ := ret[0].(os.FileInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PullImageConfig indicates an expected call of PullImageConfig
+func (mr *MockClientMockRecorder) PullImageConfig(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullImageConfig", reflect.TypeOf((*MockClient)(nil).PullImageConfig), arg0)
+}
+
+// PullLayer mocks base method
+func (m *MockClient) PullLayer(arg0 image.Digest) (os.FileInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PullLayer", arg0)
+	ret0, _ := ret[0].(os.FileInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PullLayer indicates an expected call of PullLayer
+func (mr *MockClientMockRecorder) PullLayer(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullLayer", reflect.TypeOf((*MockClient)(nil).PullLayer), arg0)
+}
+
+// PullManifest mocks base method
+func (m *MockClient) PullManifest(arg0 string) (*image.DistributionManifest, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PullManifest", arg0)
+	ret0, _ := ret[0].(*image.DistributionManifest)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PullManifest indicates an expected call of PullManifest
+func (mr *MockClientMockRecorder) PullManifest(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullManifest", reflect.TypeOf((*MockClient)(nil).PullManifest), arg0)
+}
+
+// Push mocks base method
+func (m *MockClient) Push(arg0 string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Push", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Push indicates an expected call of Push
+func (mr *MockClientMockRecorder) Push(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Push", reflect.TypeOf((*MockClient)(nil).Push), arg0)
+}
+
+// PushImageConfig mocks base method
+func (m *MockClient) PushImageConfig(arg0 image.Digest) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PushImageConfig", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PushImageConfig indicates an expected call of PushImageConfig
+func (mr *MockClientMockRecorder) PushImageConfig(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PushImageConfig", reflect.TypeOf((*MockClient)(nil).PushImageConfig), arg0)
+}
+
+// PushLayer mocks base method
+func (m *MockClient) PushLayer(arg0 image.Digest) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PushLayer", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PushLayer indicates an expected call of PushLayer
+func (mr *MockClientMockRecorder) PushLayer(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PushLayer", reflect.TypeOf((*MockClient)(nil).PushLayer), arg0)
+}
+
+// PushManifest mocks base method
+func (m *MockClient) PushManifest(arg0 string, arg1 *image.DistributionManifest) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PushManifest", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PushManifest indicates an expected call of PushManifest
+func (mr *MockClientMockRecorder) PushManifest(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PushManifest", reflect.TypeOf((*MockClient)(nil).PushManifest), arg0, arg1)
+}