refactor: unify filesystem access by introducing a fileSystem interface and updating tools to use it directly, removing os.Root dependency from sandboxFs.

pkg/tools/edit.go

@@ -5,23 +5,24 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
-	"os"
 	"strings"
 )
 
 // EditFileTool edits a file by replacing old_text with new_text.
 // The old_text must exist exactly in the file.
 type EditFileTool struct {
-	allowedDir string
+	fs fileSystem
-	restrict   bool
 }
 
 // NewEditFileTool creates a new EditFileTool with optional directory restriction.
-func NewEditFileTool(allowedDir string, restrict bool) *EditFileTool {
+func NewEditFileTool(workspace string, restrict bool) *EditFileTool {
-	return &EditFileTool{
+	var fs fileSystem
-		allowedDir: allowedDir,
+	if restrict {
-		restrict:   restrict,
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
 	}
+	return &EditFileTool{fs: fs}
 }
 
 func (t *EditFileTool) Name() string {
@@ -69,29 +70,24 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
 		return ErrorResult("new_text is required")
 	}
 
-	if t.restrict {
+	if err := editFile(t.fs, path, oldText, newText); err != nil {
-		return executeInWorkspace(t.allowedDir, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			rw := &sandboxFs{root: root}
-			if err := editFile(rw, relPath, oldText, newText); err != nil {
-				return nil, err
-			}
-			return SilentResult(fmt.Sprintf("File edited: %s", path)), nil
-		})
-	}
-
-	if err := editFile(&hostFs{}, path, oldText, newText); err != nil {
 		return ErrorResult(err.Error())
 	}
 	return SilentResult(fmt.Sprintf("File edited: %s", path))
 }
 
 type AppendFileTool struct {
-	workspace string
+	fs fileSystem
-	restrict  bool
 }
 
 func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
-	return &AppendFileTool{workspace: workspace, restrict: restrict}
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &AppendFileTool{fs: fs}
 }
 
 func (t *AppendFileTool) Name() string {
@@ -130,27 +126,16 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool
 		return ErrorResult("content is required")
 	}
 
-	var rw fileReadWriter
+	if err := appendFile(t.fs, path, content); err != nil {
-	if t.restrict {
-		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			if err := appendFile(&sandboxFs{root: root}, relPath, content); err != nil {
-				return nil, err
-			}
-			return SilentResult(fmt.Sprintf("Appended to %s", path)), nil
-		})
-	}
-
-	rw = &hostFs{}
-	if err := appendFile(rw, path, content); err != nil {
 		return ErrorResult(err.Error())
 	}
 	return SilentResult(fmt.Sprintf("Appended to %s", path))
 }
 
-// editFile reads the file via rw, performs the replacement, and writes back.
+// editFile reads the file via sysFs, performs the replacement, and writes back.
-// It uses a fileReadWriter, allowing the same logic for both restricted and unrestricted modes.
+// It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes.
-func editFile(rw fileReadWriter, path, oldText, newText string) error {
+func editFile(sysFs fileSystem, path, oldText, newText string) error {
-	content, err := rw.Read(path)
+	content, err := sysFs.ReadFile(path)
 	if err != nil {
 		return err
 	}
@@ -160,18 +145,18 @@ func editFile(rw fileReadWriter, path, oldText, newText string) error {
 		return err
 	}
 
-	return rw.Write(path, newContent)
+	return sysFs.WriteFile(path, newContent)
 }
 
-// appendFile reads the existing content (if any) via rw, appends new content, and writes back.
+// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back.
-func appendFile(rw fileReadWriter, path, appendContent string) error {
+func appendFile(sysFs fileSystem, path, appendContent string) error {
-	content, err := rw.Read(path)
+	content, err := sysFs.ReadFile(path)
 	if err != nil && !errors.Is(err, fs.ErrNotExist) {
 		return err
 	}
 
 	newContent := append(content, []byte(appendContent)...)
-	return rw.Write(path, newContent)
+	return sysFs.WriteFile(path, newContent)
 }
 
 // replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText.

pkg/tools/filesystem.go

@@ -82,12 +82,17 @@ func isWithinWorkspace(candidate, workspace string) bool {
 }
 
 type ReadFileTool struct {
-	workspace string
+	fs fileSystem
-	restrict  bool
 }
 
 func NewReadFileTool(workspace string, restrict bool) *ReadFileTool {
-	return &ReadFileTool{workspace: workspace, restrict: restrict}
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &ReadFileTool{fs: fs}
 }
 
 func (t *ReadFileTool) Name() string {
@@ -117,17 +122,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
 		return ErrorResult("path is required")
 	}
 
-	if t.restrict {
+	content, err := t.fs.ReadFile(path)
-		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			content, err := (&sandboxFs{root: root}).Read(relPath)
-			if err != nil {
-				return nil, err
-			}
-			return NewToolResult(string(content)), nil
-		})
-	}
-
-	content, err := (&hostFs{}).Read(path)
 	if err != nil {
 		return ErrorResult(err.Error())
 	}
@@ -135,12 +130,17 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
 }
 
 type WriteFileTool struct {
-	workspace string
+	fs fileSystem
-	restrict  bool
 }
 
 func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool {
-	return &WriteFileTool{workspace: workspace, restrict: restrict}
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &WriteFileTool{fs: fs}
 }
 
 func (t *WriteFileTool) Name() string {
@@ -179,16 +179,7 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR
 		return ErrorResult("content is required")
 	}
 
-	if t.restrict {
+	if err := t.fs.WriteFile(path, []byte(content)); err != nil {
-		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			if err := (&sandboxFs{root: root}).Write(relPath, []byte(content)); err != nil {
-				return nil, err
-			}
-			return SilentResult(fmt.Sprintf("File written: %s", path)), nil
-		})
-	}
-
-	if err := (&hostFs{}).Write(path, []byte(content)); err != nil {
 		return ErrorResult(err.Error())
 	}
 
@@ -196,12 +187,17 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR
 }
 
 type ListDirTool struct {
-	workspace string
+	fs fileSystem
-	restrict  bool
 }
 
 func NewListDirTool(workspace string, restrict bool) *ListDirTool {
-	return &ListDirTool{workspace: workspace, restrict: restrict}
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &ListDirTool{fs: fs}
 }
 
 func (t *ListDirTool) Name() string {
@@ -231,22 +227,11 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolRes
 		path = "."
 	}
 
-	if !t.restrict {
+	entries, err := t.fs.ReadDir(path)
-		entries, err := os.ReadDir(path)
+	if err != nil {
-		if err != nil {
+		return ErrorResult(fmt.Sprintf("failed to read directory: %v", err))
-			return ErrorResult(fmt.Sprintf("failed to read directory: %v", err))
-		}
-		return formatDirEntries(entries)
 	}
-
+	return formatDirEntries(entries)
-	return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-		entries, err := fs.ReadDir(root.FS(), relPath)
-		if err != nil {
-			return nil, fmt.Errorf("failed to read directory: %w", err)
-		}
-
-		return formatDirEntries(entries), nil
-	})
 }
 
 func formatDirEntries(entries []os.DirEntry) *ToolResult {
@@ -261,17 +246,18 @@ func formatDirEntries(entries []os.DirEntry) *ToolResult {
 	return NewToolResult(result.String())
 }
 
-// fileReadWriter abstracts reading and writing files, allowing both unrestricted
+// fileSystem abstracts reading, writing, and listing files, allowing both
-// (host filesystem) and sandbox (os.Root) implementations to share the same logic.
+// unrestricted (host filesystem) and sandbox (os.Root) implementations to share the same polymorphic interface.
-type fileReadWriter interface {
+type fileSystem interface {
-	Read(path string) ([]byte, error)
+	ReadFile(path string) ([]byte, error)
-	Write(path string, data []byte) error
+	WriteFile(path string, data []byte) error
+	ReadDir(path string) ([]os.DirEntry, error)
 }
 
 // hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem.
 type hostFs struct{}
 
-func (h *hostFs) Read(path string) ([]byte, error) {
+func (h *hostFs) ReadFile(path string) ([]byte, error) {
 	content, err := os.ReadFile(path)
 	if err != nil {
 		if os.IsNotExist(err) {
@@ -285,7 +271,11 @@ func (h *hostFs) Read(path string) ([]byte, error) {
 	return content, nil
 }
 
-func (h *hostFs) Write(path string, data []byte) error {
+func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) {
+	return os.ReadDir(path)
+}
+
+func (h *hostFs) WriteFile(path string, data []byte) error {
 	dir := filepath.Dir(path)
 	if err := os.MkdirAll(dir, 0755); err != nil {
 		return fmt.Errorf("failed to create parent directories: %w", err)
@@ -307,50 +297,88 @@ func (h *hostFs) Write(path string, data []byte) error {
 	return nil
 }
 
-// sandboxFs is a sandboxed fileReadWriter that operates within an os.Root boundary.
+// sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root.
-// All paths passed to Read/Write must be relative to the root.
 type sandboxFs struct {
-	root *os.Root
+	workspace string
 }
 
-func (r *sandboxFs) Read(path string) ([]byte, error) {
+func (r *sandboxFs) execute(path string, fn func(root *os.Root, relPath string) error) error {
-	content, err := r.root.ReadFile(path)
+	if r.workspace == "" {
+		return fmt.Errorf("workspace is not defined")
+	}
+
+	root, err := os.OpenRoot(r.workspace)
 	if err != nil {
-		if os.IsNotExist(err) {
+		return fmt.Errorf("failed to open workspace: %w", err)
-			return nil, fmt.Errorf("failed to read file: file not found: %w", err)
-		}
-		// os.Root returns "escapes from parent" for paths outside the root
-		if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || strings.Contains(err.Error(), "permission denied") {
-			return nil, fmt.Errorf("failed to read file: access denied: %w", err)
-		}
-		return nil, fmt.Errorf("failed to read file: %w", err)
 	}
-	return content, nil
+	defer root.Close()
+
+	relPath, err := getSafeRelPath(r.workspace, path)
+	if err != nil {
+		return err
+	}
+
+	return fn(root, relPath)
 }
 
-func (r *sandboxFs) Write(path string, data []byte) error {
+func (r *sandboxFs) ReadFile(path string) ([]byte, error) {
-	dir := filepath.Dir(path)
+	var content []byte
-	if dir != "." && dir != "/" {
+	err := r.execute(path, func(root *os.Root, relPath string) error {
-		if err := r.root.MkdirAll(dir, 0755); err != nil {
+		fileContent, err := root.ReadFile(relPath)
-			return fmt.Errorf("failed to create parent directories: %w", err)
+		if err != nil {
+			if os.IsNotExist(err) {
+				return fmt.Errorf("failed to read file: file not found: %w", err)
+			}
+			// os.Root returns "escapes from parent" for paths outside the root
+			if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || strings.Contains(err.Error(), "permission denied") {
+				return fmt.Errorf("failed to read file: access denied: %w", err)
+			}
+			return fmt.Errorf("failed to read file: %w", err)
 		}
-	}
+		content = fileContent
+		return nil
+	})
+	return content, err
+}
 
-	// We use a "write-then-rename" pattern here to ensure an atomic write.
+func (r *sandboxFs) WriteFile(path string, data []byte) error {
-	// This prevents the target file from being left in a truncated or partial state
+	return r.execute(path, func(root *os.Root, relPath string) error {
-	// if the operation is interrupted, as the rename operation is atomic on Linux.
+		dir := filepath.Dir(relPath)
-	tmpRelPath := fmt.Sprintf("%s.%d.tmp", path, time.Now().UnixNano())
+		if dir != "." && dir != "/" {
+			if err := root.MkdirAll(dir, 0755); err != nil {
+				return fmt.Errorf("failed to create parent directories: %w", err)
+			}
+		}
 
-	if err := r.root.WriteFile(tmpRelPath, data, 0644); err != nil {
+		// We use a "write-then-rename" pattern here to ensure an atomic write.
-		r.root.Remove(tmpRelPath) // Ensure cleanup of partial/empty temp file
+		// This prevents the target file from being left in a truncated or partial state
-		return fmt.Errorf("failed to write to temp file: %w", err)
+		// if the operation is interrupted, as the rename operation is atomic on Linux.
-	}
+		tmpRelPath := fmt.Sprintf("%s.%d.tmp", relPath, time.Now().UnixNano())
 
-	if err := r.root.Rename(tmpRelPath, path); err != nil {
+		if err := root.WriteFile(tmpRelPath, data, 0644); err != nil {
-		r.root.Remove(tmpRelPath)
+			root.Remove(tmpRelPath) // Ensure cleanup of partial/empty temp file
-		return fmt.Errorf("failed to rename temp file over target: %w", err)
+			return fmt.Errorf("failed to write to temp file: %w", err)
-	}
+		}
-	return nil
+
+		if err := root.Rename(tmpRelPath, relPath); err != nil {
+			root.Remove(tmpRelPath)
+			return fmt.Errorf("failed to rename temp file over target: %w", err)
+		}
+		return nil
+	})
+}
+
+func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) {
+	var entries []os.DirEntry
+	err := r.execute(path, func(root *os.Root, relPath string) error {
+		dirEntries, err := fs.ReadDir(root.FS(), relPath)
+		if err != nil {
+			return err
+		}
+		entries = dirEntries
+		return nil
+	})
+	return entries, err
 }
 
 // Helper to get a safe relative path for os.Root usage
@@ -374,31 +402,3 @@ func getSafeRelPath(workspace, path string) (string, error) {
 
 	return rel, nil
 }
-
-// executeInWorkspace executes a function within the safety of os.Root
-func executeInWorkspace(workspace string, path string, fn func(root *os.Root, relPath string) (*ToolResult, error)) *ToolResult {
-	if workspace == "" {
-		return ErrorResult("workspace is not defined")
-	}
-
-	// 1. Open the Root
-	root, err := os.OpenRoot(workspace)
-	if err != nil {
-		return ErrorResult(fmt.Sprintf("failed to open workspace: %v", err))
-	}
-	defer root.Close()
-
-	// 2. Calculate relative path
-	relPath, err := getSafeRelPath(workspace, path)
-	if err != nil {
-		return ErrorResult(err.Error())
-	}
-
-	// 3. Execute the operation
-	result, err := fn(root, relPath)
-	if err != nil {
-		return ErrorResult(err.Error())
-	}
-
-	return result
-}

pkg/tools/filesystem_test.go

@@ -17,7 +17,7 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) {
 	testFile := filepath.Join(tmpDir, "test.txt")
 	os.WriteFile(testFile, []byte("test content"), 0o644)
 
-	tool := &ReadFileTool{}
+	tool := NewReadFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path": testFile,
@@ -44,7 +44,7 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) {
 
 // TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file
 func TestFilesystemTool_ReadFile_NotFound(t *testing.T) {
-	tool := &ReadFileTool{}
+	tool := NewReadFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path": "/nonexistent_file_12345.txt",
@@ -87,7 +87,7 @@ func TestFilesystemTool_WriteFile_Success(t *testing.T) {
 	tmpDir := t.TempDir()
 	testFile := filepath.Join(tmpDir, "newfile.txt")
 
-	tool := &WriteFileTool{}
+	tool := NewWriteFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path":    testFile,
@@ -126,7 +126,7 @@ func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) {
 	tmpDir := t.TempDir()
 	testFile := filepath.Join(tmpDir, "subdir", "newfile.txt")
 
-	tool := &WriteFileTool{}
+	tool := NewWriteFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path":    testFile,
@@ -152,7 +152,7 @@ func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) {
 
 // TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path
 func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) {
-	tool := &WriteFileTool{}
+	tool := NewWriteFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"content": "test",
@@ -168,7 +168,7 @@ func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) {
 
 // TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content
 func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) {
-	tool := &WriteFileTool{}
+	tool := NewWriteFileTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path": "/tmp/test.txt",
@@ -195,7 +195,7 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) {
 	os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644)
 	os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755)
 
-	tool := &ListDirTool{}
+	tool := NewListDirTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path": tmpDir,
@@ -219,7 +219,7 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) {
 
 // TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory
 func TestFilesystemTool_ListDir_NotFound(t *testing.T) {
-	tool := &ListDirTool{}
+	tool := NewListDirTool("", false)
 	ctx := context.Background()
 	args := map[string]any{
 		"path": "/nonexistent_directory_12345",
@@ -240,7 +240,7 @@ func TestFilesystemTool_ListDir_NotFound(t *testing.T) {
 
 // TestFilesystemTool_ListDir_DefaultPath verifies default to current directory
 func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) {
-	tool := &ListDirTool{}
+	tool := NewListDirTool("", false)
 	ctx := context.Background()
 	args := map[string]any{}
 
@@ -371,7 +371,7 @@ func TestHostRW_Read_PermissionDenied(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.Chmod(protected, 0644) // ensure cleanup
 
-	_, err = (&hostFs{}).Read(protected)
+	_, err = (&hostFs{}).ReadFile(protected)
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "access denied")
 }
@@ -380,7 +380,7 @@ func TestHostRW_Read_PermissionDenied(t *testing.T) {
 func TestHostRW_Read_Directory(t *testing.T) {
 	tmpDir := t.TempDir()
 
-	_, err := (&hostFs{}).Read(tmpDir)
+	_, err := (&hostFs{}).ReadFile(tmpDir)
 	assert.Error(t, err, "expected error when reading a directory as a file")
 }
 
@@ -395,7 +395,7 @@ func TestRootRW_Read_Directory(t *testing.T) {
 	err = root.Mkdir("subdir", 0755)
 	assert.NoError(t, err)
 
-	_, err = (&sandboxFs{root: root}).Read("subdir")
+	_, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir")
 	assert.Error(t, err, "expected error when reading a directory as a file")
 }
 
@@ -404,7 +404,7 @@ func TestHostRW_Write_ParentDirMissing(t *testing.T) {
 	tmpDir := t.TempDir()
 	target := filepath.Join(tmpDir, "a", "b", "c", "file.txt")
 
-	err := (&hostFs{}).Write(target, []byte("hello"))
+	err := (&hostFs{}).WriteFile(target, []byte("hello"))
 	assert.NoError(t, err)
 
 	data, err := os.ReadFile(target)
@@ -416,12 +416,9 @@ func TestHostRW_Write_ParentDirMissing(t *testing.T) {
 // nested parent directories automatically within the sandbox.
 func TestRootRW_Write_ParentDirMissing(t *testing.T) {
 	workspace := t.TempDir()
-	root, err := os.OpenRoot(workspace)
-	assert.NoError(t, err)
-	defer root.Close()
 
 	relPath := "x/y/z/file.txt"
-	err = (&sandboxFs{root: root}).Write(relPath, []byte("nested"))
+	err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested"))
 	assert.NoError(t, err)
 
 	data, err := os.ReadFile(filepath.Join(workspace, relPath))
@@ -435,7 +432,7 @@ func TestHostRW_Write(t *testing.T) {
 	testFile := filepath.Join(tmpDir, "atomic_test.txt")
 	testData := []byte("atomic test content")
 
-	err := (&hostFs{}).Write(testFile, testData)
+	err := (&hostFs{}).WriteFile(testFile, testData)
 	assert.NoError(t, err)
 
 	content, err := os.ReadFile(testFile)
@@ -444,7 +441,7 @@ func TestHostRW_Write(t *testing.T) {
 
 	// Verify it overwrites correctly
 	newData := []byte("new atomic content")
-	err = (&hostFs{}).Write(testFile, newData)
+	err = (&hostFs{}).WriteFile(testFile, newData)
 	assert.NoError(t, err)
 
 	content, err = os.ReadFile(testFile)
@@ -455,33 +452,36 @@ func TestHostRW_Write(t *testing.T) {
 // TestRootRW_Write verifies the rootRW.Write helper function
 func TestRootRW_Write(t *testing.T) {
 	tmpDir := t.TempDir()
-	root, err := os.OpenRoot(tmpDir)
-	assert.NoError(t, err)
-	defer root.Close()
 
 	relPath := "atomic_root_test.txt"
 	testData := []byte("atomic root test content")
 
-	erw := &sandboxFs{root: root}
+	erw := &sandboxFs{workspace: tmpDir}
-	err = erw.Write(relPath, testData)
+	err := erw.WriteFile(relPath, testData)
 	assert.NoError(t, err)
 
+	root, err := os.OpenRoot(tmpDir)
+	assert.NoError(t, err)
+	defer root.Close()
+
 	f, err := root.Open(relPath)
 	assert.NoError(t, err)
+	defer f.Close()
+
 	content, err := io.ReadAll(f)
 	assert.NoError(t, err)
-	f.Close()
 	assert.Equal(t, testData, content)
 
 	// Verify it overwrites correctly
 	newData := []byte("new root atomic content")
-	err = erw.Write(relPath, newData)
+	err = erw.WriteFile(relPath, newData)
 	assert.NoError(t, err)
 
-	f, err = root.Open(relPath)
+	f2, err := root.Open(relPath)
 	assert.NoError(t, err)
-	content, err = io.ReadAll(f)
+	defer f2.Close()
+
+	content, err = io.ReadAll(f2)
 	assert.NoError(t, err)
-	f.Close()
 	assert.Equal(t, newData, content)
 }