Skip to content

Commit d523cd4

Browse files
somanshreddysomanshreddy
committed
cmd: support captioned video downloads
1 parent 2e4fd74 commit d523cd4

2 files changed

Lines changed: 146 additions & 17 deletions

File tree

‎cmd/heygen/video_download.go‎

Lines changed: 60 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,17 +18,35 @@ import (
1818

1919
var downloadClient = &http.Client{Timeout: 10 * time.Minute}
2020

21+
type assetInfo struct {
22+
field string
23+
ext string
24+
label string
25+
}
26+
27+
var assetTypes = map[string]assetInfo{
28+
"video": {field: "video_url", ext: ".mp4", label: "video"},
29+
"captioned": {field: "captioned_video_url", ext: ".mp4", label: "captioned video"},
30+
}
31+
2132
func newVideoDownloadCmd(ctx *cmdContext) *cobra.Command {
2233
var outputPath string
34+
var asset string
2335

2436
cmd := &cobra.Command{
2537
Use: "download <video-id>",
26-
Short: "Download a video file to disk",
38+
Short: "Download a video file or related asset to disk",
2739
Args: cobra.ExactArgs(1),
2840
Example: "heygen video download <video-id>\n" +
41+
"heygen video download <video-id> --asset captioned\n" +
2942
"heygen video download <video-id> --output-path my-video.mp4",
3043
RunE: func(cmd *cobra.Command, args []string) error {
3144
videoID := args[0]
45+
info, ok := assetTypes[asset]
46+
if !ok {
47+
return clierrors.NewUsage(
48+
fmt.Sprintf("invalid --asset value %q: must be one of: video, captioned", asset))
49+
}
3250

3351
spec := &command.Spec{
3452
Endpoint: "/v3/videos/{video_id}",
@@ -43,7 +61,7 @@ func newVideoDownloadCmd(ctx *cmdContext) *cobra.Command {
4361
return err
4462
}
4563

46-
videoURL, err := extractVideoURL(result, videoID)
64+
assetURL, err := extractAssetURL(result, videoID, info)
4765
if err != nil {
4866
return err
4967
}
@@ -52,15 +70,16 @@ func newVideoDownloadCmd(ctx *cmdContext) *cobra.Command {
5270
if dest == "" {
5371
// Sanitize: strip directory components from video ID to prevent
5472
// path traversal. Handles IDs with / or \ safely.
55-
dest = filepath.Base(videoID) + ".mp4"
73+
dest = filepath.Base(videoID) + info.ext
5674
}
5775

58-
if err := downloadFile(cmd.Context(), videoURL, dest); err != nil {
76+
if err := downloadFile(cmd.Context(), assetURL, dest); err != nil {
5977
return err
6078
}
6179

6280
data, err := json.Marshal(map[string]string{
63-
"message": "Downloaded to " + dest,
81+
"asset": asset,
82+
"message": fmt.Sprintf("Downloaded %s to %s", info.label, dest),
6483
"path": dest,
6584
})
6685
if err != nil {
@@ -71,34 +90,49 @@ func newVideoDownloadCmd(ctx *cmdContext) *cobra.Command {
7190
},
7291
}
7392

93+
cmd.Flags().StringVar(&asset, "asset", "video", "Asset to download: video, captioned")
7494
cmd.Flags().StringVar(&outputPath, "output-path", "", "Output file path (default: {video-id}.mp4)")
7595
return cmd
7696
}
7797

78-
func extractVideoURL(raw json.RawMessage, videoID string) (string, error) {
98+
func extractAssetURL(raw json.RawMessage, videoID string, info assetInfo) (string, error) {
7999
var resp struct {
80-
Data struct {
81-
VideoURL string `json:"video_url"`
82-
Status string `json:"status"`
83-
} `json:"data"`
100+
Data map[string]json.RawMessage `json:"data"`
84101
}
85102
if err := json.Unmarshal(raw, &resp); err != nil {
86103
return "", clierrors.New("failed to parse video response")
87104
}
88105

89-
if resp.Data.VideoURL == "" {
90-
switch resp.Data.Status {
106+
var status string
107+
if rawStatus, ok := resp.Data["status"]; ok {
108+
_ = json.Unmarshal(rawStatus, &status)
109+
}
110+
111+
var assetURL string
112+
if rawURL, ok := resp.Data[info.field]; ok {
113+
_ = json.Unmarshal(rawURL, &assetURL)
114+
}
115+
116+
if assetURL == "" {
117+
switch status {
91118
case "failed", "error":
92119
return "", &clierrors.CLIError{
93120
Code: "video_failed",
94-
Message: fmt.Sprintf("video rendering failed (status: %s)", resp.Data.Status),
121+
Message: fmt.Sprintf("video rendering failed (status: %s)", status),
95122
Hint: "Check details with: heygen video get " + videoID,
96123
ExitCode: clierrors.ExitGeneral,
97124
}
125+
case "completed":
126+
return "", &clierrors.CLIError{
127+
Code: "asset_not_available",
128+
Message: fmt.Sprintf("%s not available for this video", info.label),
129+
Hint: assetHint(info.field),
130+
ExitCode: clierrors.ExitGeneral,
131+
}
98132
default:
99-
msg := "video URL not available"
100-
if resp.Data.Status != "" {
101-
msg = fmt.Sprintf("video URL not available (status: %s)", resp.Data.Status)
133+
msg := fmt.Sprintf("%s URL not available", info.label)
134+
if status != "" {
135+
msg = fmt.Sprintf("%s URL not available (status: %s)", info.label, status)
102136
}
103137
return "", &clierrors.CLIError{
104138
Code: "video_not_ready",
@@ -109,7 +143,16 @@ func extractVideoURL(raw json.RawMessage, videoID string) (string, error) {
109143
}
110144
}
111145

112-
return resp.Data.VideoURL, nil
146+
return assetURL, nil
147+
}
148+
149+
func assetHint(field string) string {
150+
switch field {
151+
case "captioned_video_url":
152+
return "Video may not have been created with captions enabled."
153+
default:
154+
return ""
155+
}
113156
}
114157

115158
func downloadFile(ctx context.Context, videoURL, dest string) error {

‎cmd/heygen/video_download_test.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ func TestVideoDownload_Success(t *testing.T) {
5454
if payload["path"] != dest {
5555
t.Fatalf("path = %q, want %q", payload["path"], dest)
5656
}
57+
if payload["asset"] != "video" {
58+
t.Fatalf("asset = %q, want %q", payload["asset"], "video")
59+
}
5760
}
5861

5962
func TestVideoDownload_DefaultFilename(t *testing.T) {
@@ -131,6 +134,89 @@ func TestVideoDownload_WithOutputPath(t *testing.T) {
131134
}
132135
}
133136

137+
func TestVideoDownload_AssetCaptioned(t *testing.T) {
138+
tmpDir := t.TempDir()
139+
dest := filepath.Join(tmpDir, "captioned.mp4")
140+
141+
var srv *httptest.Server
142+
srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
143+
switch {
144+
case r.Method == http.MethodGet && r.URL.Path == "/v3/videos/vid_123":
145+
writeJSON(t, w, map[string]any{
146+
"data": map[string]any{
147+
"captioned_video_url": srv.URL + "/download/captioned.mp4",
148+
"status": "completed",
149+
},
150+
})
151+
case r.Method == http.MethodGet && r.URL.Path == "/download/captioned.mp4":
152+
_, _ = w.Write([]byte("captioned-bytes"))
153+
default:
154+
t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path)
155+
}
156+
}))
157+
defer srv.Close()
158+
159+
res := runCommand(t, srv.URL, "test-key", "video", "download", "vid_123", "--asset", "captioned", "--output-path", dest)
160+
if res.ExitCode != 0 {
161+
t.Fatalf("ExitCode = %d, want 0\nstderr: %s", res.ExitCode, res.Stderr)
162+
}
163+
164+
data, err := os.ReadFile(dest)
165+
if err != nil {
166+
t.Fatalf("ReadFile: %v", err)
167+
}
168+
if string(data) != "captioned-bytes" {
169+
t.Fatalf("file contents = %q, want %q", string(data), "captioned-bytes")
170+
}
171+
172+
var payload map[string]string
173+
if err := json.Unmarshal([]byte(res.Stdout), &payload); err != nil {
174+
t.Fatalf("stdout is not valid JSON: %v\nstdout: %s", err, res.Stdout)
175+
}
176+
if payload["asset"] != "captioned" {
177+
t.Fatalf("asset = %q, want %q", payload["asset"], "captioned")
178+
}
179+
}
180+
181+
func TestVideoDownload_AssetInvalid(t *testing.T) {
182+
srv := setupTestServer(t, map[string]testHandler{})
183+
defer srv.Close()
184+
185+
res := runCommand(t, srv.URL, "test-key", "video", "download", "vid_123", "--asset", "foo")
186+
if res.ExitCode != clierrors.ExitUsage {
187+
t.Fatalf("ExitCode = %d, want %d\nstderr: %s", res.ExitCode, clierrors.ExitUsage, res.Stderr)
188+
}
189+
if !strings.Contains(res.Stderr, "must be one of: video, captioned") {
190+
t.Fatalf("stderr = %q, want valid asset list", res.Stderr)
191+
}
192+
}
193+
194+
func TestVideoDownload_AssetNotAvailable(t *testing.T) {
195+
srv := setupTestServer(t, map[string]testHandler{
196+
"GET /v3/videos/vid_123": {
197+
StatusCode: http.StatusOK,
198+
Body: `{"data":{"status":"completed","captioned_video_url":""}}`,
199+
},
200+
})
201+
defer srv.Close()
202+
203+
res := runCommand(t, srv.URL, "test-key", "video", "download", "vid_123", "--asset", "captioned")
204+
if res.ExitCode != clierrors.ExitGeneral {
205+
t.Fatalf("ExitCode = %d, want %d\nstderr: %s", res.ExitCode, clierrors.ExitGeneral, res.Stderr)
206+
}
207+
208+
var envelope map[string]map[string]any
209+
if err := json.Unmarshal([]byte(res.Stderr), &envelope); err != nil {
210+
t.Fatalf("stderr is not valid JSON: %v\nstderr: %s", err, res.Stderr)
211+
}
212+
if envelope["error"]["code"] != "asset_not_available" {
213+
t.Fatalf("error.code = %v, want %q", envelope["error"]["code"], "asset_not_available")
214+
}
215+
if !strings.Contains(res.Stderr, "captions enabled") {
216+
t.Fatalf("stderr = %q, want captions hint", res.Stderr)
217+
}
218+
}
219+
134220
func TestVideoDownload_PreservesExistingFileOnFailure(t *testing.T) {
135221
tmpDir := t.TempDir()
136222
dest := filepath.Join(tmpDir, "existing.mp4")

0 commit comments

Comments
 (0)