From 0c78935a30cd8de990a6973ff6443fb7b6711a7b Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Tue, 10 Dec 2024 17:30:16 +0000
Subject: [PATCH] fix panic

---
 pkg/tfshim/sdk-v2/provider2.go      |  8 ++++++-
 pkg/tfshim/sdk-v2/provider2_test.go | 37 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/pkg/tfshim/sdk-v2/provider2.go b/pkg/tfshim/sdk-v2/provider2.go
index 347496e15..ce38816ec 100644
--- a/pkg/tfshim/sdk-v2/provider2.go
+++ b/pkg/tfshim/sdk-v2/provider2.go
@@ -89,6 +89,9 @@ func (s *v2InstanceState2) Type() string {
 }
 
 func (s *v2InstanceState2) ID() string {
+	if s.stateValue.IsNull() {
+		return ""
+	}
 	id := s.stateValue.GetAttr("id")
 	if !id.IsKnown() {
 		return ""
@@ -621,7 +624,7 @@ func (s *grpcServer) ApplyResourceChange(
 	config, priorState, plannedState cty.Value,
 	plannedMeta map[string]interface{},
 	providerMeta *cty.Value,
-) (*v2InstanceState2, error) {
+) (shim.InstanceState, error) {
 	configVal, err := msgpack.Marshal(config, ty)
 	if err != nil {
 		return nil, err
@@ -671,6 +674,9 @@ func (s *grpcServer) ApplyResourceChange(
 		}
 	}
 	returnErr := handleDiagnostics(ctx, resp.Diagnostics, applyErr)
+	if newState.IsNull() {
+		return nil, returnErr
+	}
 	return &v2InstanceState2{
 		resourceType: typeName,
 		stateValue:   newState,
diff --git a/pkg/tfshim/sdk-v2/provider2_test.go b/pkg/tfshim/sdk-v2/provider2_test.go
index 47ad7ac77..a96c7553b 100644
--- a/pkg/tfshim/sdk-v2/provider2_test.go
+++ b/pkg/tfshim/sdk-v2/provider2_test.go
@@ -641,3 +641,40 @@ func (l *testLogger) Error(msg string) {
 func (*testLogger) StatusUntyped() any {
 	return "?"
 }
+
+func TestInstanceStateId(t *testing.T) {
+	t.Parallel()
+
+	state := v2InstanceState2{
+		stateValue: cty.ObjectVal(map[string]cty.Value{
+			"id": cty.StringVal("1"),
+		}),
+	}
+	assert.Equal(t, "1", state.ID())
+
+	state.stateValue = cty.ObjectVal(map[string]cty.Value{
+		"id": cty.UnknownVal(cty.String),
+	})
+	assert.Equal(t, "", state.ID())
+
+	state.stateValue = cty.Value{}
+	assert.Equal(t, "", state.ID())
+}
+
+func TestInstanceStateObject(t *testing.T) {
+	t.Parallel()
+
+	state := v2InstanceState2{
+		stateValue: cty.ObjectVal(map[string]cty.Value{
+			"id": cty.StringVal("1"),
+		}),
+	}
+	actual, err := state.Object(nil)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]interface{}{"id": "1"}, actual)
+
+	state.stateValue = cty.Value{}
+	actual, err = state.Object(nil)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]interface{}(nil), actual)
+}