AddLocalDirectory helper method

image_writer.go

@@ -10,6 +10,7 @@ import (
 	"math"
 	"os"
 	"path"
+	"path/filepath"
 	"runtime"
 	"strings"
 	"sync/atomic"
@@ -107,6 +108,42 @@ func (iw *ImageWriter) AddLocalFile(origin, target string) error {
 	return iw.AddFile(f, target)
 }
 
+func ensureIsDirectory(path string) error {
+	f, err := os.Open(path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	fileinfo, err := f.Stat()
+	if err != nil {
+		return err
+	}
+
+	if !fileinfo.IsDir() {
+		return fmt.Errorf("%q is not a directory", path)
+	}
+
+	return nil
+}
+
+// AddLocalDirectory adds a directory recursively to the ImageWriter's staging area.
+func (iw *ImageWriter) AddLocalDirectory(origin, target string) error {
+	if err := ensureIsDirectory(origin); err != nil {
+		return err
+	}
+
+	walkfn := func(path string, info os.FileInfo, err error) error {
+		if info.IsDir() {
+			return nil
+		}
+		relPath := path[len(origin):] // We need the path to be relative to the origin.
+		return iw.AddLocalFile(path, filepath.Join(target, relPath))
+	}
+
+	return filepath.Walk(origin, walkfn)
+}
+
 func manglePath(input string) (string, string) {
 	nonEmptySegments := splitPath(input)
 

image_writer_test.go

@@ -1,7 +1,6 @@
 package iso9660
 
 import (
-	"bytes"
 	"io/ioutil"
 	"os"
 	"path"
@@ -84,7 +83,7 @@ func TestWriterStaging(t *testing.T) {
 	assert.Equal(t, testFileContents, string(readData))
 }
 
-func TestWriter(t *testing.T) {
+func TestWriterAddLocalDirectory(t *testing.T) {
 	w, err := NewWriter()
 	assert.NoError(t, err)
 	defer func() {
@@ -93,13 +92,7 @@ func TestWriter(t *testing.T) {
 		}
 	}()
 
-	err = w.AddFile(strings.NewReader("hrh2309hr320h"), "someDirectory-Path/dir1/somefile.dat")
-	assert.NoError(t, err)
-
-	largeFileData, err := ioutil.ReadFile("fixtures/test.iso_source/dir2/large.txt")
-	assert.NoError(t, err)
-
-	err = w.AddFile(bytes.NewReader(largeFileData), "anotherDir/large.txt")
+	err = w.AddLocalDirectory("fixtures/test.iso_source", "foo")
 	assert.NoError(t, err)
 
 	f, err := ioutil.TempFile(os.TempDir(), "iso9660_golang_test")
@@ -124,16 +117,45 @@ func TestWriter(t *testing.T) {
 
 	children, err := root.GetChildren()
 	assert.NoError(t, err)
-	assert.Len(t, children, 2)
-	assert.Equal(t, "anotherdir", children[0].Name())
+	assert.Len(t, children, 1)
+	assert.Equal(t, "foo", children[0].Name())
 
 	children, err = children[0].GetChildren()
 	assert.NoError(t, err)
-	assert.Len(t, children, 1)
-	assert.Equal(t, "large.txt", children[0].Name())
+	assert.Len(t, children, 4)
+	assert.Equal(t, "dir4", children[3].Name())
+
+	children, err = children[3].GetChildren()
+	assert.NoError(t, err)
+	assert.Len(t, children, 1000)
+}
 
-	readData, err := ioutil.ReadAll(children[0].Reader())
+func TestWriterAddLocalDirectoryNonExisting(t *testing.T) {
+	w, err := NewWriter()
 	assert.NoError(t, err)
+	defer func() {
+		if cleanupErr := w.Cleanup(); cleanupErr != nil {
+			t.Fatalf("failed to cleanup writer: %v", cleanupErr)
+		}
+	}()
+
+	err = w.AddLocalDirectory("/etc/nonexistent", "foo")
+	assert.Error(t, err)
+
+	assert.True(t, os.IsNotExist(err))
+}
+
+func TestWriterAddLocalDirectoryWithFile(t *testing.T) {
+	w, err := NewWriter()
+	assert.NoError(t, err)
+	defer func() {
+		if cleanupErr := w.Cleanup(); cleanupErr != nil {
+			t.Fatalf("failed to cleanup writer: %v", cleanupErr)
+		}
+	}()
+
+	err = w.AddLocalDirectory("/etc/hosts", "foo")
+	assert.Error(t, err)
 
-	assert.Equal(t, largeFileData, readData)
+	assert.False(t, os.IsNotExist(err))
 }