Skip to content

Commit

Permalink
Added Support for Unit Tests for Terraform Plugin Framework
Browse files Browse the repository at this point in the history
  • Loading branch information
@tanmay-db
tanmay-db committed Aug 13, 2024
1 parent 19502ad commit 48c5e96
Show file tree
Hide file tree
Showing 5 changed files with 389 additions and 57 deletions.
2 changes: 2 additions & 0 deletions pluginframework/resource_quality_monitor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ func (r *QualityMonitorResource) Create(ctx context.Context, req resource.Create
return
}
var monitorInfoTfSDK MonitorInfoExtended

// "mismatch between struct and object: Struct defines fields not found in object: baseline_table_name, time_series, monitor_version, status, warehouse_id, dashboard_id, data_classification_config, drift_metrics_table_name, snapshot, skip_builtin_dashboard, notifications, schedule, table_name, custom_metrics, latest_monitor_failure_msg, inference_log, profile_metrics_table_name, and slicing_exprs."
diags := req.Plan.Get(ctx, &monitorInfoTfSDK)
resp.Diagnostics.Append(diags...)
if resp.Diagnostics.HasError() {
Expand Down
304 changes: 260 additions & 44 deletions pluginframework/resource_quality_monitor_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,68 +6,284 @@ import (

"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/catalog"
"github.com/databricks/terraform-provider-databricks/common"
pluginFrameworkResource "github.com/hashicorp/terraform-plugin-framework/resource"
"github.com/hashicorp/terraform-plugin-framework/tfsdk"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)

func TestQualityMonitorCreateTimeseriesPluginFramework(t *testing.T) {
// Setup mock workspace or account client whichever is applicable
// This will be use to mock the databricks-sdk-go calls
MockWorkspaceClientFunc := func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.Create(mock.Anything, catalog.CreateMonitor{
TableName: "test_table",
OutputSchemaName: "output.schema",
AssetsDir: "sample.dir",
TimeSeries: &catalog.MonitorTimeSeries{
Granularities: []string{"1 day"},
TimestampCol: "timestamp",
},
}).Return(&catalog.MonitorInfo{
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusPending,
DriftMetricsTableName: "test_table_drift",
}, nil)
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
DriftMetricsTableName: "test_table_drift",
}, nil)
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.Create(mock.Anything, catalog.CreateMonitor{
TableName: "test_table",
OutputSchemaName: "output.schema",
AssetsDir: "sample.dir",
TimeSeries: &catalog.MonitorTimeSeries{
Granularities: []string{"1 day"},
TimestampCol: "timestamp",
},
}).Return(&catalog.MonitorInfo{
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusPending,
DriftMetricsTableName: "test_table_drift",
}, nil)
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
DriftMetricsTableName: "test_table_drift",
}, nil)
},
}.Start(t)
assert.NoError(t, err)

qualityMonitorResource := QualityMonitorResource{
Client: client,
}

schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)
createRequest := pluginFrameworkResource.CreateRequest{
Plan: tfsdk.Plan{
Raw: MapToTfTypesValue(map[string]any{
"table_name": "test_table",
"assets_dir": "sample.dir",
"output_schema_name": "output.schema",
"time_series": `{"granularities":["1 day"],"timestamp_col":"timestamp"}`,
}),
Schema: schema.Schema,
},
}
_ = MockWorkspaceClientFunc
createResponse := pluginFrameworkResource.CreateResponse{}

qualityMonitorResource.Create(ctx, createRequest, &createResponse)

assert.False(t, createResponse.Diagnostics.HasError())
}

// Start the server
// tanmaytodo
func TestQualityMonitorCreateInferencePluginFramework(t *testing.T) {
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.Create(mock.Anything, catalog.CreateMonitor{
TableName: "test_table",
OutputSchemaName: "output.schema",
AssetsDir: "sample.dir",
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 day"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
ProblemType: catalog.MonitorInferenceLogProblemTypeProblemTypeRegression,
},
}).Return(&catalog.MonitorInfo{
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
}, nil)
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 day"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
ProblemType: catalog.MonitorInferenceLogProblemTypeProblemTypeRegression,
},
}, nil)
},
}.Start(t)
assert.NoError(t, err)

// Create resource using the above mock workspace client
qualityMonitorResource := QualityMonitorResource{
Client: &common.DatabricksClient{},
Client: client,
}

// Create request and response for the respective CRUD operation and pass the config
schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)
createRequest := pluginFrameworkResource.CreateRequest{
Config: GetPluginFrameworkConfig(`
table_name = "test_table",
assets_dir = "sample.dir",
output_schema_name = "output.schema",
time_series = {
granularities = ["1 day"],
timestamp_col = "timestamp"
}
`),
Plan: tfsdk.Plan{
Raw: MapToTfTypesValue(map[string]any{
"table_name": "test_table",
"assets_dir": "sample.dir",
"output_schema_name": "output.schema",
"inference_log": `{
granularities: ["1 day"],
timestamp_col: "timestamp",
prediction_col: "prediction",
model_id_col: "model_id",
problem_type: "PROBLEM_TYPE_REGRESSION"
}`,
}),
Schema: schema.Schema,
},
}
createResponse := pluginFrameworkResource.CreateResponse{}

qualityMonitorResource.Create(ctx, createRequest, &createResponse)

assert.False(t, createResponse.Diagnostics.HasError())
}

func TestQualityMonitorCreateSnapshot(t *testing.T) {
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.Create(mock.Anything, catalog.CreateMonitor{
TableName: "test_table",
OutputSchemaName: "output.schema",
AssetsDir: "sample.dir",
Snapshot: &catalog.MonitorSnapshot{},
}).Return(&catalog.MonitorInfo{
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
}, nil)
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "sample.dir",
OutputSchemaName: "output.schema",
Snapshot: &catalog.MonitorSnapshot{},
}, nil)
},
}.Start(t)
assert.NoError(t, err)

qualityMonitorResource := QualityMonitorResource{
Client: client,
}

schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)

createRequest := pluginFrameworkResource.CreateRequest{
Plan: tfsdk.Plan{
Raw: MapToTfTypesValue(map[string]any{
"table_name": "test_table",
"assets_dir": "sample.dir",
"output_schema_name": "output.schema",
"snapshot": `{}`,
}),
Schema: schema.Schema,
},
}
createResponse := pluginFrameworkResource.CreateResponse{}

// Call the method to unit test
qualityMonitorResource.Create(ctx, createRequest, &createResponse)

// Add assertions
assert.False(t, createResponse.Diagnostics.HasError())
}

func TestQualityMonitorGet(t *testing.T) {
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "new_assets.dir",
OutputSchemaName: "output.schema",
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 week"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
},
DriftMetricsTableName: "test_table_drift"}, nil)
},
}.Start(t)
assert.NoError(t, err)

qualityMonitorResource := QualityMonitorResource{
Client: client,
}

schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)
}

func TestQualityMonitorUpdate(t *testing.T) {
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.Update(mock.Anything, catalog.UpdateMonitor{
TableName: "test_table",
OutputSchemaName: "output.schema",
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 week"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
ProblemType: catalog.MonitorInferenceLogProblemTypeProblemTypeRegression,
},
}).Return(&catalog.MonitorInfo{
AssetsDir: "new_assets.dir",
OutputSchemaName: "output.schema",
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 week"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
},
DriftMetricsTableName: "test_table_drift",
}, nil)
e.GetByTableName(mock.Anything, "test_table").Return(&catalog.MonitorInfo{
TableName: "test_table",
Status: catalog.MonitorInfoStatusMonitorStatusActive,
AssetsDir: "new_assets.dir",
OutputSchemaName: "output.schema",
InferenceLog: &catalog.MonitorInferenceLog{
Granularities: []string{"1 week"},
TimestampCol: "timestamp",
PredictionCol: "prediction",
ModelIdCol: "model_id",
},
DriftMetricsTableName: "test_table_drift",
}, nil)
},
}.Start(t)
assert.NoError(t, err)

qualityMonitorResource := QualityMonitorResource{
Client: client,
}

schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)
}

func TestQualityMonitorDelete(t *testing.T) {
ctx := context.Background()
client, err := ResourceFixturePluginFramework{
MockWorkspaceClientFunc: func(w *mocks.MockWorkspaceClient) {
e := w.GetMockQualityMonitorsAPI().EXPECT()
e.DeleteByTableName(mock.Anything, "test_table").Return(nil)
},
}.Start(t)
assert.NoError(t, err)

qualityMonitorResource := QualityMonitorResource{
Client: client,
}

schema := pluginFrameworkResource.SchemaResponse{}
qualityMonitorResource.Schema(ctx, pluginFrameworkResource.SchemaRequest{}, &schema)
}
Loading

0 comments on commit 48c5e96

Please sign in to comment.