feat: add suite_status filter for cloud drift tests (#210)

Author: sohil-kshirsagar

cmd/list.go

@@ -15,6 +15,7 @@ import (
 	"github.com/Use-Tusk/tusk-cli/internal/runner"
 	"github.com/Use-Tusk/tusk-cli/internal/tui"
 	"github.com/Use-Tusk/tusk-cli/internal/utils"
+	backend "github.com/Use-Tusk/tusk-drift-schemas/generated/go/backend"
 )
 
 //go:embed short_docs/drift/drift_list.md
@@ -93,14 +94,26 @@ func listTests(cmd *cobra.Command, args []string) error {
 			return formatApiError(err)
 		}
 
-		all, err := api.FetchAllTraceTestsWithCache(
-			context.Background(),
-			client,
-			authOptions,
-			cfg.Service.ID,
-			false,
-			false,
-		)
+		var all []*backend.TraceTest
+		usedStatusFilter := false
+		if val, ok := runner.ExtractSuiteStatusFromFilter(filter); ok {
+			if statusFilter := runner.ParseTraceTestStatusFilter(val); statusFilter != nil {
+				all, err = api.FetchAllTraceTests(context.Background(), client, authOptions, cfg.Service.ID, &api.FetchAllTraceTestsOptions{
+					StatusFilter: statusFilter,
+				})
+				usedStatusFilter = true
+			}
+		}
+		if !usedStatusFilter && err == nil {
+			all, err = api.FetchAllTraceTestsWithCache(
+				context.Background(),
+				client,
+				authOptions,
+				cfg.Service.ID,
+				false,
+				false,
+			)
+		}
 		if err != nil {
 			return formatApiError(err)
 		}

cmd/run.go

@@ -598,7 +598,11 @@ func runTests(cmd *cobra.Command, args []string) error {
 			if isValidation {
 				preloadedTests, err = fetchValidationTraceTests(context.Background(), client, authOptions, cfg.Service.ID)
 			} else {
-				preloadedTests, err = loadCloudTests(context.Background(), client, authOptions, cfg.Service.ID, driftRunID, traceTestID, allCloudTraceTests || !ci, quiet)
+				var suiteStatusFilter *backend.TraceTestStatus
+				if val, ok := runner.ExtractSuiteStatusFromFilter(filter); ok {
+					suiteStatusFilter = runner.ParseTraceTestStatusFilter(val)
+				}
+				preloadedTests, err = loadCloudTests(context.Background(), client, authOptions, cfg.Service.ID, driftRunID, traceTestID, allCloudTraceTests || !ci, quiet, suiteStatusFilter)
 			}
 			if err != nil {
 				return formatApiError(fmt.Errorf("failed to load cloud tests: %w", err))
@@ -896,7 +900,7 @@ func runTests(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-func loadCloudTests(ctx context.Context, client *api.TuskClient, auth api.AuthOptions, serviceID, driftRunID, traceTestID string, allCloud bool, quiet bool) ([]runner.Test, error) {
+func loadCloudTests(ctx context.Context, client *api.TuskClient, auth api.AuthOptions, serviceID, driftRunID, traceTestID string, allCloud bool, quiet bool, suiteStatusFilter *backend.TraceTestStatus) ([]runner.Test, error) {
 	if traceTestID != "" {
 		req := &backend.GetTraceTestRequest{
 			ObservableServiceId: serviceID,
@@ -912,7 +916,14 @@ func loadCloudTests(ctx context.Context, client *api.TuskClient, auth api.AuthOp
 	var all []*backend.TraceTest
 	var err error
 
-	if allCloud {
+	switch {
+	case suiteStatusFilter != nil:
+		// When filtering by suite status, bypass cache and use GetAllTraceTests
+		// with the status filter directly
+		all, err = api.FetchAllTraceTests(ctx, client, auth, serviceID, &api.FetchAllTraceTestsOptions{
+			StatusFilter: suiteStatusFilter,
+		})
+	case allCloud:
 		all, err = api.FetchAllTraceTestsWithCache(
 			ctx,
 			client,
@@ -921,7 +932,7 @@ func loadCloudTests(ctx context.Context, client *api.TuskClient, auth api.AuthOp
 			false,
 			quiet,
 		)
-	} else {
+	default:
 		all, err = api.FetchDriftRunTraceTests(
 			ctx,
 			client,
@@ -958,7 +969,11 @@ func makeLoadTestsFunc(
 			if traceID != "" && traceTestID == "" {
 				return nil, fmt.Errorf("specify --trace-test-id to run against a single trace test in Tusk Drift Cloud")
 			}
-			tests, err = loadCloudTests(ctx, client, auth, serviceID, driftRunID, traceTestID, allCloud, quiet)
+			var suiteStatusFilter *backend.TraceTestStatus
+			if val, ok := runner.ExtractSuiteStatusFromFilter(filter); ok {
+				suiteStatusFilter = runner.ParseTraceTestStatusFilter(val)
+			}
+			tests, err = loadCloudTests(ctx, client, auth, serviceID, driftRunID, traceTestID, allCloud, quiet, suiteStatusFilter)
 			if err != nil {
 				return nil, err
 			}

cmd/short_docs/drift/drift_filter.md

@@ -2,15 +2,19 @@
 
 Filter tests with `-f`/`--filter`.
 
-Fields: `path=...,name=...,type=...,method=...,status=...,id=...`.
+Fields: `path=...,name=...,type=...,method=...,status=...,id=...,suite_status=...`.
 Comma-separated, values are regex.
 
+Use `suite_status` to filter cloud tests by suite status (`draft` or `in_suite`).
+When `suite_status=draft` is set, draft tests are fetched directly from the backend.
+
 Examples:
 
 ```bash
 tusk drift <list/run> -f 'type=GRAPHQL,op=^GetUser$'
 tusk drift <list/run> -f 'method=POST,path=/checkout'
 tusk drift <list/run> -f 'file=2025-09-24.*trace.*\\.jsonl'
+tusk drift run --cloud -f 'suite_status=draft'
 ```
 
 See <https://github.com/Use-Tusk/tusk-cli/blob/main/docs/drift/filter.md> for more details.

docs/drift/filter.md

@@ -16,6 +16,7 @@ Keys (case-insensitive; aliases in parentheses):
 - `status` (`s`) – test status label for display (e.g., `success`, `error`)
 - `id` (`trace`, `trace_id`) – trace ID
 - `file` (`filename`, `f`) – source file name
+- `suite_status` (`suite`) – cloud suite status: `draft` or `in_suite` (exact values only, not regex)
 
 Notes:
 
@@ -38,6 +39,11 @@ HTTP:
 - By method + route: `tusk drift run -f 'method=POST,path=/checkout'`
 - By type: `tusk drift list -f 'type=HTTP'`
 
+Suite status (cloud only):
+
+- Draft tests only: `tusk drift run --cloud -f 'suite_status=draft'`
+- In-suite tests only: `tusk drift run --cloud -f 'suite_status=in_suite'`
+
 Trace/file:
 
 - Specific trace: `tusk drift run -f 'id=84d0de6b4e4498e996c7f8b8c0f35230'`

go.mod

@@ -4,7 +4,7 @@ go 1.25.0
 
 require (
 	github.com/Use-Tusk/fence v0.1.36
-	github.com/Use-Tusk/tusk-drift-schemas v0.1.30
+	github.com/Use-Tusk/tusk-drift-schemas v0.1.32
 	github.com/agnivade/levenshtein v1.0.3
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1
 	github.com/charmbracelet/bubbles v0.21.1-0.20250623103423-23b8fd6302d7

go.sum

@@ -4,6 +4,8 @@ github.com/Use-Tusk/fence v0.1.36 h1:8S15y8cp3X+xXukx6AN0Ky/aX9/dZyW3fLw5XOQ8YtE
 github.com/Use-Tusk/fence v0.1.36/go.mod h1:YkowBDzXioVKJE16vg9z3gSVC6vhzkIZZw2dFf7MW/o=
 github.com/Use-Tusk/tusk-drift-schemas v0.1.30 h1:A45pJ/Za6BLIfTLF53BhuzKHHSJ9L7dXEisnuKT5dTc=
 github.com/Use-Tusk/tusk-drift-schemas v0.1.30/go.mod h1:pa3EvTj9kKxl9f904RVFkj9YK1zB75QogboKi70zalM=
+github.com/Use-Tusk/tusk-drift-schemas v0.1.32 h1:9+q1RH0036rG3RDjVEeUf0ejMsVP7AxqJ8uQ+XPPCH8=
+github.com/Use-Tusk/tusk-drift-schemas v0.1.32/go.mod h1:pa3EvTj9kKxl9f904RVFkj9YK1zB75QogboKi70zalM=
 github.com/agnivade/levenshtein v1.0.3 h1:M5ZnqLOoZR8ygVq0FfkXsNOKzMCk0xRiow0R5+5VkQ0=
 github.com/agnivade/levenshtein v1.0.3/go.mod h1:4SFRZbbXWLF4MU1T9Qg0pGgH3Pjs+t6ie5efyrwRJXs=
 github.com/alecthomas/assert/v2 v2.7.0 h1:QtqSACNS3tF7oasA8CU6A6sXZSBDqnm7RfpLl9bZqbE=

internal/api/fetch_tests.go

@@ -31,6 +31,9 @@ type FetchAllTraceTestsOptions struct {
 	Message string
 	// PageSize for pagination (default 25)
 	PageSize int32
+	// StatusFilter filters by trace test suite status (e.g., DRAFT, IN_SUITE).
+	// If nil, the server defaults to IN_SUITE.
+	StatusFilter *backend.TraceTestStatus
 }
 
 // FetchAllTraceTests fetches all trace tests from the cloud with a progress bar.
@@ -64,6 +67,7 @@ func FetchAllTraceTests(
 		req := &backend.GetAllTraceTestsRequest{
 			ObservableServiceId: serviceID,
 			PageSize:            opts.PageSize,
+			StatusFilter:        opts.StatusFilter,
 		}
 		if cursor != "" {
 			req.PaginationCursor = &cursor

internal/runner/convert.go

@@ -31,6 +31,7 @@ func ConvertTraceTestToRunnerTest(tt *backend.TraceTest) Test {
 		DisplayType: "HTTP", // will be overridden if we detect GraphQL/etc.
 		DisplayName: fmt.Sprintf("Trace %s", tt.TraceId),
 		Status:      "pending",
+		SuiteStatus: protoTraceTestStatusToString(tt.GetStatus()),
 	}
 
 	// Extract environment from any span that has it
@@ -287,6 +288,34 @@ func ConvertRunnerResultToTraceTestResult(result TestResult, test Test) *backend
 	return out
 }
 
+func protoTraceTestStatusToString(s backend.TraceTestStatus) string {
+	switch s {
+	case backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT:
+		return "draft"
+	case backend.TraceTestStatus_TRACE_TEST_STATUS_IN_SUITE:
+		return "in_suite"
+	case backend.TraceTestStatus_TRACE_TEST_STATUS_REMOVED:
+		return "removed"
+	default:
+		return ""
+	}
+}
+
+// ParseTraceTestStatusFilter converts a user-provided filter value to a proto TraceTestStatus.
+// Returns nil if the value doesn't match a known status.
+func ParseTraceTestStatusFilter(val string) *backend.TraceTestStatus {
+	switch strings.ToLower(val) {
+	case "draft":
+		s := backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT
+		return &s
+	case "in_suite":
+		s := backend.TraceTestStatus_TRACE_TEST_STATUS_IN_SUITE
+		return &s
+	default:
+		return nil
+	}
+}
+
 func getStringFromStruct(s *structpb.Struct, key string) (string, bool) {
 	if s == nil || s.Fields == nil {
 		return "", false

internal/runner/convert_test.go

@@ -7,6 +7,7 @@ import (
 
 	backend "github.com/Use-Tusk/tusk-drift-schemas/generated/go/backend"
 	core "github.com/Use-Tusk/tusk-drift-schemas/generated/go/core"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -174,6 +175,57 @@ func TestConvertTraceTestsToRunnerTests(t *testing.T) {
 	require.Equal(t, ConvertTraceTestToRunnerTest(tt2), got[1])
 }
 
+func TestProtoTraceTestStatusToString(t *testing.T) {
+	t.Parallel()
+
+	assert.Equal(t, "draft", protoTraceTestStatusToString(backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT))
+	assert.Equal(t, "in_suite", protoTraceTestStatusToString(backend.TraceTestStatus_TRACE_TEST_STATUS_IN_SUITE))
+	assert.Equal(t, "removed", protoTraceTestStatusToString(backend.TraceTestStatus_TRACE_TEST_STATUS_REMOVED))
+	assert.Equal(t, "", protoTraceTestStatusToString(backend.TraceTestStatus_TRACE_TEST_STATUS_UNSPECIFIED))
+}
+
+func TestParseTraceTestStatusFilter(t *testing.T) {
+	t.Parallel()
+
+	draft := ParseTraceTestStatusFilter("draft")
+	require.NotNil(t, draft)
+	require.Equal(t, backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT, *draft)
+
+	inSuite := ParseTraceTestStatusFilter("in_suite")
+	require.NotNil(t, inSuite)
+	require.Equal(t, backend.TraceTestStatus_TRACE_TEST_STATUS_IN_SUITE, *inSuite)
+
+	// Case insensitive
+	draftUpper := ParseTraceTestStatusFilter("DRAFT")
+	require.NotNil(t, draftUpper)
+	require.Equal(t, backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT, *draftUpper)
+
+	// Unknown returns nil
+	require.Nil(t, ParseTraceTestStatusFilter("removed"))
+	require.Nil(t, ParseTraceTestStatusFilter("unknown"))
+	require.Nil(t, ParseTraceTestStatusFilter(""))
+}
+
+func TestConvertTraceTestToRunnerTest_SuiteStatus(t *testing.T) {
+	t.Parallel()
+
+	tt := &backend.TraceTest{
+		Id:      "tt-1",
+		TraceId: "trace-1",
+		Status:  backend.TraceTestStatus_TRACE_TEST_STATUS_DRAFT,
+	}
+	got := ConvertTraceTestToRunnerTest(tt)
+	require.Equal(t, "draft", got.SuiteStatus)
+
+	tt.Status = backend.TraceTestStatus_TRACE_TEST_STATUS_IN_SUITE
+	got = ConvertTraceTestToRunnerTest(tt)
+	require.Equal(t, "in_suite", got.SuiteStatus)
+
+	tt.Status = backend.TraceTestStatus_TRACE_TEST_STATUS_UNSPECIFIED
+	got = ConvertTraceTestToRunnerTest(tt)
+	require.Equal(t, "", got.SuiteStatus)
+}
+
 func TestConvertRunnerResultToTraceTestResult(t *testing.T) {
 	t.Parallel()
 

internal/runner/filter.go

@@ -61,14 +61,18 @@ func parseFieldedFilter(q string) ([]fieldMatcher, error) {
 		if len(val) >= 2 && ((val[0] == '\'' && val[len(val)-1] == '\'') || (val[0] == '"' && val[len(val)-1] == '"')) {
 			val = val[1 : len(val)-1]
 		}
-		re, err := regexp.Compile(val)
-		if err != nil {
-			return nil, fmt.Errorf("invalid regex for %s: %w", key, err)
-		}
 		field := normalizeFilterFieldKey(key)
 		if field == "" {
 			return nil, fmt.Errorf("unknown filter field: %s", key)
 		}
+		// suite_status values are always lowercase; make matching case-insensitive
+		if field == "suite_status" {
+			val = "(?i)" + val
+		}
+		re, err := regexp.Compile(val)
+		if err != nil {
+			return nil, fmt.Errorf("invalid regex for %s: %w", key, err)
+		}
 		out = append(out, fieldMatcher{field: field, re: re})
 	}
 	return out, nil
@@ -123,6 +127,8 @@ func normalizeFilterFieldKey(k string) string {
 		return "id"
 	case "file", "filename", "f":
 		return "file"
+	case "suite_status", "suite":
+		return "suite_status"
 	default:
 		return ""
 	}
@@ -155,6 +161,8 @@ func getFieldValueForFilter(t Test, field string) string {
 		return t.TraceID
 	case "file":
 		return t.FileName
+	case "suite_status":
+		return t.SuiteStatus
 	default:
 		return ""
 	}
@@ -172,6 +180,23 @@ func extractGraphQLOperationName(displayName string) string {
 	return displayName
 }
 
+// ExtractSuiteStatusFromFilter extracts the suite_status value from a filter string.
+// Returns the value and true if found, empty string and false otherwise.
+func ExtractSuiteStatusFromFilter(filter string) (string, bool) {
+	matchers, err := parseFieldedFilter(filter)
+	if err != nil {
+		return "", false
+	}
+	for _, m := range matchers {
+		if m.field == "suite_status" {
+			val := m.re.String()
+			val = strings.TrimPrefix(val, "(?i)")
+			return val, true
+		}
+	}
+	return "", false
+}
+
 // FilterLocalTestsForExecution filters out local tests with HTTP status >= 300.
 // These tests are skipped for replay but their spans remain available for mock matching.
 // Returns (testsToExecute, excludedCount).