From ca95a6a334769750ea4b6f2e1b680a5ad5961340 Mon Sep 17 00:00:00 2001
From: VenelinMartinov <venelin.v.martinov@gmail.com>
Date: Thu, 14 Nov 2024 16:53:47 +0000
Subject: [PATCH] Add datasources to pf providerbuilder (#2615)

This adds data sources to the providerbuilder test module. It allows us
to test data sources in PF
---
 .../providerbuilder/build_datasource.go       | 43 +++++++++++++++++++
 .../providerbuilder/build_provider.go         | 26 ++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 pkg/pf/tests/internal/providerbuilder/build_datasource.go

diff --git a/pkg/pf/tests/internal/providerbuilder/build_datasource.go b/pkg/pf/tests/internal/providerbuilder/build_datasource.go
new file mode 100644
index 000000000..49f70a91d
--- /dev/null
+++ b/pkg/pf/tests/internal/providerbuilder/build_datasource.go
@@ -0,0 +1,43 @@
+// Copyright 2016-2023, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package providerbuilder
+
+import (
+	"context"
+
+	"github.com/hashicorp/terraform-plugin-framework/datasource"
+	"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
+)
+
+type DataSource struct {
+	Name             string
+	DataSourceSchema schema.Schema
+
+	ReadFunc func(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse)
+}
+
+func (r *DataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, re *datasource.MetadataResponse) {
+	re.TypeName = req.ProviderTypeName + "_" + r.Name
+}
+
+func (r *DataSource) Schema(ctx context.Context, _ datasource.SchemaRequest, re *datasource.SchemaResponse) {
+	re.Schema = r.DataSourceSchema
+}
+
+func (r *DataSource) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
+	r.ReadFunc(ctx, req, resp)
+}
+
+var _ datasource.DataSource = &DataSource{}
diff --git a/pkg/pf/tests/internal/providerbuilder/build_provider.go b/pkg/pf/tests/internal/providerbuilder/build_provider.go
index cdd1262fc..d3ffcb42b 100644
--- a/pkg/pf/tests/internal/providerbuilder/build_provider.go
+++ b/pkg/pf/tests/internal/providerbuilder/build_provider.go
@@ -17,6 +17,7 @@ import (
 	"context"
 
 	"github.com/hashicorp/terraform-plugin-framework/datasource"
+	dschema "github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
 	"github.com/hashicorp/terraform-plugin-framework/provider/schema"
@@ -36,6 +37,7 @@ type Provider struct {
 	Version        string
 	ProviderSchema schema.Schema
 	AllResources   []Resource
+	AllDataSources []DataSource
 
 	configureFunc func(ctx context.Context, req provider.ConfigureRequest, resp *provider.ConfigureResponse)
 }
@@ -58,7 +60,14 @@ func (impl *Provider) Configure(ctx context.Context, req provider.ConfigureReque
 }
 
 func (impl *Provider) DataSources(ctx context.Context) []func() datasource.DataSource {
-	return []func() datasource.DataSource{}
+	d := make([]func() datasource.DataSource, len(impl.AllDataSources))
+	for i := 0; i < len(impl.AllDataSources); i++ {
+		i := i
+		d[i] = func() datasource.DataSource {
+			return &impl.AllDataSources[i]
+		}
+	}
+	return d
 }
 
 func (impl *Provider) Resources(ctx context.Context) []func() resource.Resource {
@@ -81,6 +90,7 @@ type NewProviderArgs struct {
 	Version        string
 	ProviderSchema schema.Schema
 	AllResources   []Resource
+	AllDataSources []DataSource
 
 	ConfigureFunc func(ctx context.Context, req provider.ConfigureRequest, resp *provider.ConfigureResponse)
 }
@@ -92,6 +102,7 @@ func NewProvider(params NewProviderArgs) *Provider {
 		Version:        params.Version,
 		ProviderSchema: params.ProviderSchema,
 		AllResources:   params.AllResources,
+		AllDataSources: params.AllDataSources,
 
 		configureFunc: params.ConfigureFunc,
 	}
@@ -130,5 +141,18 @@ func NewProvider(params NewProviderArgs) *Provider {
 		}
 	}
 
+	for i := range prov.AllDataSources {
+		d := &prov.AllDataSources[i]
+		if d.DataSourceSchema.Attributes == nil {
+			d.DataSourceSchema.Attributes = map[string]dschema.Attribute{}
+		}
+
+		if d.ReadFunc == nil {
+			d.ReadFunc = func(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
+				resp.State = tfsdk.State(req.Config)
+			}
+		}
+	}
+
 	return prov
 }