Add test for timeouts on ML job state

Author: tobio

internal/elasticsearch/ml/datafeed_state/create.go

@@ -2,8 +2,10 @@ package datafeed_state
 
 import (
 	"context"
+	"fmt"
 	"time"
 
+	"github.com/elastic/terraform-provider-elasticstack/internal/diagutil"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
 )
 
@@ -23,5 +25,9 @@ func (r *mlDatafeedStateResource) Create(ctx context.Context, req resource.Creat
 	}
 
 	diags = r.update(ctx, req.Plan, &resp.State, createTimeout)
+	if diags.Contains(diagutil.FrameworkDiagFromError(context.DeadlineExceeded)[0]) {
+		diags.AddError("Operation timed out", fmt.Sprintf("The operation to create the ML datafeed state timed out after %s. You may need to allocate more free memory within ML nodes by either closing other jobs, or increasing the overall ML memory. You may retry the operation.", createTimeout))
+	}
+
 	resp.Diagnostics.Append(diags...)
 }

internal/elasticsearch/ml/datafeed_state/update.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/elastic/terraform-provider-elasticstack/internal/clients"
 	"github.com/elastic/terraform-provider-elasticstack/internal/clients/elasticsearch"
+	"github.com/elastic/terraform-provider-elasticstack/internal/diagutil"
 	"github.com/elastic/terraform-provider-elasticstack/internal/elasticsearch/ml/datafeed"
 	"github.com/hashicorp/terraform-plugin-framework-timetypes/timetypes"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
@@ -32,6 +33,10 @@ func (r *mlDatafeedStateResource) Update(ctx context.Context, req resource.Updat
 	}
 
 	diags = r.update(ctx, req.Plan, &resp.State, updateTimeout)
+	if diags.Contains(diagutil.FrameworkDiagFromError(context.DeadlineExceeded)[0]) {
+		diags.AddError("Operation timed out", fmt.Sprintf("The operation to update the ML datafeed state timed out after %s. You may need to allocate more free memory within ML nodes by either closing other jobs, or increasing the overall ML memory. You may retry the operation.", updateTimeout))
+	}
+
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
 		return

internal/elasticsearch/ml/job_state/acc_test.go

@@ -115,3 +115,23 @@ func TestAccResourceMLJobStateImport(t *testing.T) {
 		},
 	})
 }
+
+func TestAccResourceMLJobState_timeouts(t *testing.T) {
+	jobID := fmt.Sprintf("test-job-%s", sdkacctest.RandStringFromCharSet(10, sdkacctest.CharSetAlphaNum))
+	indexName := fmt.Sprintf("test-datafeed-index-%s", sdkacctest.RandStringFromCharSet(10, sdkacctest.CharSetAlphaNum))
+
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() { acctest.PreCheck(t) },
+		Steps: []resource.TestStep{
+			{
+				ProtoV6ProviderFactories: acctest.Providers,
+				ConfigDirectory:          acctest.NamedTestCaseDirectory("timeouts"),
+				ConfigVariables: config.Variables{
+					"job_id":     config.StringVariable(jobID),
+					"index_name": config.StringVariable(indexName),
+				},
+				ExpectError: regexp.MustCompile("Operation timed out"),
+			},
+		},
+	})
+}

internal/elasticsearch/ml/job_state/create.go

@@ -2,8 +2,10 @@ package job_state
 
 import (
 	"context"
+	"fmt"
 	"time"
 
+	"github.com/elastic/terraform-provider-elasticstack/internal/diagutil"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
 )
 
@@ -23,5 +25,9 @@ func (r *mlJobStateResource) Create(ctx context.Context, req resource.CreateRequ
 	}
 
 	diags = r.update(ctx, req.Plan, &resp.State, createTimeout)
+	if diags.Contains(diagutil.FrameworkDiagFromError(context.DeadlineExceeded)[0]) {
+		diags.AddError("Operation timed out", fmt.Sprintf("The operation to create the ML job state timed out after %s. You may need to allocate more free memory within ML nodes by either closing other jobs, or increasing the overall ML memory. You may retry the operation.", createTimeout))
+	}
+
 	resp.Diagnostics.Append(diags...)
 }

internal/elasticsearch/ml/job_state/testdata/TestAccResourceMLJobState_timeouts/timeouts/job_state.tf

@@ -0,0 +1,58 @@
+variable "job_id" {
+  description = "The ML job ID"
+  type        = string
+}
+
+variable "index_name" {
+  description = "The index name"
+  type        = string
+}
+
+provider "elasticstack" {
+  elasticsearch {}
+}
+
+resource "elasticstack_elasticsearch_index" "test" {
+  name                = var.index_name
+  deletion_protection = false
+  mappings = jsonencode({
+    properties = {
+      "@timestamp" = {
+        type = "date"
+      }
+      value = {
+        type = "double"
+      }
+    }
+  })
+}
+
+resource "elasticstack_elasticsearch_ml_anomaly_detection_job" "test" {
+  job_id      = var.job_id
+  description = "Test job for datafeed state timeout testing with large memory"
+  analysis_config = {
+    bucket_span = "1h"
+    detectors = [{
+      function             = "count"
+      detector_description = "count"
+    }]
+  }
+  data_description = {
+    time_field  = "@timestamp"
+    time_format = "epoch_ms"
+  }
+  analysis_limits = {
+    model_memory_limit = "784mb"  # Large memory requirement close to cluster limit
+  }
+  allow_lazy_open = true  # This should cause datafeed to wait for available node
+}
+
+resource "elasticstack_elasticsearch_ml_job_state" "test" {
+  job_id = elasticstack_elasticsearch_ml_anomaly_detection_job.test.job_id
+  state  = "opened"
+  
+  timeouts = {
+    create = "10s"
+    update = "10s"
+  }
+}

internal/elasticsearch/ml/job_state/update.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/elastic/terraform-provider-elasticstack/internal/clients"
 	"github.com/elastic/terraform-provider-elasticstack/internal/clients/elasticsearch"
+	"github.com/elastic/terraform-provider-elasticstack/internal/diagutil"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
 	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
@@ -30,6 +31,10 @@ func (r *mlJobStateResource) Update(ctx context.Context, req resource.UpdateRequ
 	}
 
 	diags = r.update(ctx, req.Plan, &resp.State, updateTimeout)
+	if diags.Contains(diagutil.FrameworkDiagFromError(context.DeadlineExceeded)[0]) {
+		diags.AddError("Operation timed out", fmt.Sprintf("The operation to update the ML job state timed out after %s. You may need to allocate more free memory within ML nodes by either closing other jobs, or increasing the overall ML memory. You may retry the operation.", updateTimeout))
+	}
+
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
 		return