From b7eac0bdbd1c461920c537c02877bccc0b790c91 Mon Sep 17 00:00:00 2001
From: Baraa Basata <bbasata@gmail.com>
Date: Wed, 26 Feb 2025 08:25:21 -0500
Subject: [PATCH] more green

---
 statecheck/expect_known_value_test.go | 42 +++++++++------------------
 statecheck/test_provider_test.go      | 12 +++++++-
 2 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/statecheck/expect_known_value_test.go b/statecheck/expect_known_value_test.go
index eb93ee43..7eced141 100644
--- a/statecheck/expect_known_value_test.go
+++ b/statecheck/expect_known_value_test.go
@@ -225,10 +225,8 @@ func TestExpectKnownValue_CheckState_Bool_KnownValueWrongType(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -253,10 +251,8 @@ func TestExpectKnownValue_CheckState_Bool_KnownValueWrongValue(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -281,10 +277,8 @@ func TestExpectKnownValue_CheckState_Float64(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -309,10 +303,8 @@ func TestExpectKnownValue_CheckState_Float64_KnownValueWrongType(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -337,10 +329,8 @@ func TestExpectKnownValue_CheckState_Float64_KnownValueWrongValue(t *testing.T)
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -365,10 +355,8 @@ func TestExpectKnownValue_CheckState_Int64(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
@@ -392,10 +380,8 @@ func TestExpectKnownValue_CheckState_Int64_KnownValueWrongValue(t *testing.T) {
 	t.Parallel()
 
 	r.Test(t, r.TestCase{
-		ProviderFactories: map[string]func() (*schema.Provider, error){
-			"test": func() (*schema.Provider, error) { //nolint:unparam // required signature
-				return testProvider(), nil
-			},
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"test": providerserver.NewProviderServer(anTestProvider),
 		},
 		Steps: []r.TestStep{
 			{
diff --git a/statecheck/test_provider_test.go b/statecheck/test_provider_test.go
index 294c44e4..3bc4641d 100644
--- a/statecheck/test_provider_test.go
+++ b/statecheck/test_provider_test.go
@@ -20,7 +20,17 @@ var anTestProvider = testprovider.Provider{
 							{
 								Name:     "bool_attribute",
 								Type:     tftypes.Bool,
-								Computed: true,
+								Optional: true,
+							},
+							{
+								Name:     "float_attribute",
+								Type:     tftypes.Number,
+								Optional: true,
+							},
+							{
+								Name:     "int_attribute",
+								Type:     tftypes.Number,
+								Optional: true,
 							},
 						},
 					},