hse-2026/homework/03-http-server/tests/env.go

package hw3test

import (
	"fmt"
	"io"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
)

// Env contains information about all files in the test environment.
// Can be written to disk.
type Env struct {
	RootDir *EnvDir
}

// Lookup returns the node for the given path.
func (e *Env) Lookup(p string) (*EnvDir, EnvNode) {
	p = path.Clean(strings.TrimPrefix(p, "/"))
	if p == "." {
		return nil, e.RootDir
	}
	dir, file := path.Split(p)

	dirs := strings.Split(dir, "/")
	dirs = dirs[:len(dirs)-1]

	parent := e.RootDir
	for _, d := range dirs {
		nxt, ok := parent.Listing[d]
		if !ok {
			return nil, nil
		}
		parent, ok = nxt.(*EnvDir)
		if !ok {
			// not a directory
			return nil, nil
		}
	}

	nxt := parent.Listing[file]
	return parent, nxt
}

// Clone returns deep clone of an Env.
func (e *Env) Clone() *Env {
	return &Env{
		RootDir: e.RootDir.Clone().(*EnvDir),
	}
}

// EnvNode is a file/dir.
type EnvNode interface {
	// WriteToDisk persists node and its children to disk.
	WriteToDisk(path string) error

	// Stats aggregates stats of the node and its children.
	Stats(path string, stats *Stats)

	// Clone creates a deep copy of the node.
	Clone() EnvNode
}

// EnvDir is a virtual directory, that can be written to disk.
type EnvDir struct {
	Listing map[string]EnvNode
	Depth   int
}

func (d *EnvDir) WriteToDisk(p string) error {
	err := os.Mkdir(p, 0777)
	if err != nil {
		return err
	}

	for name, writter := range d.Listing {
		nxt := path.Join(p, name)
		err := writter.WriteToDisk(nxt)
		if err != nil {
			return err
		}
	}
	return nil
}

// WriteToDiskSelected keeps the size and directory structure of every initial
// file, but only materializes contents for paths that a query may read.
func (d *EnvDir) WriteToDiskSelected(p string, needed map[string]bool) error {
	return d.writeToDiskSelected(p, "", needed)
}

func (d *EnvDir) writeToDiskSelected(p, relative string, needed map[string]bool) error {
	if err := os.Mkdir(p, 0777); err != nil {
		return err
	}
	for name, node := range d.Listing {
		childPath := filepath.Join(p, name)
		childRelative := path.Join(relative, name)
		switch child := node.(type) {
		case *EnvDir:
			if err := child.writeToDiskSelected(childPath, childRelative, needed); err != nil {
				return err
			}
		case *EnvFile:
			if needed[childRelative] {
				if err := child.WriteToDisk(childPath); err != nil {
					return err
				}
			} else if err := child.WriteSparseToDisk(childPath); err != nil {
				return err
			}
		default:
			return fmt.Errorf("unsupported environment node %T", child)
		}
	}
	return nil
}

func (d *EnvDir) Stats(p string, stats *Stats) {
	stats.Dirs++
	stats.DirPaths = append(stats.DirPaths, p)
	for name, node := range d.Listing {
		node.Stats(path.Join(p, name), stats)
	}
}

func (d *EnvDir) Clone() EnvNode {
	newDir := &EnvDir{
		Listing: map[string]EnvNode{},
		Depth:   d.Depth,
	}

	for name, node := range d.Listing {
		newDir.Listing[name] = node.Clone()
	}
	return newDir
}

func (d *EnvDir) CreateDir(name string) (dir *EnvDir, exist bool) {
	_, exist = d.Listing[name]
	if exist {
		return nil, exist
	}

	newDir := &EnvDir{
		Listing: map[string]EnvNode{},
		Depth:   d.Depth + 1,
	}
	d.Listing[name] = newDir
	return newDir, false
}

// EnvFile is a virtual file, that can be written to disk.
type EnvFile struct {
	GenSeed  int64
	Size     int64
	TextOnly bool
}

func (f *EnvFile) Open() io.Reader {
	var gen io.Reader = rand.New(rand.NewSource(f.GenSeed))
	if f.TextOnly {
		gen = &TextReader{gen}
	}
	gen = io.LimitReader(gen, f.Size)
	return gen
}

func (f *EnvFile) WriteToDisk(p string) error {
	file, err := os.OpenFile(p, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		return err
	}
	defer file.Close()

	reader := f.Open()
	_, err = io.CopyBuffer(file, reader, make([]byte, 256*1024))
	return err
}

func (f *EnvFile) WriteSparseToDisk(p string) error {
	file, err := os.OpenFile(p, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0666)
	if err != nil {
		return err
	}
	defer file.Close()
	return file.Truncate(f.Size)
}

func (f *EnvFile) Stats(p string, stats *Stats) {
	stats.Files++
	stats.Size += f.Size
	stats.FilePaths = append(stats.FilePaths, p)
}

func (f *EnvFile) Clone() EnvNode {
	return &EnvFile{
		GenSeed:  f.GenSeed,
		Size:     f.Size,
		TextOnly: f.TextOnly,
	}
}

// Stats is a helper for listing all files/dirs in environment.
type Stats struct {
	// Number of directories in the environment.
	Files int

	// Count of directories in the environment.
	Dirs int

	// Size of all files in the environment.
	Size int64

	// FilePaths to all files in the environment.
	FilePaths []string

	// DirPaths to all directories in the environment.
	DirPaths []string
}

// Normalize will fix stats to be deterministic and don't contain root dir.
func (s *Stats) Normalize() {
	sort.Strings(s.FilePaths)
	sort.Strings(s.DirPaths)

	if len(s.DirPaths) > 0 && s.DirPaths[0] == "" {
		s.DirPaths = s.DirPaths[1:]
	}
}