Skip to content

Commit 5980c6f

Browse files
somanshreddysomanshreddy
committed
cmd: add --wait polling framework (PRINFRA-125)
1 parent f7a59e9 commit 5980c6f

7 files changed

Lines changed: 781 additions & 0 deletions

File tree

‎cmd/heygen/builder.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package main
22

33
import (
4+
"context"
45
"encoding/json"
56
"errors"
67
"fmt"
78
"io"
89
"os"
910
"strconv"
1011
"strings"
12+
"time"
1113

1214
"github.com/heygen-com/heygen-cli/internal/client"
1315
"github.com/heygen-com/heygen-cli/internal/command"
@@ -70,6 +72,34 @@ func buildCobraCommand(spec *command.Spec, ctx *cmdContext) *cobra.Command {
7072
}
7173
}
7274

75+
if spec.PollConfig != nil {
76+
wait, _ := cmd.Flags().GetBool("wait")
77+
if wait {
78+
timeout, _ := cmd.Flags().GetDuration("timeout")
79+
pollCtx, cancel := context.WithTimeout(cmd.Context(), timeout)
80+
defer cancel()
81+
82+
var lastStatus string
83+
result, err := ctx.client.ExecuteAndPoll(pollCtx, spec, inv, client.PollOptions{
84+
Timeout: timeout,
85+
BaseDelay: 2 * time.Second,
86+
MaxDelay: 30 * time.Second,
87+
OnStatus: func(status string, elapsed time.Duration) {
88+
if status == lastStatus {
89+
return
90+
}
91+
fmt.Fprintf(cmd.ErrOrStderr(), "Polling: status=%s (elapsed %s)\n", status, elapsed.Round(time.Second))
92+
lastStatus = status
93+
},
94+
})
95+
if err != nil {
96+
return err
97+
}
98+
99+
return ctx.formatter.Data(result)
100+
}
101+
}
102+
73103
result, err := ctx.client.Execute(spec, inv)
74104
if err != nil {
75105
return err
@@ -87,6 +117,10 @@ func buildCobraCommand(spec *command.Spec, ctx *cmdContext) *cobra.Command {
87117
if spec.Paginated {
88118
cmd.Flags().Bool("all", false, "Fetch all pages (returns flat JSON array instead of API envelope)")
89119
}
120+
if spec.PollConfig != nil {
121+
cmd.Flags().Bool("wait", false, "Poll until the operation completes or fails")
122+
cmd.Flags().Duration("timeout", 10*time.Minute, "Max time to wait when using --wait")
123+
}
90124

91125
// Add -d/--data for commands with JSON request bodies
92126
if spec.BodyEncoding == "json" {

‎cmd/heygen/builder_test.go‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,23 @@ var videoListSpec = &command.Spec{
3535
Examples: []string{"heygen video list --limit 10"},
3636
}
3737

38+
var videoCreateWaitSpec = &command.Spec{
39+
Group: "video",
40+
Name: "create",
41+
Summary: "Create a video",
42+
Endpoint: "/v3/videos",
43+
Method: "POST",
44+
BodyEncoding: "json",
45+
PollConfig: &command.PollConfig{
46+
StatusEndpoint: "/v3/videos/{video_id}",
47+
StatusField: "data.status",
48+
TerminalOK: []string{"completed"},
49+
TerminalFail: []string{"failed", "error"},
50+
IDField: "data.video_id",
51+
},
52+
Examples: []string{"heygen video create --wait"},
53+
}
54+
3855
func TestGenBuilder_VideoList_Success(t *testing.T) {
3956
srv := setupTestServer(t, map[string]testHandler{
4057
"GET /v3/videos": {
@@ -262,6 +279,166 @@ func TestGenBuilder_VideoList_AllPages_Truncated(t *testing.T) {
262279
}
263280
}
264281

282+
func TestGenBuilder_VideoCreate_Wait_Success(t *testing.T) {
283+
var statusCalls int
284+
srv := setupTestServer(t, map[string]testHandler{
285+
"POST /v3/videos": {
286+
StatusCode: 200,
287+
Body: `{"data":{"video_id":"vid_123"}}`,
288+
},
289+
"GET /v3/videos/vid_123": {
290+
StatusCode: 200,
291+
ValidateRequest: func(t *testing.T, r *http.Request) {
292+
t.Helper()
293+
statusCalls++
294+
},
295+
Body: `{"data":{"video_id":"vid_123","status":"processing"}}`,
296+
},
297+
})
298+
defer srv.Close()
299+
300+
originalHandler := srv.Config.Handler
301+
srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
302+
if r.Method == http.MethodGet && r.URL.Path == "/v3/videos/vid_123" {
303+
statusCalls++
304+
w.WriteHeader(http.StatusOK)
305+
if statusCalls < 2 {
306+
_, _ = w.Write([]byte(`{"data":{"video_id":"vid_123","status":"processing"}}`))
307+
return
308+
}
309+
_, _ = w.Write([]byte(`{"data":{"video_id":"vid_123","status":"completed","video_url":"https://cdn.test/video.mp4"}}`))
310+
return
311+
}
312+
originalHandler.ServeHTTP(w, r)
313+
})
314+
315+
res := runGenCommand(t, srv.URL, "test-key", videoCreateWaitSpec, "create", "--wait")
316+
317+
if res.ExitCode != 0 {
318+
t.Fatalf("ExitCode = %d, want 0\nstderr: %s", res.ExitCode, res.Stderr)
319+
}
320+
321+
var parsed map[string]any
322+
if err := json.Unmarshal([]byte(res.Stdout), &parsed); err != nil {
323+
t.Fatalf("stdout is not valid JSON: %v\nstdout: %s", err, res.Stdout)
324+
}
325+
data := parsed["data"].(map[string]any)
326+
if data["status"] != "completed" {
327+
t.Fatalf("status = %v, want completed", data["status"])
328+
}
329+
if strings.Count(res.Stderr, "Polling: status=processing") != 1 {
330+
t.Fatalf("stderr = %s, want exactly one processing line", res.Stderr)
331+
}
332+
}
333+
334+
func TestGenBuilder_VideoCreate_Wait_Failure(t *testing.T) {
335+
srv := setupTestServer(t, map[string]testHandler{
336+
"POST /v3/videos": {
337+
StatusCode: 200,
338+
Body: `{"data":{"video_id":"vid_123"}}`,
339+
},
340+
"GET /v3/videos/vid_123": {
341+
StatusCode: 200,
342+
Body: `{"data":{"video_id":"vid_123","status":"failed"}}`,
343+
},
344+
})
345+
defer srv.Close()
346+
347+
res := runGenCommand(t, srv.URL, "test-key", videoCreateWaitSpec, "create", "--wait")
348+
349+
if res.ExitCode != 1 {
350+
t.Fatalf("ExitCode = %d, want 1\nstderr: %s", res.ExitCode, res.Stderr)
351+
}
352+
if !strings.Contains(res.Stderr, "operation reached terminal failure state: failed") {
353+
t.Fatalf("stderr = %s, want failure message", res.Stderr)
354+
}
355+
}
356+
357+
func TestGenBuilder_VideoCreate_Wait_Timeout(t *testing.T) {
358+
srv := setupTestServer(t, map[string]testHandler{
359+
"POST /v3/videos": {
360+
StatusCode: 200,
361+
Body: `{"data":{"video_id":"vid_123"}}`,
362+
},
363+
"GET /v3/videos/vid_123": {
364+
StatusCode: 200,
365+
Body: `{"data":{"video_id":"vid_123","status":"processing"}}`,
366+
},
367+
})
368+
defer srv.Close()
369+
370+
res := runGenCommand(t, srv.URL, "test-key", videoCreateWaitSpec, "create", "--wait", "--timeout", "20ms")
371+
372+
if res.ExitCode != 1 {
373+
t.Fatalf("ExitCode = %d, want 1\nstderr: %s", res.ExitCode, res.Stderr)
374+
}
375+
if !strings.Contains(res.Stderr, "polling timed out before the operation completed") {
376+
t.Fatalf("stderr = %s, want timeout message", res.Stderr)
377+
}
378+
}
379+
380+
func TestGenBuilder_VideoCreate_NoWait(t *testing.T) {
381+
var statusCalled bool
382+
srv := setupTestServer(t, map[string]testHandler{
383+
"POST /v3/videos": {
384+
StatusCode: 200,
385+
Body: `{"data":{"video_id":"vid_123","status":"pending"}}`,
386+
},
387+
"GET /v3/videos/vid_123": {
388+
StatusCode: 200,
389+
ValidateRequest: func(t *testing.T, r *http.Request) {
390+
t.Helper()
391+
statusCalled = true
392+
},
393+
Body: `{"data":{"video_id":"vid_123","status":"completed"}}`,
394+
},
395+
})
396+
defer srv.Close()
397+
398+
res := runGenCommand(t, srv.URL, "test-key", videoCreateWaitSpec, "create")
399+
400+
if res.ExitCode != 0 {
401+
t.Fatalf("ExitCode = %d, want 0\nstderr: %s", res.ExitCode, res.Stderr)
402+
}
403+
if statusCalled {
404+
t.Fatal("status endpoint should not be called without --wait")
405+
}
406+
var parsed map[string]any
407+
if err := json.Unmarshal([]byte(res.Stdout), &parsed); err != nil {
408+
t.Fatalf("stdout is not valid JSON: %v\nstdout: %s", err, res.Stdout)
409+
}
410+
data := parsed["data"].(map[string]any)
411+
if data["status"] != "pending" {
412+
t.Fatalf("status = %v, want pending", data["status"])
413+
}
414+
}
415+
416+
func TestGenBuilder_VideoCreate_WaitNotAvailable(t *testing.T) {
417+
nonPollable := &command.Spec{
418+
Group: "video",
419+
Name: "get",
420+
Summary: "Get video",
421+
Endpoint: "/v3/videos/{video_id}",
422+
Method: "GET",
423+
Args: []command.ArgSpec{
424+
{Name: "video-id", Param: "video_id"},
425+
},
426+
Examples: []string{"heygen video get <video-id>"},
427+
}
428+
429+
srv := setupTestServer(t, map[string]testHandler{})
430+
defer srv.Close()
431+
432+
res := runGenCommand(t, srv.URL, "test-key", nonPollable, "get", "vid_123", "--wait")
433+
434+
if res.ExitCode != 2 {
435+
t.Fatalf("ExitCode = %d, want 2\nstderr: %s", res.ExitCode, res.Stderr)
436+
}
437+
if !strings.Contains(res.Stderr, "unknown flag: --wait") {
438+
t.Fatalf("stderr = %s, want unknown flag error", res.Stderr)
439+
}
440+
}
441+
265442
func TestGenBuilder_PostWithBodyFlags(t *testing.T) {
266443
var gotBody map[string]any
267444

@@ -579,6 +756,9 @@ func runGeneratedRootCommand(t *testing.T, serverURL, apiKey string, groups map[
579756

580757
t.Setenv("HEYGEN_API_KEY", apiKey)
581758
t.Setenv("HEYGEN_API_BASE", serverURL)
759+
if _, ok := os.LookupEnv("HEYGEN_CONFIG_DIR"); !ok {
760+
t.Setenv("HEYGEN_CONFIG_DIR", t.TempDir())
761+
}
582762

583763
root := newRootCmdWithSpecs("test", formatter, groups)
584764
root.SetOut(&stdout)

‎cmd/heygen/poll_configs.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package main
2+
3+
import "github.com/heygen-com/heygen-cli/internal/command"
4+
5+
// pollConfigs maps "group/spec.Name" to PollConfig for async commands.
6+
// Keys use the full spec name to avoid collisions within a group.
7+
var pollConfigs = map[string]*command.PollConfig{
8+
"video/create": {
9+
StatusEndpoint: "/v3/videos/{video_id}",
10+
StatusField: "data.status",
11+
TerminalOK: []string{"completed"},
12+
TerminalFail: []string{"failed", "error"},
13+
IDField: "data.video_id",
14+
},
15+
"video-translate/create": {
16+
StatusEndpoint: "/v3/video-translations/{video_translation_id}",
17+
StatusField: "data.status",
18+
TerminalOK: []string{"completed"},
19+
TerminalFail: []string{"failed", "error"},
20+
IDField: "data.video_translate_id",
21+
},
22+
"video-agent/create": {
23+
StatusEndpoint: "/v3/videos/{video_id}",
24+
StatusField: "data.status",
25+
TerminalOK: []string{"completed"},
26+
TerminalFail: []string{"failed", "error"},
27+
IDField: "data.video_id",
28+
},
29+
}

‎cmd/heygen/poll_configs_test.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package main
2+
3+
import (
4+
"testing"
5+
6+
"github.com/heygen-com/heygen-cli/gen"
7+
)
8+
9+
func TestPollConfigs_AllReferencedCommandsExist(t *testing.T) {
10+
for key := range pollConfigs {
11+
found := false
12+
for group, specs := range gen.Groups {
13+
for _, spec := range specs {
14+
if key == group+"/"+spec.Name {
15+
found = true
16+
break
17+
}
18+
}
19+
if found {
20+
break
21+
}
22+
}
23+
24+
if !found {
25+
t.Errorf("poll config key %q does not match any generated spec", key)
26+
}
27+
}
28+
}

‎cmd/heygen/root.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313

1414
func newRootCmd(version string, formatter output.Formatter) *cobra.Command {
1515
ctx := &cmdContext{formatter: formatter}
16+
applyPollConfigs(gen.Groups)
1617

1718
root := &cobra.Command{
1819
Use: "heygen",
@@ -54,6 +55,7 @@ Environment Variables:
5455
// verify the generic builder produces correct behavior.
5556
func newRootCmdWithSpecs(version string, formatter output.Formatter, groups map[string][]*command.Spec) *cobra.Command {
5657
ctx := &cmdContext{formatter: formatter}
58+
applyPollConfigs(groups)
5759

5860
root := &cobra.Command{
5961
Use: "heygen",
@@ -99,6 +101,17 @@ func registerGroups(root *cobra.Command, ctx *cmdContext, groups map[string][]*c
99101
}
100102
}
101103

104+
func applyPollConfigs(groups map[string][]*command.Spec) {
105+
for group, specs := range groups {
106+
for _, spec := range specs {
107+
key := group + "/" + spec.Name
108+
if pc, ok := pollConfigs[key]; ok {
109+
spec.PollConfig = pc
110+
}
111+
}
112+
}
113+
}
114+
102115
func registerSpecCommand(groupCmd *cobra.Command, spec *command.Spec, ctx *cmdContext) {
103116
path := commandPathParts(spec)
104117
if len(path) == 0 {

0 commit comments

Comments
 (0)