[WIP] - Testing ReadStateBytes and WriteStateBytes (#37464)

* Fix nil pointer error

* Add WIP test for ReadStateBytes

* Move test mock to separate testing file

* Update mock to send unexpected EOF when there's a problem returning data and it's not a true EOF

* Add test case for when length != expected length

* Add test for when trying to read state from a store type that doesn't exist

* Change symbol names to lowercase

* Add ability to force a diagnostic to be returned from `mockReadStateBytesClient`'s `Recv` method

* Add test showing error diagnostics raised by the ReadStateBytes client are returned

* Add missing header

* Simplify mock by using an embedded type

* Rename `mockOpts` to `mockReadStateBytesOpts`

* Update existing tests to assert what arguments are passed to the RPC method call

* Add mock WriteStateBytesClient which uses `go.uber.org/mock/gomock` to enable assertions about calls to Send

* Add a test for WriteStateBytes that makes assertions about calls to the Send method

* Update test case to explicitly test writing data smaller than the chunk size

* Implement chunking in WriteStateBytes, add test case to assert expected chunking behaviour

* Add generated mock for Provider_WriteStateBytesClient in protocol v6

* Update tests to use new `MockProvider_WriteStateBytesClient`, remove handwritten mock

* Update code comments in test

* Add tests for diagnostics and errors returned during WriteStateBytes

* Add generated mock for Provider_ReadStateBytesClient in protocol v6, replace old mock

* Add test case for grpc errors in ReadStateBytes, fix how error is returned

* Typo in comment

* Add missing warning test, rename some test cases

Author: SarahFrench

internal/plugin6/grpc_provider.go

@@ -1510,7 +1510,7 @@ func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) (resp p
 		return resp
 	}
 
-	var buf *bytes.Buffer
+	buf := &bytes.Buffer{}
 	var expectedTotalLength int
 	for {
 		chunk, err := client.Recv()
@@ -1519,7 +1519,7 @@ func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) (resp p
 			break
 		}
 		if err != nil {
-			resp.Diagnostics = resp.Diagnostics.Append(err)
+			resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
 			break
 		}
 		resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(chunk.Diagnostics))
@@ -1576,39 +1576,52 @@ func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) (resp
 	// TODO: Configurable chunk size
 	chunkSize := 4 * 1_000_000 // 4MB
 
-	if len(r.Bytes) < chunkSize {
+	client, err := p.client.WriteStateBytes(ctx)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
+		return resp
+	}
+
+	buf := bytes.NewBuffer(r.Bytes)
+	var totalLength int64 = int64(len(r.Bytes))
+	var totalBytesProcessed int
+	for {
+		chunk := buf.Next(chunkSize)
+
+		if len(chunk) == 0 {
+			// The previous iteration read the last of the data. Now we finish up.
+			protoResp, err := client.CloseAndRecv()
+			if err != nil {
+				resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
+				return resp
+			}
+			resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
+			if resp.Diagnostics.HasErrors() {
+				return resp
+			}
+			break
+		}
+
+		// There is more data to write
 		protoReq := &proto6.WriteStateBytes_RequestChunk{
 			TypeName:    r.TypeName,
 			StateId:     r.StateId,
-			Bytes:       r.Bytes,
-			TotalLength: int64(len(r.Bytes)),
+			Bytes:       chunk,
+			TotalLength: totalLength,
 			Range: &proto6.StateRange{
-				Start: 0,
-				End:   int64(len(r.Bytes)),
+				Start: int64(totalBytesProcessed),
+				End:   int64(totalBytesProcessed + len(chunk)),
 			},
 		}
-		client, err := p.client.WriteStateBytes(ctx)
-		if err != nil {
-			resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
-			return resp
-		}
 		err = client.Send(protoReq)
 		if err != nil {
 			resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
 			return resp
 		}
-		protoResp, err := client.CloseAndRecv()
-		if err != nil {
-			resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
-			return resp
-		}
-		resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
-		if resp.Diagnostics.HasErrors() {
-			return resp
-		}
-	}
 
-	// TODO: implement chunking for state files larger than chunkSize
+		// Track progress before next iteration
+		totalBytesProcessed += len(chunk)
+	}
 
 	return resp
 }