feat: testutil generator enhancements

Author: rvagg

testutil/generator.go

@@ -7,7 +7,6 @@ import (
 	"math/big"
 	"sort"
 	"strings"
-	"testing"
 
 	"github.com/ipfs/go-cid"
 	"github.com/ipfs/go-unixfsnode/data/builder"
@@ -23,7 +22,7 @@ import (
 // GenerateFile generates a random unixfs file of the given size, storing the
 // blocks in the provided LinkSystem and returns a DirEntry representation of
 // the file.
-func GenerateFile(t *testing.T, linkSys *linking.LinkSystem, randReader io.Reader, size int) DirEntry {
+func GenerateFile(t require.TestingT, linkSys *linking.LinkSystem, randReader io.Reader, size int) DirEntry {
 	// a file of `size` random bytes, packaged into unixfs DAGs, stored in the remote blockstore
 	delimited := io.LimitReader(randReader, int64(size))
 	var buf bytes.Buffer
@@ -54,15 +53,15 @@ func GenerateFile(t *testing.T, linkSys *linking.LinkSystem, randReader io.Reade
 // the root directory will be built as HAMT sharded (with a low "width" to
 // maximise the chance of collisions and therefore greater depth for smaller
 // number of files).
-func GenerateDirectory(t *testing.T, linkSys *linking.LinkSystem, randReader io.Reader, targetSize int, rootSharded bool) DirEntry {
+func GenerateDirectory(t require.TestingT, linkSys *linking.LinkSystem, randReader io.Reader, targetSize int, rootSharded bool) DirEntry {
 	return GenerateDirectoryFrom(t, linkSys, randReader, targetSize, "", rootSharded)
 }
 
 // GenerateDirectoryFrom is the same as GenerateDirectory but allows the caller
 // to specify a directory path to start from. This is useful for generating
 // nested directories.
 func GenerateDirectoryFrom(
-	t *testing.T,
+	t require.TestingT,
 	linkSys *linking.LinkSystem,
 	randReader io.Reader,
 	targetSize int,
@@ -91,7 +90,8 @@ func GenerateDirectoryFrom(
 					break
 				}
 			}
-			child := GenerateDirectoryFrom(t, linkSys, randReader, targetSize-curSize, dir+"/"+newDir, false)
+			sharded := rndInt(randReader, 6) == 0
+			child := GenerateDirectoryFrom(t, linkSys, randReader, targetSize-curSize, dir+"/"+newDir, sharded)
 			children = append(children, child)
 			curSize += int(child.TSize)
 		default: // 4 in 6 chance of making a new file
@@ -124,12 +124,46 @@ func GenerateDirectoryFrom(
 	return dirEntry
 }
 
+// GenerateDirectoryFor will build a directory with name `dir` whose children
+// are generated by the caller using the provided mkchild function. This is
+// useful for generating directories with a specific structure where you need
+// to have greater control over the children.
+func GenerateDirectoryFor(
+	t require.TestingT,
+	linkSys *linking.LinkSystem,
+	randReader io.Reader,
+	dir string,
+	sharded bool,
+	mkchild func(name string) *DirEntry,
+) DirEntry {
+	children := make([]DirEntry, 0)
+	for {
+		var name string
+		for {
+			var err error
+			name, err = namegen.RandomDirectoryName(randReader)
+			require.NoError(t, err)
+			if !isDupe(children, name) {
+				break
+			}
+		}
+		child := mkchild(dir + "/" + name)
+		if child == nil {
+			break
+		}
+		children = append(children, *child)
+	}
+	dirEntry := BuildDirectory(t, linkSys, children, sharded)
+	dirEntry.Path = dir
+	return dirEntry
+}
+
 // BuildDirectory builds a directory from the given children, storing the
 // blocks in the provided LinkSystem and returns a DirEntry representation of
 // the directory. If sharded is true, the root directory will be built as HAMT
 // sharded (with a low "width" to maximise the chance of collisions and
 // therefore greater depth for smaller number of files).
-func BuildDirectory(t *testing.T, linkSys *linking.LinkSystem, children []DirEntry, sharded bool) DirEntry {
+func BuildDirectory(t require.TestingT, linkSys *linking.LinkSystem, children []DirEntry, sharded bool) DirEntry {
 	// create stable sorted children, which should match the encoded form
 	// in dag-pb
 	sort.Slice(children, func(i, j int) bool {
@@ -197,8 +231,16 @@ func cidCollector(ls *ipld.LinkSystem, cids *[]cid.Cid) (ipld.BlockWriteOpener,
 }
 
 func isDupe(children []DirEntry, name string) bool {
+	if strings.Contains(name, ".") {
+		name = name[:strings.LastIndex(name, ".")]
+	}
 	for _, child := range children {
-		if strings.HasSuffix(child.Path, "/"+name) {
+		childName := child.Path[strings.LastIndex(child.Path, "/")+1:]
+		// remove suffix
+		if strings.Contains(childName, ".") {
+			childName = childName[:strings.LastIndex(childName, ".")]
+		}
+		if childName == name {
 			return true
 		}
 	}

testutil/namegen/namegen.go

@@ -33,11 +33,21 @@ func RandomFileName(r io.Reader) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	extIndex, err := getRandomIndex(r, len(extensions))
+	ext, err := RandomFileExtension(r)
 	if err != nil {
 		return "", err
 	}
-	return words[wordIndex] + extensions[extIndex], nil
+	return words[wordIndex] + ext, nil
+}
+
+// RandomFileExtension returns a random file extension, including '.'. This may
+// also return an empty string.
+func RandomFileExtension(r io.Reader) (string, error) {
+	index, err := getRandomIndex(r, len(extensions))
+	if err != nil {
+		return "", err
+	}
+	return extensions[index], nil
 }
 
 const wordData = `jabberwocky Snark whiffling borogoves mome raths brillig slithy toves outgrabe