Use iteration on map to baseline vfs from map

Author: sheetalkamat

internal/execute/testsys_test.go

@@ -1,7 +1,6 @@
 package execute_test
 
 import (
-	"errors"
 	"fmt"
 	"io"
 	"io/fs"
@@ -385,43 +384,22 @@ func (s *testSys) baselineFSwithDiff(baseline io.Writer) {
 
 	testFs := s.testFs()
 	diffs := map[string]string{}
-	err := s.fsFromFileMap().WalkDir("/", func(path string, d vfs.DirEntry, e error) error {
-		if e != nil {
-			return e
-		}
 
-		fileInfo := d.Type()
-		if fileInfo&fs.ModeSymlink != 0 {
+	for path, file := range s.mapFs().Entries() {
+		if file.Mode&fs.ModeSymlink != 0 {
 			target, ok := s.mapFs().GetTargetOfSymlink(path)
 			if !ok {
 				panic("Failed to resolve symlink target: " + path)
 			}
 			newEntry := &diffEntry{symlinkTarget: target}
 			snap[path] = newEntry
 			s.addFsEntryDiff(diffs, newEntry, path)
-			return nil
-		}
-
-		if !fileInfo.IsRegular() {
-			return nil
-		}
-
-		newContents, ok := s.fsFromFileMap().ReadFile(path)
-		if !ok {
-			return nil
-		}
-		stat := s.fsFromFileMap().Stat(path)
-		if stat == nil {
-			panic("stat is nil: " + path)
+			continue
+		} else if file.Mode.IsRegular() {
+			newEntry := &diffEntry{content: string(file.Data), mTime: file.ModTime, isWritten: testFs.writtenFiles.Has(path)}
+			snap[path] = newEntry
+			s.addFsEntryDiff(diffs, newEntry, path)
 		}
-		newEntry := &diffEntry{content: newContents, mTime: stat.ModTime(), isWritten: testFs.writtenFiles.Has(path)}
-		snap[path] = newEntry
-		s.addFsEntryDiff(diffs, newEntry, path)
-
-		return nil
-	})
-	if err != nil && !errors.Is(err, fs.ErrNotExist) {
-		panic("walkdir error during diff: " + err.Error())
 	}
 	if s.serializedDiff != nil {
 		for path := range s.serializedDiff.snap {

internal/vfs/vfstest/vfstest.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"iter"
 	"maps"
 	"path"
 	"slices"
@@ -550,6 +551,26 @@ func (m *MapFS) GetTargetOfSymlink(path string) (string, bool) {
 	return "", false
 }
 
+func (m *MapFS) Entries() iter.Seq2[string, *fstest.MapFile] {
+	return func(yield func(string, *fstest.MapFile) bool) {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		inputKeys := slices.Collect(maps.Keys(m.m))
+		slices.SortFunc(inputKeys, comparePathsByParts)
+
+		for _, p := range inputKeys {
+			file := m.m[p]
+			path := file.Sys.(*sys).realpath
+			if !tspath.PathIsAbsolute(path) {
+				path = "/" + path
+			}
+			if !yield(path, file) {
+				break
+			}
+		}
+	}
+}
+
 func must[T any](v T, err error) T {
 	if err != nil {
 		panic(err)