add workers process files to completion

Author: ybirader

main.go

@@ -40,15 +40,24 @@ func (a *Archiver) ArchiveDir(root string) error {
 type FileProcessPool struct {
 	tasks    chan File
 	executor func(f File)
+	wg       *sync.WaitGroup
 }
 
 func (f *FileProcessPool) Start() {
 	for i := 0; i < 1; i++ {
+		f.wg.Add(1)
 		go f.listen()
 	}
 }
 
+func (f *FileProcessPool) Close() {
+	close(f.tasks)
+	f.wg.Wait()
+}
+
 func (f *FileProcessPool) listen() {
+	defer f.wg.Done()
+
 	for file := range f.tasks {
 		f.executor(file)
 	}

main_test.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"io/fs"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
@@ -139,17 +140,23 @@ func TestFileProcessPool(t *testing.T) {
 		assert.Equal(t, 1, fileProcessPool.PendingFiles())
 	})
 
-	t.Run("can have workers process files", func(t *testing.T) {
+	t.Run("can have workers process files to completion", func(t *testing.T) {
+		output := bytes.Buffer{}
 		executor := func(_ File) {
+			time.Sleep(5 * time.Millisecond)
+			output.WriteString("hello, world!")
 		}
 
-		fileProcessPool := &FileProcessPool{tasks: make(chan File), executor: executor}
+		fileProcessPool := &FileProcessPool{tasks: make(chan File), executor: executor, wg: new(sync.WaitGroup)}
 		fileProcessPool.Start()
 
 		info := getFileInfo(t, helloTxtFileFixture)
 		fileProcessPool.Enqueue(File{Path: helloTxtFileFixture, Info: info})
 
+		fileProcessPool.Close()
+
 		assert.Equal(t, 0, fileProcessPool.PendingFiles())
+		assert.Equal(t, "hello, world!", output.String())
 	})
 }