From e8f852af86bb30c95b747df91da38409d11ce23a Mon Sep 17 00:00:00 2001 From: Haven Xia Date: Wed, 29 Jul 2026 09:37:10 -0700 Subject: [PATCH] kubectl-ate: filter get workers by namespace, assigned actor's atespace and selector --- cmd/kubectl-ate/README.md | 6 + cmd/kubectl-ate/internal/cmd/get_workers.go | 87 +++++++++----- .../internal/cmd/get_workers_test.go | 107 ++++++++++++++++++ cmd/kubectl-ate/internal/cmd/top_workers.go | 56 +-------- cmd/kubectl-ate/internal/cmd/workers.go | 79 +++++++++++++ 5 files changed, 257 insertions(+), 78 deletions(-) create mode 100644 cmd/kubectl-ate/internal/cmd/get_workers_test.go create mode 100644 cmd/kubectl-ate/internal/cmd/workers.go diff --git a/cmd/kubectl-ate/README.md b/cmd/kubectl-ate/README.md index a83fd8745..af69c8a63 100644 --- a/cmd/kubectl-ate/README.md +++ b/cmd/kubectl-ate/README.md @@ -89,6 +89,12 @@ kubectl ate get actor --atespace -o yaml # List all physical workers and see which actors are assigned to them kubectl ate get workers + +# Filter workers by Kubernetes namespace, assigned-actor atespace, or +# worker pool labels (same flags as `top workers`) +kubectl ate get workers -n +kubectl ate get workers -a +kubectl ate get workers -l ``` > **Note:** `get actors` requires either `--atespace ` / `-a ` (one atespace) or `-A`/`--all-atespaces` (all atespaces) — there is no default atespace. Getting a single actor always requires `--atespace`/`-a`, since an actor is addressed by `(atespace, name)`. `-a` (lower-case) scopes to one atespace; `-A` (upper-case) spans all. diff --git a/cmd/kubectl-ate/internal/cmd/get_workers.go b/cmd/kubectl-ate/internal/cmd/get_workers.go index f821d79f3..6f2956e43 100644 --- a/cmd/kubectl-ate/internal/cmd/get_workers.go +++ b/cmd/kubectl-ate/internal/cmd/get_workers.go @@ -15,48 +15,79 @@ package cmd import ( + "context" "fmt" + "io" + "os" "github.com/agent-substrate/substrate/cmd/kubectl-ate/internal/printer" "github.com/agent-substrate/substrate/internal/ateclient" - "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/spf13/cobra" ) +var ( + getWorkerNamespaceFlag string + getWorkerAtespaceFlag string + getWorkerSelectorFlag string +) + var getWorkersCmd = &cobra.Command{ Use: "workers", Aliases: []string{"worker"}, Short: "List all workers", Args: cobra.NoArgs, - RunE: func(cmd *cobra.Command, args []string) error { - ctx := cmd.Context() - apiClient, err := ateclient.NewClient(ctx, kubeconfig, k8sContext, endpoint, traceEnabled) - if err != nil { - return fmt.Errorf("failed to connect to ate-api-server: %w", err) - } - defer apiClient.Close() - - var allWorkers []*ateapipb.Worker - pageToken := "" - for { - resp, err := apiClient.ListWorkers(ctx, &ateapipb.ListWorkersRequest{ - PageSize: 1000, - PageToken: pageToken, - }) - if err != nil { - return fmt.Errorf("failed to list workers: %w", err) - } - allWorkers = append(allWorkers, resp.GetWorkers()...) - - pageToken = resp.GetNextPageToken() - if pageToken == "" { - break - } - } - return printer.PrintWorkers(allWorkers, outputFmt) - }, + RunE: runGetWorkers, } func init() { + getWorkersCmd.Flags().StringVarP(&getWorkerNamespaceFlag, "namespace", "n", "", "Scope output to a specific Kubernetes namespace") + getWorkersCmd.Flags().StringVarP(&getWorkerAtespaceFlag, "atespace", "a", "", "Filter worker pods hosting actors in a specific atespace") + getWorkersCmd.Flags().StringVarP(&getWorkerSelectorFlag, "selector", "l", "", "Filter by worker pool labels") getCmd.AddCommand(getWorkersCmd) } + +// GetWorkersRunner executes the get workers command logic. +type GetWorkersRunner struct { + workerLister WorkerLister + namespace string + atespace string + selector string + outputFmt string + out io.Writer +} + +func (r *GetWorkersRunner) Run(ctx context.Context) error { + workers, err := listAllWorkers(ctx, r.workerLister) + if err != nil { + return err + } + filtered, err := filterWorkers(workers, r.namespace, r.atespace, r.selector) + if err != nil { + return err + } + + outWriter := r.out + if outWriter == nil { + outWriter = os.Stdout + } + return printer.PrintWorkersTo(outWriter, filtered, r.outputFmt) +} + +func runGetWorkers(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + apiClient, err := ateclient.NewClient(ctx, kubeconfig, k8sContext, endpoint, traceEnabled) + if err != nil { + return fmt.Errorf("failed to connect to ate-api-server: %w", err) + } + defer apiClient.Close() + + runner := &GetWorkersRunner{ + workerLister: apiClient, + namespace: getWorkerNamespaceFlag, + atespace: getWorkerAtespaceFlag, + selector: getWorkerSelectorFlag, + outputFmt: outputFmt, + out: os.Stdout, + } + return runner.Run(ctx) +} diff --git a/cmd/kubectl-ate/internal/cmd/get_workers_test.go b/cmd/kubectl-ate/internal/cmd/get_workers_test.go new file mode 100644 index 000000000..8cf288f28 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/get_workers_test.go @@ -0,0 +1,107 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "bytes" + "context" + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/go-cmp/cmp" +) + +func TestGetWorkersRunner_Filters(t *testing.T) { + workers := []*ateapipb.Worker{ + { + WorkerNamespace: "ns-1", + WorkerPool: "counter", + WorkerPod: "pod-1", + Assignment: &ateapipb.Assignment{ + ActorTemplate: &ateapipb.KubeNamespacedObjectRef{Namespace: "ns-1", Name: "counter"}, + Actor: &ateapipb.ObjectRef{Atespace: "space-a", Name: "actor-a"}, + }, + Labels: map[string]string{"ate.dev/worker-pool": "counter"}, + }, + { + WorkerNamespace: "ns-1", + WorkerPool: "other", + WorkerPod: "pod-2", + Labels: map[string]string{"ate.dev/worker-pool": "other"}, + }, + { + WorkerNamespace: "ns-2", + WorkerPool: "counter", + WorkerPod: "pod-3", + Assignment: &ateapipb.Assignment{ + ActorTemplate: &ateapipb.KubeNamespacedObjectRef{Namespace: "ns-2", Name: "counter"}, + Actor: &ateapipb.ObjectRef{Atespace: "space-b", Name: "actor-b"}, + }, + Labels: map[string]string{"ate.dev/worker-pool": "counter"}, + }, + } + + header := "NAMESPACE POOL POD STATUS ASSIGNED ACTOR\n" + row1 := "ns-1 counter pod-1 ASSIGNED ns-1/counter/space-a/actor-a\n" + row2 := "ns-1 other pod-2 FREE \n" + row3 := "ns-2 counter pod-3 ASSIGNED ns-2/counter/space-b/actor-b\n" + + tests := []struct { + name string + namespace string + atespace string + selector string + expected string + }{ + {name: "no filter", expected: header + row1 + row2 + row3}, + {name: "namespace", namespace: "ns-1", expected: header + row1 + row2}, + {name: "atespace", atespace: "space-a", expected: header + row1}, + // With no matching rows the tabwriter sizes columns to the header alone. + {name: "atespace excludes free workers", atespace: "no-such-space", expected: "NAMESPACE POOL POD STATUS ASSIGNED ACTOR\n"}, + {name: "selector", selector: "ate.dev/worker-pool=counter", expected: header + row1 + row3}, + {name: "combined", namespace: "ns-1", selector: "ate.dev/worker-pool=counter", expected: header + row1}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var buf bytes.Buffer + runner := &GetWorkersRunner{ + workerLister: &mockWorkerLister{workers: workers}, + namespace: test.namespace, + atespace: test.atespace, + selector: test.selector, + outputFmt: "table", + out: &buf, + } + if err := runner.Run(context.Background()); err != nil { + t.Fatalf("Run() unexpected error: %v", err) + } + if diff := cmp.Diff(test.expected, buf.String()); diff != "" { + t.Errorf("output mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestGetWorkersRunner_InvalidSelector(t *testing.T) { + runner := &GetWorkersRunner{ + workerLister: &mockWorkerLister{workers: nil}, + selector: "invalid==selector==", + } + + if err := runner.Run(context.Background()); err == nil { + t.Errorf("expected error for invalid label selector, got nil") + } +} diff --git a/cmd/kubectl-ate/internal/cmd/top_workers.go b/cmd/kubectl-ate/internal/cmd/top_workers.go index 238279ea9..43c7b1b6f 100644 --- a/cmd/kubectl-ate/internal/cmd/top_workers.go +++ b/cmd/kubectl-ate/internal/cmd/top_workers.go @@ -22,13 +22,10 @@ import ( "github.com/agent-substrate/substrate/cmd/kubectl-ate/internal/printer" "github.com/agent-substrate/substrate/internal/ateclient" - "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/spf13/cobra" - "google.golang.org/grpc" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/labels" metricsv1beta1 "k8s.io/metrics/pkg/apis/metrics/v1beta1" metricsclient "k8s.io/metrics/pkg/client/clientset/versioned" ) @@ -54,11 +51,6 @@ func init() { topCmd.AddCommand(topWorkersCmd) } -// WorkerLister abstracts ListWorkers RPC calls. -type WorkerLister interface { - ListWorkers(ctx context.Context, req *ateapipb.ListWorkersRequest, opts ...grpc.CallOption) (*ateapipb.ListWorkersResponse, error) -} - // PodMetricsLister abstracts fetching Kubernetes pod metrics. type PodMetricsLister interface { ListPodMetrics(ctx context.Context, namespace string, opts metav1.ListOptions) ([]metricsv1beta1.PodMetrics, error) @@ -88,49 +80,13 @@ type TopWorkersRunner struct { } func (r *TopWorkersRunner) Run(ctx context.Context) error { - var allWorkers []*ateapipb.Worker - pageToken := "" - for { - resp, err := r.workerLister.ListWorkers(ctx, &ateapipb.ListWorkersRequest{ - PageSize: 1000, - PageToken: pageToken, - }) - if err != nil { - return fmt.Errorf("failed to list workers: %w", err) - } - allWorkers = append(allWorkers, resp.GetWorkers()...) - pageToken = resp.GetNextPageToken() - if pageToken == "" { - break - } - } - - var labelSel labels.Selector - if r.selector != "" { - var err error - labelSel, err = labels.Parse(r.selector) - if err != nil { - return fmt.Errorf("invalid label selector %q: %w", r.selector, err) - } + allWorkers, err := listAllWorkers(ctx, r.workerLister) + if err != nil { + return err } - - var filtered []*ateapipb.Worker - for _, w := range allWorkers { - if r.namespace != "" && w.GetWorkerNamespace() != r.namespace { - continue - } - if r.atespace != "" { - wass := w.GetAssignment() - if wass == nil || wass.GetActor() == nil || wass.GetActor().GetAtespace() != r.atespace { - continue - } - } - if labelSel != nil { - if !labelSel.Matches(labels.Set(w.GetLabels())) { - continue - } - } - filtered = append(filtered, w) + filtered, err := filterWorkers(allWorkers, r.namespace, r.atespace, r.selector) + if err != nil { + return err } metricsMap := make(map[string]metricsv1beta1.PodMetrics) diff --git a/cmd/kubectl-ate/internal/cmd/workers.go b/cmd/kubectl-ate/internal/cmd/workers.go new file mode 100644 index 000000000..2d4b6eff0 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/workers.go @@ -0,0 +1,79 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "context" + "fmt" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" + "k8s.io/apimachinery/pkg/labels" +) + +// WorkerLister abstracts ListWorkers RPC calls. +type WorkerLister interface { + ListWorkers(ctx context.Context, req *ateapipb.ListWorkersRequest, opts ...grpc.CallOption) (*ateapipb.ListWorkersResponse, error) +} + +// listAllWorkers pages through ListWorkers and returns all workers. +func listAllWorkers(ctx context.Context, lister WorkerLister) ([]*ateapipb.Worker, error) { + var workers []*ateapipb.Worker + pageToken := "" + for { + resp, err := lister.ListWorkers(ctx, &ateapipb.ListWorkersRequest{ + PageSize: 1000, + PageToken: pageToken, + }) + if err != nil { + return nil, fmt.Errorf("failed to list workers: %w", err) + } + workers = append(workers, resp.GetWorkers()...) + pageToken = resp.GetNextPageToken() + if pageToken == "" { + break + } + } + return workers, nil +} + +// filterWorkers filters workers by Kubernetes namespace, assigned-actor +// atespace, and worker pool label selector. Empty values match everything; +// an atespace filter only matches workers with an assigned actor. +func filterWorkers(workers []*ateapipb.Worker, namespace, atespace, selector string) ([]*ateapipb.Worker, error) { + var labelSel labels.Selector + if selector != "" { + var err error + labelSel, err = labels.Parse(selector) + if err != nil { + return nil, fmt.Errorf("invalid label selector %q: %w", selector, err) + } + } + + var filtered []*ateapipb.Worker + for _, w := range workers { + if namespace != "" && w.GetWorkerNamespace() != namespace { + continue + } + if atespace != "" && w.GetAssignment().GetActor().GetAtespace() != atespace { + continue + } + if labelSel != nil && !labelSel.Matches(labels.Set(w.GetLabels())) { + continue + } + filtered = append(filtered, w) + } + return filtered, nil +}