diff --git a/codec_test.go b/codec_test.go
index e6b1f49..4ada166 100644
--- a/codec_test.go
+++ b/codec_test.go
@@ -14,13 +14,13 @@ var (
 
 func BenchmarkEncoding(b *testing.B) {
 	// generate some fake data
-	data := generateRandData(128)
+	data := generateRandData(128, shareSize)
 	for codecName, codec := range codecs {
 		// For some implementations we want to ensure the encoder for this data length
 		// is already cached and initialized. For this run with same sized arbitrary data.
-		_, _ = codec.Encode(generateRandData(128))
+		_, _ = codec.Encode(generateRandData(128, shareSize))
 		b.Run(
-			fmt.Sprintf("%s 128 shares", codecName),
+			fmt.Sprintf("%s 128 shares %d", codecName, shareSize),
 			func(b *testing.B) {
 				for n := 0; n < b.N; n++ {
 					encodedData, err := codec.Encode(data)
@@ -34,10 +34,10 @@ func BenchmarkEncoding(b *testing.B) {
 	}
 }
 
-func generateRandData(count int) [][]byte {
+func generateRandData(count int, chunkSize int) [][]byte {
 	out := make([][]byte, count)
 	for i := 0; i < count; i++ {
-		randData := make([]byte, count)
+		randData := make([]byte, chunkSize)
 		_, err := cryptorand.Read(randData)
 		if err != nil {
 			panic(err)
@@ -52,11 +52,11 @@ func BenchmarkDecoding(b *testing.B) {
 	for codecName, codec := range codecs {
 		// For some implementations we want to ensure the encoder for this data length
 		// is already cached and initialized. For this run with same sized arbitrary data.
-		_, _ = codec.Decode(generateMissingData(128, codec))
+		_, _ = codec.Decode(generateMissingData(128, shareSize, codec))
 
-		data := generateMissingData(128, codec)
+		data := generateMissingData(128, shareSize, codec)
 		b.Run(
-			fmt.Sprintf("%s 128 shares", codecName),
+			fmt.Sprintf("%s 128 shares %d", codecName, shareSize),
 			func(b *testing.B) {
 				for n := 0; n < b.N; n++ {
 					decodedData, err := codec.Decode(data)
@@ -70,8 +70,8 @@ func BenchmarkDecoding(b *testing.B) {
 	}
 }
 
-func generateMissingData(count int, codec Codec) [][]byte {
-	randData := generateRandData(count)
+func generateMissingData(count int, chunkSize int, codec Codec) [][]byte {
+	randData := generateRandData(count, chunkSize)
 	encoded, err := codec.Encode(randData)
 	if err != nil {
 		panic(err)
diff --git a/datasquare_test.go b/datasquare_test.go
index 45a449a..928f4fa 100644
--- a/datasquare_test.go
+++ b/datasquare_test.go
@@ -403,8 +403,7 @@ func Test_setColSlice(t *testing.T) {
 
 func BenchmarkEDSRoots(b *testing.B) {
 	for i := 32; i < 513; i *= 2 {
-		chunkSize := uint(256)
-		square, err := newDataSquare(genRandDS(i*2, int(chunkSize)), NewDefaultTree, chunkSize)
+		square, err := newDataSquare(genRandDS(i*2, int(shareSize)), NewDefaultTree, shareSize)
 		if err != nil {
 			b.Errorf("Failure to create square of size %d: %s", i, err)
 		}
diff --git a/extendeddatasquare_test.go b/extendeddatasquare_test.go
index ddb3a93..0815418 100644
--- a/extendeddatasquare_test.go
+++ b/extendeddatasquare_test.go
@@ -278,7 +278,6 @@ var dump *ExtendedDataSquare
 // BenchmarkExtension benchmarks extending datasquares sizes 4-128 using all
 // supported codecs (encoding only)
 func BenchmarkExtensionEncoding(b *testing.B) {
-	chunkSize := 256
 	for i := 4; i < 513; i *= 2 {
 		for codecName, codec := range codecs {
 			if codec.MaxChunks() < i*i {
@@ -286,7 +285,7 @@ func BenchmarkExtensionEncoding(b *testing.B) {
 				continue
 			}
 
-			square := genRandDS(i, chunkSize)
+			square := genRandDS(i, shareSize)
 			b.Run(
 				fmt.Sprintf("%s %dx%dx%d ODS", codecName, i, i, len(square[0])),
 				func(b *testing.B) {
@@ -306,7 +305,6 @@ func BenchmarkExtensionEncoding(b *testing.B) {
 // BenchmarkExtension benchmarks extending datasquares sizes 4-128 using all
 // supported codecs (both encoding and root computation)
 func BenchmarkExtensionWithRoots(b *testing.B) {
-	chunkSize := 256
 	for i := 4; i < 513; i *= 2 {
 		for codecName, codec := range codecs {
 			if codec.MaxChunks() < i*i {
@@ -314,7 +312,7 @@ func BenchmarkExtensionWithRoots(b *testing.B) {
 				continue
 			}
 
-			square := genRandDS(i, chunkSize)
+			square := genRandDS(i, shareSize)
 			b.Run(
 				fmt.Sprintf("%s %dx%dx%d ODS", codecName, i, i, len(square[0])),
 				func(b *testing.B) {