From 01329d5ac16043c0f7aa45a4be3202302bf695d5 Mon Sep 17 00:00:00 2001
From: Andrew Rynhard <andrew@andrewrynhard.com>
Date: Mon, 10 Dec 2018 11:34:06 -0800
Subject: [PATCH] Fix intermediate layer caching (#474)

* Fix intermediate layer caching

* Move the if statement into the ShouldTakeSnapshot function.

Also add some unit tests.
---
 pkg/executor/build.go      |  14 +++--
 pkg/executor/build_test.go | 106 +++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 4 deletions(-)

diff --git a/pkg/executor/build.go b/pkg/executor/build.go
index bd151d2b13..f2dd9c78f6 100644
--- a/pkg/executor/build.go
+++ b/pkg/executor/build.go
@@ -279,13 +279,18 @@ func (s *stageBuilder) takeSnapshot(files []string) (string, error) {
 func (s *stageBuilder) shouldTakeSnapshot(index int, files []string) bool {
 	isLastCommand := index == len(s.stage.Commands)-1
 
-	// We only snapshot the very end of intermediate stages.
-	if !s.stage.Final {
+	// We only snapshot the very end with single snapshot mode on.
+	if s.opts.SingleSnapshot {
 		return isLastCommand
 	}
 
-	// We only snapshot the very end with single snapshot mode on.
-	if s.opts.SingleSnapshot {
+	// Always take snapshots if we're using the cache.
+	if s.opts.Cache {
+		return true
+	}
+
+	// We only snapshot the very end of intermediate stages.
+	if !s.stage.Final {
 		return isLastCommand
 	}
 
@@ -298,6 +303,7 @@ func (s *stageBuilder) shouldTakeSnapshot(index int, files []string) bool {
 	if len(files) == 0 {
 		return false
 	}
+
 	return true
 }
 
diff --git a/pkg/executor/build_test.go b/pkg/executor/build_test.go
index 63dbd8a4b1..e212096f1a 100644
--- a/pkg/executor/build_test.go
+++ b/pkg/executor/build_test.go
@@ -19,6 +19,8 @@ package executor
 import (
 	"testing"
 
+	"github.com/moby/buildkit/frontend/dockerfile/instructions"
+
 	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/dockerfile"
 	"github.com/GoogleContainerTools/kaniko/testutil"
@@ -74,3 +76,107 @@ func stage(t *testing.T, d string) config.KanikoStage {
 		Stage: stages[0],
 	}
 }
+
+type MockCommand struct {
+	name string
+}
+
+func (m *MockCommand) Name() string {
+	return m.name
+}
+
+func Test_stageBuilder_shouldTakeSnapshot(t *testing.T) {
+	commands := []instructions.Command{
+		&MockCommand{name: "command1"},
+		&MockCommand{name: "command2"},
+		&MockCommand{name: "command3"},
+	}
+
+	stage := instructions.Stage{
+		Commands: commands,
+	}
+
+	type fields struct {
+		stage config.KanikoStage
+		opts  *config.KanikoOptions
+	}
+	type args struct {
+		index int
+		files []string
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			name: "final stage not last command",
+			fields: fields{
+				stage: config.KanikoStage{
+					Final: true,
+					Stage: stage,
+				},
+			},
+			args: args{
+				index: 1,
+			},
+			want: true,
+		},
+		{
+			name: "not final stage last command",
+			fields: fields{
+				stage: config.KanikoStage{
+					Final: false,
+					Stage: stage,
+				},
+			},
+			args: args{
+				index: len(commands) - 1,
+			},
+			want: true,
+		},
+		{
+			name: "not final stage not last command",
+			fields: fields{
+				stage: config.KanikoStage{
+					Final: false,
+					Stage: stage,
+				},
+			},
+			args: args{
+				index: 0,
+			},
+			want: false,
+		},
+		{
+			name: "caching enabled intermediate container",
+			fields: fields{
+				stage: config.KanikoStage{
+					Final: false,
+					Stage: stage,
+				},
+				opts: &config.KanikoOptions{Cache: true},
+			},
+			args: args{
+				index: 0,
+			},
+			want: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			if tt.fields.opts == nil {
+				tt.fields.opts = &config.KanikoOptions{}
+			}
+			s := &stageBuilder{
+				stage: tt.fields.stage,
+				opts:  tt.fields.opts,
+			}
+			if got := s.shouldTakeSnapshot(tt.args.index, tt.args.files); got != tt.want {
+				t.Errorf("stageBuilder.shouldTakeSnapshot() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}