PSS: Implement `Get` and `Put` methods on remote grpc state, using Read/WriteStateBytes RPCs (#37717)

* Implement Get method on remote gRPC state client, add tests

* Implement Put method on remote gRPC state client, add tests

internal/states/remote/remote_grpc.go

@@ -54,7 +54,24 @@ type grpcClient struct {
 //
 // Implementation of remote.Client
 func (g *grpcClient) Get() (*Payload, tfdiags.Diagnostics) {
-	panic("not implemented yet")
+	req := providers.ReadStateBytesRequest{
+		TypeName: g.typeName,
+		StateId:  g.stateId,
+	}
+	resp := g.provider.ReadStateBytes(req)
+
+	if len(resp.Bytes) == 0 {
+		// No state to return
+		return nil, resp.Diagnostics
+	}
+
+	// TODO: Remove or replace use of MD5?
+	// The MD5 value here is never used.
+	payload := &Payload{
+		Data: resp.Bytes,
+		MD5:  []byte{}, // empty, as this is unused downstream
+	}
+	return payload, resp.Diagnostics
 }
 
 // Put invokes the WriteStateBytes gRPC method in the plugin protocol
@@ -62,7 +79,14 @@ func (g *grpcClient) Get() (*Payload, tfdiags.Diagnostics) {
 //
 // Implementation of remote.Client
 func (g *grpcClient) Put(state []byte) tfdiags.Diagnostics {
-	panic("not implemented yet")
+	req := providers.WriteStateBytesRequest{
+		TypeName: g.typeName,
+		StateId:  g.stateId,
+		Bytes:    state,
+	}
+	resp := g.provider.WriteStateBytes(req)
+
+	return resp.Diagnostics
 }
 
 // Delete invokes the DeleteState gRPC method in the plugin protocol

internal/states/remote/remote_grpc_test.go

@@ -4,16 +4,228 @@
 package remote
 
 import (
+	"bytes"
+	"strings"
 	"testing"
 
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/providers"
 	testing_provider "github.com/hashicorp/terraform/internal/providers/testing"
+	"github.com/hashicorp/terraform/internal/states"
+	"github.com/hashicorp/terraform/internal/states/statefile"
+	"github.com/hashicorp/terraform/internal/tfdiags"
+	"github.com/zclconf/go-cty/cty"
 )
 
+// Testing grpcClient's Get method, via the state manager made using a grpcClient.
+// The RefreshState method on a state manager calls the Get method of the underlying client.
+func Test_grpcClient_Get(t *testing.T) {
+	typeName := "foo_bar" // state store 'bar' in provider 'foo'
+	stateId := "production"
+	stateString := `{
+    "version": 4,
+    "terraform_version": "0.13.0",
+    "serial": 0,
+    "lineage": "",
+    "outputs": {
+        "foo": {
+            "value": "bar",
+            "type": "string"
+        }
+    }
+}`
+
+	t.Run("state manager made using grpcClient returns expected state", func(t *testing.T) {
+		provider := testing_provider.MockProvider{
+			// Mock a provider and internal state store that
+			// have both been configured
+			ConfigureProviderCalled:   true,
+			ConfigureStateStoreCalled: true,
+
+			// Check values received by the provider from the Get method.
+			ReadStateBytesFn: func(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+				if req.TypeName != typeName || req.StateId != stateId {
+					t.Fatalf("expected provider ReadStateBytes method to receive TypeName %q and StateId %q, instead got TypeName %q and StateId %q",
+						typeName,
+						stateId,
+						req.TypeName,
+						req.StateId)
+				}
+				return providers.ReadStateBytesResponse{
+					Bytes: []byte(stateString),
+					// no diags
+				}
+			},
+		}
+
+		// This package will be consumed in a statemgr.Full, so we test using NewRemoteGRPC
+		// and invoke the method on that interface that uses Get.
+		c := NewRemoteGRPC(&provider, typeName, stateId)
+
+		err := c.RefreshState() // Calls Get
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		if !provider.ReadStateBytesCalled {
+			t.Fatal("expected remote grpc state manager's RefreshState method to, via Get, call ReadStateBytes method on underlying provider, but it has not been called")
+		}
+		s := c.State()
+		v, ok := s.RootOutputValues["foo"]
+		if !ok {
+			t.Fatal("state manager doesn't contain the state returned by the mock")
+		}
+		if v.Value.AsString() != "bar" {
+			t.Fatal("state manager doesn't contain the correct output value in the state")
+		}
+	})
+
+	t.Run("state manager made using grpcClient returns expected error from error diagnostic", func(t *testing.T) {
+		var diags tfdiags.Diagnostics
+		diags = diags.Append(&hcl.Diagnostic{
+			Severity: hcl.DiagError,
+			Summary:  "error forced from test",
+			Detail:   "error forced from test",
+		})
+		provider := testing_provider.MockProvider{
+			// Mock a provider and internal state store that
+			// have both been configured
+			ConfigureProviderCalled:   true,
+			ConfigureStateStoreCalled: true,
+
+			// Force an error diagnostic
+			ReadStateBytesFn: func(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+				return providers.ReadStateBytesResponse{
+					// we don't expect state to accompany an error, but this test shows that
+					// if an error us present amy state returned is ignored.
+					Bytes:       []byte(stateString),
+					Diagnostics: diags,
+				}
+			},
+		}
+
+		// This package will be consumed in a statemgr.Full, so we test using NewRemoteGRPC
+		// and invoke the method on that interface that uses Get.
+		c := NewRemoteGRPC(&provider, typeName, stateId)
+
+		err := c.RefreshState() // Calls Get
+		if err == nil {
+			t.Fatal("expected an error but got none")
+		}
+
+		if !provider.ReadStateBytesCalled {
+			t.Fatal("expected remote grpc state manager's RefreshState method to, via Get, call ReadStateBytes method on underlying provider, but it has not been called")
+		}
+		s := c.State()
+		if s != nil {
+			t.Fatalf("expected refresh to fail due to error diagnostic, but state has been refreshed: %s", s.String())
+		}
+	})
+}
+
+// Testing grpcClient's Put method, via the state manager made using a grpcClient.
+// The PersistState method on a state manager calls the Put method of the underlying client.
+func Test_grpcClient_Put(t *testing.T) {
+	typeName := "foo_bar" // state store 'bar' in provider 'foo'
+	stateId := "production"
+
+	// State with 1 output
+	s := states.NewState()
+	s.SetOutputValue(addrs.AbsOutputValue{
+		Module:      addrs.RootModuleInstance,
+		OutputValue: addrs.OutputValue{Name: "foo"},
+	}, cty.StringVal("bar"), false)
+
+	t.Run("state manager made using grpcClient writes the expected state", func(t *testing.T) {
+		provider := testing_provider.MockProvider{
+			// Mock a provider and internal state store that
+			// have both been configured
+			ConfigureProviderCalled:   true,
+			ConfigureStateStoreCalled: true,
+
+			// Check values received by the provider from the Put method.
+			WriteStateBytesFn: func(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+				if req.TypeName != typeName || req.StateId != stateId {
+					t.Fatalf("expected provider WriteStateBytes method to receive TypeName %q and StateId %q, instead got TypeName %q and StateId %q",
+						typeName,
+						stateId,
+						req.TypeName,
+						req.StateId)
+				}
+
+				r := bytes.NewReader(req.Bytes)
+				reqState, err := statefile.Read(r)
+				if err != nil {
+					t.Fatal(err)
+				}
+				if reqState.State.String() != s.String() {
+					t.Fatalf("wanted state %s got %s", s.String(), reqState.State.String())
+				}
+				return providers.WriteStateBytesResponse{
+					// no diags
+				}
+			},
+		}
+
+		// This package will be consumed in a statemgr.Full, so we test using NewRemoteGRPC
+		// and invoke the method on that interface that uses Put.
+		c := NewRemoteGRPC(&provider, typeName, stateId)
+
+		// Set internal state value that will be persisted.
+		c.WriteState(s)
+
+		// Test PersistState, which uses Put.
+		err := c.PersistState(nil)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("state manager made using grpcClient returns expected error from error diagnostic", func(t *testing.T) {
+		expectedErr := "error forced from test"
+		var diags tfdiags.Diagnostics
+		diags = diags.Append(&hcl.Diagnostic{
+			Severity: hcl.DiagError,
+			Summary:  expectedErr,
+			Detail:   expectedErr,
+		})
+		provider := testing_provider.MockProvider{
+			// Mock a provider and internal state store that
+			// have both been configured
+			ConfigureProviderCalled:   true,
+			ConfigureStateStoreCalled: true,
+
+			// Force an error diagnostic
+			WriteStateBytesFn: func(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+				return providers.WriteStateBytesResponse{
+					Diagnostics: diags,
+				}
+			},
+		}
+
+		// This package will be consumed in a statemgr.Full, so we test using NewRemoteGRPC
+		// and invoke the method on that interface that uses Get.
+		c := NewRemoteGRPC(&provider, typeName, stateId)
+
+		// Set internal state value that will be persisted.
+		c.WriteState(s)
+
+		// Test PersistState, which uses Put.
+		err := c.PersistState(nil)
+		if err == nil {
+			t.Fatalf("expected error but got none")
+		}
+		if !strings.Contains(err.Error(), expectedErr) {
+			t.Fatalf("expected error to contain %q, but got: %s", expectedErr, err.Error())
+		}
+	})
+}
+
 // Testing grpcClient's Delete method.
 // This method is needed to implement the remote.Client interface, but
 // this is not invoked by the remote state manager (remote.State) that
-// wil contain the client.
+// will contain the client.
 //
 // In future we should remove the need for a Delete method in
 // remote.Client, but for now it is implemented and tested.