Updated tests

pkg/dockerfile/dockerfile.go

@@ -123,12 +123,15 @@ func Dependencies(index int, stages []instructions.Stage, buildArgs *BuildArgs)
 				if err := util.UpdateConfigEnv(c.Env, &imageConfig.Config, replacementEnvs); err != nil {
 					return nil, err
 				}
+			case *instructions.ArgCommand:
+				buildArgs.AddArg(c.Key, c.Value)
 			case *instructions.CopyCommand:
 				if c.From != strconv.Itoa(index) {
 					continue
 				}
 				// First, resolve any environment replacement
-				resolvedEnvs, err := util.ResolveEnvironmentReplacementList(c.SourcesAndDest, imageConfig.Config.Env, true)
+				replacementEnvs := buildArgs.ReplacementEnvs(imageConfig.Config.Env)
+				resolvedEnvs, err := util.ResolveEnvironmentReplacementList(c.SourcesAndDest, replacementEnvs, true)
 				if err != nil {
 					return nil, err
 				}

pkg/dockerfile/dockerfile_test.go

@@ -89,7 +89,48 @@ func Test_Dependencies(t *testing.T) {
 	}
 
 	for index := range stages {
-		actualDeps, err := Dependencies(index, stages)
+		buildArgs := NewBuildArgs([]string{})
+		actualDeps, err := Dependencies(index, stages, buildArgs)
+		testutil.CheckErrorAndDeepEqual(t, false, err, expectedDependencies[index], actualDeps)
+	}
+}
+
+func Test_DependenciesWithArg(t *testing.T) {
+	testDir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	helloPath := filepath.Join(testDir, "hello")
+	if err := os.Mkdir(helloPath, 0755); err != nil {
+		t.Fatal(err)
+	}
+
+	dockerfile := fmt.Sprintf(`
+	FROM scratch
+	COPY %s %s
+	
+	FROM scratch AS second
+	ARG hienv
+	COPY a b
+	COPY --from=0 /$hienv %s /hi2/
+	`, helloPath, helloPath, testDir)
+
+	stages, err := Parse([]byte(dockerfile))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expectedDependencies := [][]string{
+		{
+			helloPath,
+			testDir,
+		},
+		nil,
+	}
+	buildArgs := NewBuildArgs([]string{fmt.Sprintf("hienv=%s", helloPath)})
+
+	for index := range stages {
+		actualDeps, err := Dependencies(index, stages, buildArgs)
 		testutil.CheckErrorAndDeepEqual(t, false, err, expectedDependencies[index], actualDeps)
 	}
 }