diff --git a/cmd/kubectl-ate/README.md b/cmd/kubectl-ate/README.md index 54712e826..72415b249 100644 --- a/cmd/kubectl-ate/README.md +++ b/cmd/kubectl-ate/README.md @@ -83,9 +83,15 @@ kubectl ate get actors -a # List actors across all atespaces kubectl ate get actors -A +# List actors, then print actors again when their state changes +kubectl ate get actors -A --watch + # Get a specific actor by name and output as raw YAML kubectl ate get actor --atespace -o yaml +# Get an actor, then watch it for lifecycle changes +kubectl ate get actor --atespace -w + # List all physical workers and see which actors are assigned to them kubectl ate get workers ``` @@ -162,6 +168,9 @@ kubectl ate suspend actor my-actor -a # Delete an actor. kubectl ate delete actor my-actor -a + +# Wait up to 60 seconds for an actor to finish resuming. +kubectl ate wait actor my-actor -a --for=status=running --timeout=60s ``` ### Logs diff --git a/cmd/kubectl-ate/internal/cmd/get_actors.go b/cmd/kubectl-ate/internal/cmd/get_actors.go index 0f90a71cf..9bd278045 100644 --- a/cmd/kubectl-ate/internal/cmd/get_actors.go +++ b/cmd/kubectl-ate/internal/cmd/get_actors.go @@ -15,6 +15,7 @@ package cmd import ( + "context" "fmt" "github.com/agent-substrate/substrate/cmd/kubectl-ate/internal/printer" @@ -26,6 +27,7 @@ import ( var ( getActorsAtespaceFlag string getActorsAllAtespaces bool + getActorsWatch bool ) var getActorsCmd = &cobra.Command{ @@ -53,15 +55,18 @@ var getActorsCmd = &cobra.Command{ return fmt.Errorf("--atespace is required when getting actors") } - actors := make([]*ateapipb.Actor, 0, len(args)) - for _, actorName := range args { - resp, err := apiClient.GetActor(ctx, &ateapipb.GetActorRequest{Actor: &ateapipb.ObjectRef{Atespace: getActorsAtespaceFlag, Name: actorName}}) - if err != nil { - return fmt.Errorf("failed to get actor %q: %w", actorName, err) + fetch := func(ctx context.Context) ([]*ateapipb.Actor, error) { + actors := make([]*ateapipb.Actor, 0, len(args)) + for _, actorName := range args { + resp, err := apiClient.GetActor(ctx, &ateapipb.GetActorRequest{Actor: &ateapipb.ObjectRef{Atespace: getActorsAtespaceFlag, Name: actorName}}) + if err != nil { + return nil, fmt.Errorf("failed to get actor %q: %w", actorName, err) + } + actors = append(actors, resp) } - actors = append(actors, resp) + return actors, nil } - return printer.PrintActors(actors, outputFmt) + return printOrWatch(ctx, cmd.OutOrStdout(), getActorsWatch, outputFmt, fetch, printer.PrintActorsTo, actorWatchKey) } // Listing requires exactly one of --atespace (one atespace) or -A (all @@ -74,32 +79,30 @@ var getActorsCmd = &cobra.Command{ } // 3. Handle List All Actors - var allActors []*ateapipb.Actor - pageToken := "" - - for { - resp, err := apiClient.ListActors(ctx, &ateapipb.ListActorsRequest{ - PageSize: 1000, - PageToken: pageToken, - Atespace: getActorsAtespaceFlag, + fetch := func(ctx context.Context) ([]*ateapipb.Actor, error) { + return fetchAllPages(ctx, func(ctx context.Context, pageToken string) ([]*ateapipb.Actor, string, error) { + resp, err := apiClient.ListActors(ctx, &ateapipb.ListActorsRequest{ + PageSize: 1000, + PageToken: pageToken, + Atespace: getActorsAtespaceFlag, + }) + if err != nil { + return nil, "", fmt.Errorf("failed to list actors: %w", err) + } + return resp.GetActors(), resp.GetNextPageToken(), nil }) - if err != nil { - return fmt.Errorf("failed to list actors: %w", err) - } - allActors = append(allActors, resp.GetActors()...) - - pageToken = resp.GetNextPageToken() - if pageToken == "" { - break - } } - - return printer.PrintActors(allActors, outputFmt) + return printOrWatch(ctx, cmd.OutOrStdout(), getActorsWatch, outputFmt, fetch, printer.PrintActorsTo, actorWatchKey) }, } +func actorWatchKey(actor *ateapipb.Actor) string { + return actor.GetMetadata().GetAtespace() + "/" + actor.GetMetadata().GetName() +} + func init() { getActorsCmd.Flags().StringVarP(&getActorsAtespaceFlag, "atespace", "a", "", "Atespace to list/get actors in. Required when getting actors; for listing, use this or -A.") getActorsCmd.Flags().BoolVarP(&getActorsAllAtespaces, "all-atespaces", "A", false, "List actors across all atespaces (listing only; mutually exclusive with --atespace)") + getActorsCmd.Flags().BoolVarP(&getActorsWatch, "watch", "w", false, "After listing/getting the requested actor(s), watch for changes") getCmd.AddCommand(getActorsCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/get_atespaces.go b/cmd/kubectl-ate/internal/cmd/get_atespaces.go index c25d45d2d..be851d0c3 100644 --- a/cmd/kubectl-ate/internal/cmd/get_atespaces.go +++ b/cmd/kubectl-ate/internal/cmd/get_atespaces.go @@ -15,6 +15,7 @@ package cmd import ( + "context" "fmt" "github.com/agent-substrate/substrate/cmd/kubectl-ate/internal/printer" @@ -23,6 +24,8 @@ import ( "github.com/spf13/cobra" ) +var getAtespacesWatch bool + var getAtespacesCmd = &cobra.Command{ Use: "atespaces [name ...]", Aliases: []string{"atespace"}, @@ -36,38 +39,41 @@ var getAtespacesCmd = &cobra.Command{ defer apiClient.Close() if len(args) > 0 { - atespaces := make([]*ateapipb.Atespace, 0, len(args)) - for _, atespaceName := range args { - resp, err := apiClient.GetAtespace(ctx, &ateapipb.GetAtespaceRequest{Atespace: &ateapipb.ObjectRef{Name: atespaceName}}) - if err != nil { - return fmt.Errorf("failed to get atespace %q: %w", atespaceName, err) + fetch := func(ctx context.Context) ([]*ateapipb.Atespace, error) { + atespaces := make([]*ateapipb.Atespace, 0, len(args)) + for _, atespaceName := range args { + resp, err := apiClient.GetAtespace(ctx, &ateapipb.GetAtespaceRequest{Atespace: &ateapipb.ObjectRef{Name: atespaceName}}) + if err != nil { + return nil, fmt.Errorf("failed to get atespace %q: %w", atespaceName, err) + } + atespaces = append(atespaces, resp) } - atespaces = append(atespaces, resp) + return atespaces, nil } - return printer.PrintAtespaces(atespaces, outputFmt) + return printOrWatch(ctx, cmd.OutOrStdout(), getAtespacesWatch, outputFmt, fetch, printer.PrintAtespacesTo, atespaceWatchKey) } - var allAtespaces []*ateapipb.Atespace - pageToken := "" - for { - resp, err := apiClient.ListAtespaces(ctx, &ateapipb.ListAtespacesRequest{ - PageSize: 1000, - PageToken: pageToken, + fetch := func(ctx context.Context) ([]*ateapipb.Atespace, error) { + return fetchAllPages(ctx, func(ctx context.Context, pageToken string) ([]*ateapipb.Atespace, string, error) { + resp, err := apiClient.ListAtespaces(ctx, &ateapipb.ListAtespacesRequest{ + PageSize: 1000, + PageToken: pageToken, + }) + if err != nil { + return nil, "", fmt.Errorf("failed to list atespaces: %w", err) + } + return resp.GetAtespaces(), resp.GetNextPageToken(), nil }) - if err != nil { - return fmt.Errorf("failed to list atespaces: %w", err) - } - allAtespaces = append(allAtespaces, resp.GetAtespaces()...) - - pageToken = resp.GetNextPageToken() - if pageToken == "" { - break - } } - return printer.PrintAtespaces(allAtespaces, outputFmt) + return printOrWatch(ctx, cmd.OutOrStdout(), getAtespacesWatch, outputFmt, fetch, printer.PrintAtespacesTo, atespaceWatchKey) }, } +func atespaceWatchKey(atespace *ateapipb.Atespace) string { + return atespace.GetMetadata().GetName() +} + func init() { + getAtespacesCmd.Flags().BoolVarP(&getAtespacesWatch, "watch", "w", false, "After listing/getting the requested atespace(s), watch for changes") getCmd.AddCommand(getAtespacesCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/get_test.go b/cmd/kubectl-ate/internal/cmd/get_test.go index c2e6639fe..1ffa181a8 100644 --- a/cmd/kubectl-ate/internal/cmd/get_test.go +++ b/cmd/kubectl-ate/internal/cmd/get_test.go @@ -49,3 +49,18 @@ func TestGetCommandArgs(t *testing.T) { }) } } + +func TestGetCommandsHaveWatchFlag(t *testing.T) { + commands := []*cobra.Command{getActorsCmd, getAtespacesCmd, getWorkersCmd} + for _, command := range commands { + t.Run(command.Name(), func(t *testing.T) { + flag := command.Flags().Lookup("watch") + if flag == nil { + t.Fatal("--watch flag is not registered") + } + if flag.Shorthand != "w" { + t.Errorf("--watch shorthand = %q, want w", flag.Shorthand) + } + }) + } +} diff --git a/cmd/kubectl-ate/internal/cmd/get_workers.go b/cmd/kubectl-ate/internal/cmd/get_workers.go index f821d79f3..0f3df7971 100644 --- a/cmd/kubectl-ate/internal/cmd/get_workers.go +++ b/cmd/kubectl-ate/internal/cmd/get_workers.go @@ -15,6 +15,7 @@ package cmd import ( + "context" "fmt" "github.com/agent-substrate/substrate/cmd/kubectl-ate/internal/printer" @@ -23,6 +24,8 @@ import ( "github.com/spf13/cobra" ) +var getWorkersWatch bool + var getWorkersCmd = &cobra.Command{ Use: "workers", Aliases: []string{"worker"}, @@ -36,27 +39,27 @@ var getWorkersCmd = &cobra.Command{ } defer apiClient.Close() - var allWorkers []*ateapipb.Worker - pageToken := "" - for { - resp, err := apiClient.ListWorkers(ctx, &ateapipb.ListWorkersRequest{ - PageSize: 1000, - PageToken: pageToken, + fetch := func(ctx context.Context) ([]*ateapipb.Worker, error) { + return fetchAllPages(ctx, func(ctx context.Context, pageToken string) ([]*ateapipb.Worker, string, error) { + resp, err := apiClient.ListWorkers(ctx, &ateapipb.ListWorkersRequest{ + PageSize: 1000, + PageToken: pageToken, + }) + if err != nil { + return nil, "", fmt.Errorf("failed to list workers: %w", err) + } + return resp.GetWorkers(), resp.GetNextPageToken(), nil }) - 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) + return printOrWatch(ctx, cmd.OutOrStdout(), getWorkersWatch, outputFmt, fetch, printer.PrintWorkersTo, workerWatchKey) }, } +func workerWatchKey(worker *ateapipb.Worker) string { + return worker.GetWorkerNamespace() + "/" + worker.GetWorkerPool() + "/" + worker.GetWorkerPod() +} + func init() { + getWorkersCmd.Flags().BoolVarP(&getWorkersWatch, "watch", "w", false, "After listing the requested workers, watch for changes") getCmd.AddCommand(getWorkersCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/pagination.go b/cmd/kubectl-ate/internal/cmd/pagination.go new file mode 100644 index 000000000..5e582fd5f --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/pagination.go @@ -0,0 +1,33 @@ +// 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" + +func fetchAllPages[M any](ctx context.Context, fetchPage func(context.Context, string) ([]M, string, error)) ([]M, error) { + var all []M + pageToken := "" + for { + resources, nextPageToken, err := fetchPage(ctx, pageToken) + if err != nil { + return nil, err + } + all = append(all, resources...) + if nextPageToken == "" { + return all, nil + } + pageToken = nextPageToken + } +} diff --git a/cmd/kubectl-ate/internal/cmd/pagination_test.go b/cmd/kubectl-ate/internal/cmd/pagination_test.go new file mode 100644 index 000000000..aebf87772 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/pagination_test.go @@ -0,0 +1,55 @@ +// 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" + "errors" + "slices" + "testing" +) + +func TestFetchAllPages(t *testing.T) { + var tokens []string + got, err := fetchAllPages(context.Background(), func(_ context.Context, token string) ([]string, string, error) { + tokens = append(tokens, token) + if token == "" { + return []string{"one", "two"}, "next", nil + } + return []string{"three"}, "", nil + }) + if err != nil { + t.Fatalf("fetchAllPages() error = %v", err) + } + if want := []string{"one", "two", "three"}; !slices.Equal(got, want) { + t.Errorf("fetchAllPages() = %v, want %v", got, want) + } + if want := []string{"", "next"}; !slices.Equal(tokens, want) { + t.Errorf("page tokens = %v, want %v", tokens, want) + } +} + +func TestFetchAllPagesReturnsError(t *testing.T) { + wantErr := errors.New("list failed") + got, err := fetchAllPages(context.Background(), func(context.Context, string) ([]string, string, error) { + return nil, "", wantErr + }) + if !errors.Is(err, wantErr) { + t.Fatalf("fetchAllPages() error = %v, want %v", err, wantErr) + } + if got != nil { + t.Errorf("fetchAllPages() = %v, want nil", got) + } +} diff --git a/cmd/kubectl-ate/internal/cmd/wait.go b/cmd/kubectl-ate/internal/cmd/wait.go new file mode 100644 index 000000000..493eeae11 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/wait.go @@ -0,0 +1,26 @@ +// 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 "github.com/spf13/cobra" + +var waitCmd = &cobra.Command{ + Use: "wait", + Short: "Wait for a specific condition on one or more resources", +} + +func init() { + rootCmd.AddCommand(waitCmd) +} diff --git a/cmd/kubectl-ate/internal/cmd/wait_actor.go b/cmd/kubectl-ate/internal/cmd/wait_actor.go new file mode 100644 index 000000000..74c0cd04a --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/wait_actor.go @@ -0,0 +1,181 @@ +// 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" + "io" + "strings" + "time" + + "github.com/agent-substrate/substrate/internal/ateclient" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/spf13/cobra" + "google.golang.org/grpc" +) + +const defaultWaitPollInterval = 500 * time.Millisecond + +var ( + waitActorAtespace string + waitActorFor string + waitActorTimeout time.Duration +) + +var waitActorCmd = &cobra.Command{ + Use: "actor ", + Aliases: []string{"actors"}, + Short: "Wait for an actor to reach a status", + Example: ` # Wait up to 60 seconds for an actor to reach the running status + kubectl ate wait actor my-actor --atespace team-a --for=status=running --timeout=60s + + # Check whether an actor is currently suspended without polling + kubectl ate wait actor my-actor -a team-a --for=status=suspended --timeout=0s`, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + targetStatus, err := parseWaitActorStatus(waitActorFor) + if err != nil { + return err + } + if waitActorTimeout < 0 { + return fmt.Errorf("--timeout must be greater than or equal to zero") + } + + apiClient, err := ateclient.NewClient(cmd.Context(), kubeconfig, k8sContext, endpoint, traceEnabled) + if err != nil { + return fmt.Errorf("failed to connect to ate-api-server: %w", err) + } + defer apiClient.Close() + + runner := &WaitActorRunner{ + apiClient: apiClient, + out: cmd.OutOrStdout(), + pollInterval: defaultWaitPollInterval, + } + return runner.Run(cmd.Context(), waitActorAtespace, args[0], targetStatus, waitActorTimeout) + }, +} + +// WaitActorRunner polls an actor until it reaches the requested status. +type WaitActorRunner struct { + apiClient ActorGetter + out io.Writer + pollInterval time.Duration +} + +type ActorGetter interface { + GetActor(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) +} + +func (r *WaitActorRunner) Run(ctx context.Context, atespace, actorName string, targetStatus ateapipb.Actor_Status, timeout time.Duration) error { + waitCtx := ctx + cancel := func() {} + if timeout > 0 { + waitCtx, cancel = context.WithTimeout(ctx, timeout) + } + defer cancel() + + check := func(checkCtx context.Context) (bool, error) { + actor, err := r.apiClient.GetActor(checkCtx, &ateapipb.GetActorRequest{ + Actor: &ateapipb.ObjectRef{Atespace: atespace, Name: actorName}, + }) + if err != nil { + return false, fmt.Errorf("failed to get actor %q: %w", actorName, err) + } + return actor.GetStatus() == targetStatus, nil + } + + met, err := check(waitCtx) + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + if waitCtx.Err() == context.DeadlineExceeded { + return waitActorTimeoutError(actorName, targetStatus, timeout) + } + return err + } + if met { + fmt.Fprintf(r.out, "actor/%s condition met\n", actorName) + return nil + } + if timeout == 0 { + return waitActorTimeoutError(actorName, targetStatus, timeout) + } + + interval := r.pollInterval + if interval <= 0 { + interval = defaultWaitPollInterval + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-waitCtx.Done(): + if ctx.Err() != nil { + return ctx.Err() + } + return waitActorTimeoutError(actorName, targetStatus, timeout) + case <-ticker.C: + met, err := check(waitCtx) + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + if waitCtx.Err() == context.DeadlineExceeded { + return waitActorTimeoutError(actorName, targetStatus, timeout) + } + return err + } + if met { + fmt.Fprintf(r.out, "actor/%s condition met\n", actorName) + return nil + } + } + } +} + +func parseWaitActorStatus(condition string) (ateapipb.Actor_Status, error) { + const prefix = "status=" + if !strings.HasPrefix(strings.ToLower(condition), prefix) { + return ateapipb.Actor_STATUS_UNSPECIFIED, fmt.Errorf("unsupported --for value %q: expected status=", condition) + } + + statusName := strings.ToUpper(strings.TrimSpace(condition[len(prefix):])) + if !strings.HasPrefix(statusName, "STATUS_") { + statusName = "STATUS_" + statusName + } + statusValue, ok := ateapipb.Actor_Status_value[statusName] + if !ok || statusName == "STATUS_UNSPECIFIED" { + return ateapipb.Actor_STATUS_UNSPECIFIED, fmt.Errorf("unknown actor status %q", condition[len(prefix):]) + } + return ateapipb.Actor_Status(statusValue), nil +} + +func waitActorTimeoutError(actorName string, targetStatus ateapipb.Actor_Status, timeout time.Duration) error { + status := strings.TrimPrefix(targetStatus.String(), "STATUS_") + return fmt.Errorf("timed out after %s waiting for actor %q to reach status %s", timeout, actorName, status) +} + +func init() { + waitActorCmd.Flags().StringVarP(&waitActorAtespace, "atespace", "a", "", "Atespace the actor lives in") + _ = waitActorCmd.MarkFlagRequired("atespace") + waitActorCmd.Flags().StringVar(&waitActorFor, "for", "", "Condition to wait for, in status= format") + _ = waitActorCmd.MarkFlagRequired("for") + waitActorCmd.Flags().DurationVar(&waitActorTimeout, "timeout", 30*time.Second, "Maximum time to wait; zero checks once without waiting") + waitCmd.AddCommand(waitActorCmd) +} diff --git a/cmd/kubectl-ate/internal/cmd/wait_actor_test.go b/cmd/kubectl-ate/internal/cmd/wait_actor_test.go new file mode 100644 index 000000000..d10b345b5 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/wait_actor_test.go @@ -0,0 +1,178 @@ +// 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" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" +) + +func TestWaitActorCommandContract(t *testing.T) { + for _, test := range []struct { + name string + args []string + wantErr bool + }{ + {name: "one actor", args: []string{"actor-1"}}, + {name: "missing actor", wantErr: true}, + {name: "multiple actors", args: []string{"actor-1", "actor-2"}, wantErr: true}, + } { + t.Run(test.name, func(t *testing.T) { + err := waitActorCmd.Args(waitActorCmd, test.args) + if (err != nil) != test.wantErr { + t.Fatalf("Args(%q) error = %v, wantErr %t", test.args, err, test.wantErr) + } + }) + } + + for _, name := range []string{"atespace", "for", "timeout"} { + if flag := waitActorCmd.Flags().Lookup(name); flag == nil { + t.Errorf("--%s flag is not registered", name) + } + } + if flag := waitActorCmd.Flags().ShorthandLookup("a"); flag == nil || flag.Name != "atespace" { + t.Errorf("-a flag = %v, want --atespace", flag) + } +} + +func TestParseWaitActorStatus(t *testing.T) { + tests := []struct { + condition string + want ateapipb.Actor_Status + wantErr bool + }{ + {condition: "status=running", want: ateapipb.Actor_STATUS_RUNNING}, + {condition: "STATUS=STATUS_SUSPENDED", want: ateapipb.Actor_STATUS_SUSPENDED}, + {condition: "condition=running", wantErr: true}, + {condition: "status=unknown", wantErr: true}, + {condition: "status=unspecified", wantErr: true}, + } + + for _, test := range tests { + t.Run(test.condition, func(t *testing.T) { + got, err := parseWaitActorStatus(test.condition) + if (err != nil) != test.wantErr { + t.Fatalf("parseWaitActorStatus(%q) error = %v, wantErr %t", test.condition, err, test.wantErr) + } + if got != test.want { + t.Errorf("parseWaitActorStatus(%q) = %v, want %v", test.condition, got, test.want) + } + }) + } +} + +func TestWaitActorRunnerTransition(t *testing.T) { + statuses := []ateapipb.Actor_Status{ + ateapipb.Actor_STATUS_RESUMING, + ateapipb.Actor_STATUS_RUNNING, + } + var calls int + client := &mockAteAPIClient{ + GetActorFunc: func(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + if got := in.GetActor().GetAtespace(); got != "team-a" { + t.Errorf("atespace = %q, want team-a", got) + } + status := statuses[calls] + calls++ + return &ateapipb.Actor{Status: status}, nil + }, + } + var out bytes.Buffer + runner := &WaitActorRunner{apiClient: client, out: &out, pollInterval: time.Millisecond} + + if err := runner.Run(context.Background(), "team-a", "agent-1", ateapipb.Actor_STATUS_RUNNING, time.Second); err != nil { + t.Fatalf("Run() error = %v", err) + } + if calls != 2 { + t.Errorf("GetActor calls = %d, want 2", calls) + } + if got, want := out.String(), "actor/agent-1 condition met\n"; got != want { + t.Errorf("output = %q, want %q", got, want) + } +} + +func TestWaitActorRunnerAlreadyAtTarget(t *testing.T) { + var calls int + client := &mockAteAPIClient{ + GetActorFunc: func(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + calls++ + return &ateapipb.Actor{Status: ateapipb.Actor_STATUS_RUNNING}, nil + }, + } + var out bytes.Buffer + runner := &WaitActorRunner{apiClient: client, out: &out} + + if err := runner.Run(context.Background(), "team-a", "agent-1", ateapipb.Actor_STATUS_RUNNING, time.Second); err != nil { + t.Fatalf("Run() error = %v", err) + } + if calls != 1 { + t.Errorf("GetActor calls = %d, want 1", calls) + } + if got, want := out.String(), "actor/agent-1 condition met\n"; got != want { + t.Errorf("output = %q, want %q", got, want) + } +} + +func TestWaitActorRunnerReturnsGetError(t *testing.T) { + wantErr := errors.New("get failed") + client := &mockAteAPIClient{ + GetActorFunc: func(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + return nil, wantErr + }, + } + runner := &WaitActorRunner{apiClient: client, out: io.Discard} + + err := runner.Run(context.Background(), "team-a", "agent-1", ateapipb.Actor_STATUS_RUNNING, time.Second) + if !errors.Is(err, wantErr) { + t.Fatalf("Run() error = %v, want wrapped %v", err, wantErr) + } +} + +func TestWaitActorRunnerZeroTimeout(t *testing.T) { + client := &mockAteAPIClient{ + GetActorFunc: func(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + return &ateapipb.Actor{Status: ateapipb.Actor_STATUS_SUSPENDED}, nil + }, + } + runner := &WaitActorRunner{apiClient: client, out: &bytes.Buffer{}} + + err := runner.Run(context.Background(), "team-a", "agent-1", ateapipb.Actor_STATUS_RUNNING, 0) + if err == nil || !strings.Contains(err.Error(), "timed out after 0s") { + t.Fatalf("Run() error = %v, want zero-timeout error", err) + } +} + +func TestWaitActorRunnerTimeoutCancelsGet(t *testing.T) { + client := &mockAteAPIClient{ + GetActorFunc: func(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + } + runner := &WaitActorRunner{apiClient: client, out: &bytes.Buffer{}} + + err := runner.Run(context.Background(), "team-a", "agent-1", ateapipb.Actor_STATUS_RUNNING, time.Millisecond) + if err == nil || !strings.Contains(err.Error(), "timed out after 1ms") { + t.Fatalf("Run() error = %v, want timeout error", err) + } +} diff --git a/cmd/kubectl-ate/internal/cmd/watch_runner.go b/cmd/kubectl-ate/internal/cmd/watch_runner.go new file mode 100644 index 000000000..d965e871a --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/watch_runner.go @@ -0,0 +1,138 @@ +// 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" + "io" + "strings" + "time" + + "google.golang.org/protobuf/proto" + "k8s.io/apimachinery/pkg/util/wait" +) + +const defaultGetWatchInterval = time.Second + +type getWatchRunner[M proto.Message] struct { + fetch func(context.Context) ([]M, error) + print func(io.Writer, []M) error + key func(M) string + out io.Writer + format string + pollInterval time.Duration +} + +func printOrWatch[M proto.Message]( + ctx context.Context, + out io.Writer, + watch bool, + format string, + fetch func(context.Context) ([]M, error), + print func(io.Writer, []M, string) error, + key func(M) string, +) error { + printFormatted := func(out io.Writer, resources []M) error { + return print(out, resources, format) + } + if watch { + return (&getWatchRunner[M]{ + fetch: fetch, + print: printFormatted, + key: key, + out: out, + format: format, + }).Run(ctx) + } + + resources, err := fetch(ctx) + if err != nil { + return err + } + return printFormatted(out, resources) +} + +func (r *getWatchRunner[M]) Run(ctx context.Context) error { + resources, err := r.fetch(ctx) + if err != nil { + return err + } + if err := r.print(r.out, resources); err != nil { + return err + } + previous := r.index(resources) + + interval := r.pollInterval + if interval <= 0 { + interval = defaultGetWatchInterval + } + return wait.PollUntilContextCancel(ctx, interval, false, func(ctx context.Context) (bool, error) { + resources, err := r.fetch(ctx) + if err != nil { + return false, err + } + current := r.index(resources) + changed := changedResources(previous, current) + if len(changed) > 0 { + if err := r.printUpdate(changed); err != nil { + return false, err + } + } + previous = current + return false, nil + }) +} + +func (r *getWatchRunner[M]) index(resources []M) map[string]M { + indexed := make(map[string]M, len(resources)) + for _, resource := range resources { + indexed[r.key(resource)] = resource + } + return indexed +} + +func changedResources[M proto.Message](previous, current map[string]M) []M { + changed := make([]M, 0) + for key, resource := range current { + old, found := previous[key] + if !found || !proto.Equal(old, resource) { + changed = append(changed, resource) + } + } + for key, resource := range previous { + if _, found := current[key]; !found { + changed = append(changed, resource) + } + } + return changed +} + +func (r *getWatchRunner[M]) printUpdate(resources []M) error { + if r.format != "table" { + return r.print(r.out, resources) + } + + var buf bytes.Buffer + if err := r.print(&buf, resources); err != nil { + return err + } + output := buf.String() + if newline := strings.IndexByte(output, '\n'); newline >= 0 { + output = output[newline+1:] + } + _, err := io.WriteString(r.out, output) + return err +} diff --git a/cmd/kubectl-ate/internal/cmd/watch_runner_test.go b/cmd/kubectl-ate/internal/cmd/watch_runner_test.go new file mode 100644 index 000000000..8921c00b7 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/watch_runner_test.go @@ -0,0 +1,147 @@ +// 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" + "errors" + "fmt" + "io" + "strings" + "testing" + "time" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" +) + +func TestPrintOrWatchPrintsOnce(t *testing.T) { + actor := &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Name: "actor-1"}} + var fetchCalls int + var out bytes.Buffer + + err := printOrWatch( + context.Background(), + &out, + false, + "yaml", + func(context.Context) ([]*ateapipb.Actor, error) { + fetchCalls++ + return []*ateapipb.Actor{actor}, nil + }, + func(out io.Writer, actors []*ateapipb.Actor, format string) error { + if format != "yaml" { + t.Errorf("format = %q, want yaml", format) + } + fmt.Fprint(out, actors[0].GetMetadata().GetName()) + return nil + }, + nil, + ) + if err != nil { + t.Fatalf("printOrWatch() error = %v", err) + } + if fetchCalls != 1 { + t.Errorf("fetch calls = %d, want 1", fetchCalls) + } + if got, want := out.String(), "actor-1"; got != want { + t.Errorf("output = %q, want %q", got, want) + } +} + +func TestPrintOrWatchReturnsFetchError(t *testing.T) { + wantErr := errors.New("fetch failed") + err := printOrWatch( + context.Background(), + io.Discard, + false, + "table", + func(context.Context) ([]*ateapipb.Actor, error) { return nil, wantErr }, + func(io.Writer, []*ateapipb.Actor, string) error { return nil }, + nil, + ) + if !errors.Is(err, wantErr) { + t.Fatalf("printOrWatch() error = %v, want %v", err, wantErr) + } +} + +func TestGetWatchRunnerPrintsOnlyChanges(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + responses := [][]*ateapipb.Actor{ + {{Metadata: &ateapipb.ResourceMetadata{Name: "actor-1", Version: 1}}}, + {{Metadata: &ateapipb.ResourceMetadata{Name: "actor-1", Version: 1}}}, + {{Metadata: &ateapipb.ResourceMetadata{Name: "actor-1", Version: 2}}}, + } + var fetchCalls int + var out bytes.Buffer + runner := &getWatchRunner[*ateapipb.Actor]{ + fetch: func(context.Context) ([]*ateapipb.Actor, error) { + response := responses[fetchCalls] + fetchCalls++ + if fetchCalls == len(responses) { + cancel() + } + return response, nil + }, + print: func(out io.Writer, actors []*ateapipb.Actor) error { + fmt.Fprintln(out, "HEADER") + for _, actor := range actors { + fmt.Fprintf(out, "%s %d\n", actor.GetMetadata().GetName(), actor.GetMetadata().GetVersion()) + } + return nil + }, + key: func(actor *ateapipb.Actor) string { return actor.GetMetadata().GetName() }, + out: &out, + format: "table", + pollInterval: time.Millisecond, + } + + err := runner.Run(ctx) + if err != context.Canceled { + t.Fatalf("Run() error = %v, want context.Canceled", err) + } + if got, want := out.String(), "HEADER\nactor-1 1\nactor-1 2\n"; got != want { + t.Errorf("output = %q, want %q", got, want) + } +} + +func TestChangedResourcesIncludesAddedChangedAndDeleted(t *testing.T) { + previous := map[string]*ateapipb.Actor{ + "unchanged": {Metadata: &ateapipb.ResourceMetadata{Name: "unchanged", Version: 1}}, + "changed": {Metadata: &ateapipb.ResourceMetadata{Name: "changed", Version: 1}}, + "deleted": {Metadata: &ateapipb.ResourceMetadata{Name: "deleted", Version: 1}}, + } + current := map[string]*ateapipb.Actor{ + "unchanged": {Metadata: &ateapipb.ResourceMetadata{Name: "unchanged", Version: 1}}, + "changed": {Metadata: &ateapipb.ResourceMetadata{Name: "changed", Version: 2}}, + "added": {Metadata: &ateapipb.ResourceMetadata{Name: "added", Version: 1}}, + } + + changed := changedResources(previous, current) + var names []string + for _, actor := range changed { + names = append(names, actor.GetMetadata().GetName()) + } + for _, name := range []string{"added", "changed", "deleted"} { + if !strings.Contains(strings.Join(names, ","), name) { + t.Errorf("changed resources %v do not include %q", names, name) + } + } + if len(names) != 3 { + t.Errorf("changed resource count = %d, want 3", len(names)) + } +}