From d252161328a0a1497f9c666b71bb9ed79386417b Mon Sep 17 00:00:00 2001
From: Luke Mondy <luke.mondy@audinate.com>
Date: Tue, 11 Jun 2024 11:45:29 +1000
Subject: [PATCH] Return custom_fields from data.netbox_prefix

Add a test to ensure the data is returned correctly too.
---
 netbox/data_source_netbox_prefix.go      |  8 +++++
 netbox/data_source_netbox_prefix_test.go | 40 ++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/netbox/data_source_netbox_prefix.go b/netbox/data_source_netbox_prefix.go
index 736d9a9a..be311e5c 100644
--- a/netbox/data_source_netbox_prefix.go
+++ b/netbox/data_source_netbox_prefix.go
@@ -27,6 +27,7 @@ func dataSourceNetboxPrefix() *schema.Resource {
 				ValidateFunc:  validation.IsCIDR,
 				AtLeastOneOf:  []string{"description", "family", "prefix", "vlan_vid", "vrf_id", "vlan_id", "site_id", "role_id", "cidr", "tag"},
 			},
+			customFieldsKey: customFieldsSchema,
 			"description": {
 				Type:         schema.TypeString,
 				Optional:     true,
@@ -175,6 +176,13 @@ func dataSourceNetboxPrefixRead(d *schema.ResourceData, m interface{}) error {
 	d.Set("family", int(*result.Family.Value))
 	d.Set("tags", getTagListFromNestedTagList(result.Tags))
 
+	cf := getCustomFields(result.CustomFields)
+	if cf != nil {
+		d.Set(customFieldsKey, cf)
+	}
+	if result.Role != nil {
+		d.Set("role_id", result.Role.ID)
+	}
 	if result.Vrf != nil {
 		d.Set("vrf_id", result.Vrf.ID)
 	}
diff --git a/netbox/data_source_netbox_prefix_test.go b/netbox/data_source_netbox_prefix_test.go
index 879c39ce..520dcdd3 100644
--- a/netbox/data_source_netbox_prefix_test.go
+++ b/netbox/data_source_netbox_prefix_test.go
@@ -2,6 +2,7 @@ package netbox
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
@@ -118,3 +119,42 @@ data "netbox_prefix" "by_family" {
 		},
 	})
 }
+
+func TestAccNetboxPrefixDataSource_customFields(t *testing.T) {
+	testSlug := "prefix_customfields"
+	testPrefix := "10.0.0.0/24"
+	testField := strings.ReplaceAll(testAccGetTestName(testSlug), "-", "_")
+	resource.Test(t, resource.TestCase{
+		PreCheck:  func() { testAccPreCheck(t) },
+		Providers: testAccProviders,
+		Steps: []resource.TestStep{
+			{
+				Config: fmt.Sprintf(`
+resource "netbox_custom_field" "test" {
+  name = "%[1]s"
+  type = "text"
+  content_types = ["ipam.prefix"]
+  weight        = 100
+}
+
+resource "netbox_prefix" "test" {
+  prefix = "%[2]s"
+  status = "active"
+  custom_fields = {
+    "${netbox_custom_field.test.name}" = "test value"
+  }
+}
+
+data "netbox_prefix" "test_output" {
+  depends_on = [netbox_prefix.test]
+  prefix = "%[2]s"
+}`, testField, testPrefix),
+				Check: resource.ComposeTestCheckFunc(
+					resource.TestCheckResourceAttr("data.netbox_prefix.test_output", "status", "active"),
+					resource.TestCheckResourceAttr("data.netbox_prefix.test_output", "prefix", testPrefix),
+					resource.TestCheckResourceAttr("data.netbox_prefix.test_output", "custom_fields."+testField, "test value"),
+				),
+			},
+		},
+	})
+}