diff --git a/pf/tests/internal/providerbuilder/build_resource.go b/pf/tests/internal/providerbuilder/build_resource.go
index e50d9ca82..df8f83584 100644
--- a/pf/tests/internal/providerbuilder/build_resource.go
+++ b/pf/tests/internal/providerbuilder/build_resource.go
@@ -25,10 +25,11 @@ type Resource struct {
 	Name           string
 	ResourceSchema schema.Schema
 
-	CreateFunc func(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse)
-	ReadFunc   func(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse)
-	UpdateFunc func(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse)
-	DeleteFunc func(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse)
+	CreateFunc      func(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse)
+	ReadFunc        func(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse)
+	UpdateFunc      func(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse)
+	DeleteFunc      func(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse)
+	ImportStateFunc func(ctx context.Context, req resource.ImportStateRequest, resp *resource.ImportStateResponse)
 }
 
 func (r *Resource) Metadata(ctx context.Context, req resource.MetadataRequest, re *resource.MetadataResponse) {
@@ -67,4 +68,11 @@ func (r *Resource) Delete(ctx context.Context, req resource.DeleteRequest, resp
 	r.DeleteFunc(ctx, req, resp)
 }
 
-var _ resource.Resource = &Resource{}
+func (r *Resource) ImportState(ctx context.Context, req resource.ImportStateRequest, resp *resource.ImportStateResponse) {
+	if r.ImportStateFunc == nil {
+		return
+	}
+	r.ImportStateFunc(ctx, req, resp)
+}
+
+var _ resource.ResourceWithImportState = &Resource{}
diff --git a/pf/tests/util.go b/pf/tests/util.go
index de5d3ea8a..24f8c4b00 100644
--- a/pf/tests/util.go
+++ b/pf/tests/util.go
@@ -26,6 +26,8 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
 	rschema "github.com/hashicorp/terraform-plugin-framework/resource/schema"
+	"github.com/hashicorp/terraform-plugin-framework/resource/schema/planmodifier"
+	"github.com/hashicorp/terraform-plugin-framework/resource/schema/stringplanmodifier"
 	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
@@ -63,7 +65,12 @@ func ensureProviderValid(prov *providerbuilder.Provider) {
 	for i := range prov.AllResources {
 		r := &prov.AllResources[i]
 		if r.ResourceSchema.Attributes["id"] == nil {
-			r.ResourceSchema.Attributes["id"] = rschema.StringAttribute{Computed: true}
+			r.ResourceSchema.Attributes["id"] = rschema.StringAttribute{
+				Computed: true,
+				PlanModifiers: []planmodifier.String{
+					stringplanmodifier.UseStateForUnknown(),
+				},
+			}
 		}
 		if r.CreateFunc == nil {
 			r.CreateFunc = func(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse) {
@@ -71,8 +78,12 @@ func ensureProviderValid(prov *providerbuilder.Provider) {
 				resp.State.SetAttribute(ctx, path.Root("id"), "test-id")
 			}
 		}
+		if r.UpdateFunc == nil {
+			r.UpdateFunc = func(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse) {
+				resp.State = tfsdk.State(req.Config)
+			}
+		}
 	}
-
 }
 
 func bridgedProvider(prov *providerbuilder.Provider) info.Provider {