diff --git a/image_writer.go b/image_writer.go index a9cd0f1..23bba5f 100644 --- a/image_writer.go +++ b/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) diff --git a/image_writer_test.go b/image_writer_test.go index b681ae8..6bc2198 100644 --- a/image_writer_test.go +++ b/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)) }