From 9df814cc86ead030e1ec8834579259734380c236 Mon Sep 17 00:00:00 2001
From: Rootul P <rootulp@gmail.com>
Date: Mon, 26 Jun 2023 16:24:53 -0400
Subject: [PATCH] chore!: check newChunk size (#172)

* chore!: check newChunk size

* test: Test_setCell

* improve tests for TestSetCell

* refactor: Test_setCell
---
 datasquare.go      | 21 ++++++++---
 datasquare_test.go | 91 +++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 97 insertions(+), 15 deletions(-)

diff --git a/datasquare.go b/datasquare.go
index 5b82e01..bbdd0d8 100644
--- a/datasquare.go
+++ b/datasquare.go
@@ -292,22 +292,31 @@ func (ds *dataSquare) GetCell(x uint, y uint) []byte {
 	return cell
 }
 
-// SetCell sets a specific cell. Cell to set must be `nil`.
-// Panics if attempting to set a cell that is not `nil`.
-func (ds *dataSquare) SetCell(x uint, y uint, newChunk []byte) {
+// SetCell sets a specific cell. The cell to set must be `nil`. Returns an error
+// if the cell to set is not `nil` or newChunk is not the correct size.
+func (ds *dataSquare) SetCell(x uint, y uint, newChunk []byte) error {
 	if ds.squareRow[x][y] != nil {
-		panic(fmt.Sprintf("cannot set cell (%d, %d) as it already has a value %x", x, y, ds.squareRow[x][y]))
+		return fmt.Errorf("cannot set cell (%d, %d) as it already has a value %x", x, y, ds.squareRow[x][y])
+	}
+	if len(newChunk) != int(ds.chunkSize) {
+		return fmt.Errorf("cannot set cell with chunk size %d because dataSquare chunk size is %d", len(newChunk), ds.chunkSize)
 	}
 	ds.squareRow[x][y] = newChunk
 	ds.squareCol[y][x] = newChunk
 	ds.resetRoots()
+	return nil
 }
 
-// setCell sets a specific cell.
-func (ds *dataSquare) setCell(x uint, y uint, newChunk []byte) {
+// setCell sets a specific cell. setCell will overwrite any existing value.
+// Returns an error if the newChunk is not the correct size.
+func (ds *dataSquare) setCell(x uint, y uint, newChunk []byte) error {
+	if len(newChunk) != int(ds.chunkSize) {
+		return fmt.Errorf("cannot set cell with chunk size %d because dataSquare chunk size is %d", len(newChunk), ds.chunkSize)
+	}
 	ds.squareRow[x][y] = newChunk
 	ds.squareCol[y][x] = newChunk
 	ds.resetRoots()
+	return nil
 }
 
 // Flattened returns the concatenated rows of the data square.
diff --git a/datasquare_test.go b/datasquare_test.go
index 90b28a4..8d0eca6 100644
--- a/datasquare_test.go
+++ b/datasquare_test.go
@@ -52,19 +52,92 @@ func TestInvalidDataSquareCreation(t *testing.T) {
 }
 
 func TestSetCell(t *testing.T) {
-	ds, err := newDataSquare([][]byte{{1}, {2}, {3}, {4}}, NewDefaultTree)
-	if err != nil {
-		panic(err)
+	type testCase struct {
+		name         string
+		originalCell []byte
+		newCell      []byte
+		wantErr      bool
 	}
 
-	// SetCell can only write to nil cells
-	assert.Panics(t, func() { ds.SetCell(0, 0, []byte{0}) })
+	testCases := []testCase{
+		{
+			name:         "can set cell if originally nil",
+			originalCell: nil,
+			newCell:      []byte{42},
+			wantErr:      false,
+		},
+		{
+			name:         "expect error if cell is not originally nil",
+			originalCell: []byte{1},
+			wantErr:      true,
+		},
+		{
+			name:         "expect error if new cell is not the correct chunk size",
+			originalCell: nil,
+			newCell:      []byte{1, 2}, // incorrect chunk size
+			wantErr:      true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ds, err := newDataSquare([][]byte{tc.originalCell, {2}, {3}, {4}}, NewDefaultTree)
+			assert.NoError(t, err)
 
-	// Set the cell to nil to allow modification
-	ds.setCell(0, 0, nil)
+			err = ds.SetCell(0, 0, tc.newCell)
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tc.newCell, ds.GetCell(0, 0))
+			}
+		})
+	}
+}
 
-	ds.SetCell(0, 0, []byte{42})
-	assert.Equal(t, []byte{42}, ds.GetCell(0, 0))
+func Test_setCell(t *testing.T) {
+	type testCase struct {
+		name         string
+		originalCell []byte
+		newCell      []byte
+		wantErr      bool
+	}
+
+	testCases := []testCase{
+		{
+			name:         "can set cell if originally nil",
+			originalCell: nil,
+			newCell:      []byte{42},
+			wantErr:      false,
+		},
+		{
+			name:         "can set cell if originally some value",
+			originalCell: []byte{1},
+			newCell:      []byte{42},
+			wantErr:      false,
+		},
+		{
+			name:         "expect error if new cell is not the correct chunk size",
+			originalCell: nil,
+			newCell:      []byte{1, 2}, // incorrect chunk size
+			wantErr:      true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ds, err := newDataSquare([][]byte{tc.originalCell, {2}, {3}, {4}}, NewDefaultTree)
+			assert.NoError(t, err)
+
+			err = ds.setCell(0, 0, tc.newCell)
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tc.newCell, ds.GetCell(0, 0))
+			}
+		})
+	}
 }
 
 func TestGetCell(t *testing.T) {