From 4087b40f01c29f82bd2b8b020d5ad47fe7eb5b99 Mon Sep 17 00:00:00 2001
From: vuong-nguyen <44292934+nkvuong@users.noreply.github.com>
Date: Thu, 19 Oct 2023 10:01:07 +0100
Subject: [PATCH] Fix DiffSuppressFunc for `databricks_model_serving` resource
 (#2813)

* fix model serving diff func

* add test

* fix test

* fix logic

* add comment for future
---
 serving/resource_model_serving.go      |   8 +-
 serving/resource_model_serving_test.go | 131 +++++++++++++++++++++++++
 2 files changed, 134 insertions(+), 5 deletions(-)

diff --git a/serving/resource_model_serving.go b/serving/resource_model_serving.go
index c717e2cbf9..321042dce7 100644
--- a/serving/resource_model_serving.go
+++ b/serving/resource_model_serving.go
@@ -23,12 +23,10 @@ func ResourceModelServing() *schema.Resource {
 			common.MustSchemaPath(m, "config", "served_models", "name").Computed = true
 
 			common.MustSchemaPath(m, "config", "served_models", "workload_type").Default = "CPU"
-			// if GPU serving is not enabled, workload_type will always be empty
+			// TODO: `config.served_models.workload_type` should be a `Optional+Computed` field. Also consider this for other similar fields.
+			// In this scenario, if a workspace does not have GPU serving, specifying `workload_type` = 'CPU' will get empty response from API.
 			common.MustSchemaPath(m, "config", "served_models", "workload_type").DiffSuppressFunc = func(k, old, new string, d *schema.ResourceData) bool {
-				if old == "" || new == "" {
-					return true
-				}
-				return false
+				return old == "" && new == "CPU"
 			}
 			common.MustSchemaPath(m, "config", "traffic_config").Computed = true
 
diff --git a/serving/resource_model_serving_test.go b/serving/resource_model_serving_test.go
index cbcfa18fbd..68ac231564 100644
--- a/serving/resource_model_serving_test.go
+++ b/serving/resource_model_serving_test.go
@@ -140,6 +140,137 @@ func TestModelServingCreate(t *testing.T) {
 	}.ApplyNoError(t)
 }
 
+func TestModelServingCreateGPU(t *testing.T) {
+	qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   http.MethodPost,
+				Resource: "/api/2.0/serving-endpoints",
+				ExpectedRequest: serving.CreateServingEndpoint{
+					Name: "test-endpoint",
+					Config: serving.EndpointCoreConfigInput{
+						ServedModels: []serving.ServedModelInput{
+							{
+								Name:               "prod_model",
+								ModelName:          "ads1",
+								ModelVersion:       "2",
+								WorkloadSize:       "Small",
+								WorkloadType:       "GPU_MEDIUM",
+								ScaleToZeroEnabled: true,
+							},
+							{
+								Name:               "candidate_model",
+								ModelName:          "ads1",
+								ModelVersion:       "4",
+								WorkloadSize:       "Small",
+								WorkloadType:       "GPU_MEDIUM",
+								ScaleToZeroEnabled: false,
+							},
+						},
+						TrafficConfig: &serving.TrafficConfig{
+							Routes: []serving.Route{
+								{
+									ServedModelName:   "prod_model",
+									TrafficPercentage: 90,
+								},
+								{
+									ServedModelName:   "candidate_model",
+									TrafficPercentage: 10,
+								},
+							},
+						},
+					},
+				},
+				Response: serving.ServingEndpointDetailed{
+					Name: "test-endpoint",
+				},
+			},
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/serving-endpoints/test-endpoint?",
+				Response: serving.ServingEndpointDetailed{
+					Name: "test-endpoint",
+					State: &serving.EndpointState{
+						ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating,
+					},
+				},
+			},
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/serving-endpoints/test-endpoint?",
+				Response: serving.ServingEndpointDetailed{
+					Id:   "test-endpoint",
+					Name: "test-endpoint",
+					State: &serving.EndpointState{
+						ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating,
+					},
+					Config: &serving.EndpointCoreConfigOutput{
+						ServedModels: []serving.ServedModelOutput{
+							{
+								Name:               "prod_model",
+								ModelName:          "ads1",
+								ModelVersion:       "2",
+								ScaleToZeroEnabled: true,
+							},
+							{
+								Name:               "candidate_model",
+								ModelName:          "ads1",
+								ModelVersion:       "4",
+								ScaleToZeroEnabled: false,
+							},
+						},
+						TrafficConfig: &serving.TrafficConfig{
+							Routes: []serving.Route{
+								{
+									ServedModelName:   "prod_model",
+									TrafficPercentage: 90,
+								},
+								{
+									ServedModelName:   "candidate_model",
+									TrafficPercentage: 10,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		Resource: ResourceModelServing(),
+		HCL: `
+			name = "test-endpoint"
+			config {
+				served_models {
+					name = "prod_model"
+					model_name = "ads1"
+					model_version = "2"
+					workload_size = "Small"
+					workload_type = "GPU_MEDIUM"
+					scale_to_zero_enabled = true
+				}
+				served_models {
+					name = "candidate_model"
+					model_name = "ads1"
+					model_version = "4"
+					workload_size = "Small"
+					workload_type = "GPU_MEDIUM"
+					scale_to_zero_enabled = false
+				}
+				traffic_config {
+					routes {
+						served_model_name = "prod_model"
+						traffic_percentage = 90
+					}
+					routes {
+						served_model_name = "candidate_model"
+						traffic_percentage = 10
+					}
+				}
+			}
+			`,
+		Create: true,
+	}.ApplyNoError(t)
+}
+
 func TestModelServingCreate_Error(t *testing.T) {
 	qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{