From 3993ce191a9f45ba6e35664f655860c64411a2c5 Mon Sep 17 00:00:00 2001
From: Evan Burkey <evan@burkey.co>
Date: Sat, 5 Aug 2023 12:01:26 -0700
Subject: [PATCH] add directory size calculations

---
 dir.go      | 43 ++++++++++++++++++++++++++++++++++++++++++-
 vfs_test.go | 52 +++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/dir.go b/dir.go
index eac48c7..3476056 100644
--- a/dir.go
+++ b/dir.go
@@ -20,9 +20,14 @@ type directoryHandle struct {
 }
 
 func (d *directoryHandle) Stat() (fs.FileInfo, error) {
+	sz, err := d.dir.Size()
+	if err != nil {
+		return nil, err
+	}
+
 	return &fileInfo{
 		name: d.dir.name,
-		size: 4096,
+		size: sz,
 		mode: d.dir.mode | fs.ModeDir,
 	}, nil
 }
@@ -34,3 +39,39 @@ func (d *directoryHandle) Read(data []byte) (int, error) {
 func (d *directoryHandle) Close() error {
 	return nil
 }
+
+func (d *directory) Size() (int64, error) {
+	sum := int64(0)
+	for _, e := range d.children {
+		switch v := e.(type) {
+		case *directory:
+			dir := v
+			sz, err := dir.Size()
+			if err != nil {
+				return 0, err
+			}
+			sum += sz
+
+		case *MockFile:
+			f := v
+			f.open = true
+			stat, err := f.Stat()
+			if err != nil {
+				return 0, err
+			}
+			sum += stat.Size()
+			f.open = false
+
+		case *File:
+			f := v
+			f.open = true
+			stat, err := f.Stat()
+			if err != nil {
+				return 0, err
+			}
+			sum += stat.Size()
+			f.open = false
+		}
+	}
+	return sum, nil
+}
diff --git a/vfs_test.go b/vfs_test.go
index f9798ea..6053970 100644
--- a/vfs_test.go
+++ b/vfs_test.go
@@ -44,13 +44,13 @@ func TestVFS_Basic(t *testing.T) {
 
 func TestVFS_MockFile(t *testing.T) {
 	vfs := NewVFS()
-	path := "test.txt"
+	p := "test.txt"
 	sz := int64(8096)
 
-	err := vfs.CreateMockFile(path, sz, 0666)
+	err := vfs.CreateMockFile(p, sz, 0666)
 	assert.NoError(t, err)
 
-	fp, err := vfs.Open(path)
+	fp, err := vfs.Open(p)
 	assert.NoError(t, err)
 
 	stat, err := fp.Stat()
@@ -60,3 +60,49 @@ func TestVFS_MockFile(t *testing.T) {
 	err = fp.Close()
 	assert.NoError(t, err)
 }
+
+func TestVFS_DirSize(t *testing.T) {
+	vfs := NewVFS()
+
+	err := vfs.MkdirAll("root/one", 0777)
+	assert.NoError(t, err)
+
+	err = vfs.MkdirAll("root/two", 0777)
+	assert.NoError(t, err)
+
+	files := []string{
+		"root/one/f1",
+		"root/one/f2",
+		"root/two/f1",
+		"root/two/f2",
+		"root/two/f3",
+		"root/two/f4",
+	}
+
+	sz := int64(1)
+	for _, file := range files {
+		err = vfs.CreateMockFile(file, sz, 0666)
+		assert.NoError(t, err)
+
+	}
+
+	fp, err := vfs.Open("root/one")
+	assert.NoError(t, err)
+
+	stat, err := fp.Stat()
+	assert.NoError(t, err)
+	assert.Equal(t, stat.Size(), int64(2))
+
+	err = fp.Close()
+	assert.NoError(t, err)
+
+	fp, err = vfs.Open("root/two")
+	assert.NoError(t, err)
+
+	stat, err = fp.Stat()
+	assert.NoError(t, err)
+	assert.Equal(t, stat.Size(), int64(4))
+
+	err = fp.Close()
+	assert.NoError(t, err)
+}