Skip to content

Commit baa8261

Browse files
chenxingtong-bytedancechenxingtong-bytedance
committed
(im) support im oapi range download large file
Change-Id: I38e6f6f9cf8b8711dc40650d19c77503f4e44989
1 parent 30dba35 commit baa8261

6 files changed

Lines changed: 573 additions & 54 deletions

File tree

‎shortcuts/common/runner.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -283,7 +283,13 @@ func (ctx *RuntimeContext) DoAPIStream(callCtx context.Context, req *larkcore.Ap
283283
option.Header = make(http.Header)
284284
}
285285
if shortcutHeaders := cmdutil.ShortcutHeaderOpts(ctx.ctx); shortcutHeaders != nil {
286-
shortcutHeaders(&option)
286+
var shortcutOption larkcore.RequestOption
287+
shortcutHeaders(&shortcutOption)
288+
for key, values := range shortcutOption.Header {
289+
for _, value := range values {
290+
option.Header.Add(key, value)
291+
}
292+
}
287293
}
288294

289295
accessToken, err := ctx.AccessToken()

‎shortcuts/im/helpers_network_test.go‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,15 @@ package im
66
import (
77
"bytes"
88
"context"
9+
"crypto/md5"
910
"encoding/json"
1011
"fmt"
1112
"io"
1213
"net/http"
1314
"os"
1415
"path/filepath"
1516
"reflect"
17+
"strconv"
1618
"strings"
1719
"testing"
1820
"unsafe"
@@ -289,6 +291,9 @@ func TestDownloadIMResourceToPathSuccess(t *testing.T) {
289291
if gotHeaders.Get(cmdutil.HeaderExecutionId) != "exec-123" {
290292
t.Fatalf("%s = %q, want %q", cmdutil.HeaderExecutionId, gotHeaders.Get(cmdutil.HeaderExecutionId), "exec-123")
291293
}
294+
if gotHeaders.Get("Range") != fmt.Sprintf("bytes=0-%d", probeChunkSize-1) {
295+
t.Fatalf("Range header = %q, want %q", gotHeaders.Get("Range"), fmt.Sprintf("bytes=0-%d", probeChunkSize-1))
296+
}
292297
}
293298

294299
func TestDownloadIMResourceToPathHTTPErrorBody(t *testing.T) {
@@ -313,6 +318,252 @@ func TestDownloadIMResourceToPathHTTPErrorBody(t *testing.T) {
313318
}
314319
}
315320

321+
func TestDownloadIMResourceToPathRetriesNetworkError(t *testing.T) {
322+
attempts := 0
323+
payload := []byte("retry success")
324+
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {
325+
switch {
326+
case strings.Contains(req.URL.Path, "tenant_access_token"):
327+
return shortcutJSONResponse(200, map[string]interface{}{
328+
"code": 0,
329+
"tenant_access_token": "tenant-token",
330+
"expire": 7200,
331+
}), nil
332+
case strings.Contains(req.URL.Path, "/open-apis/im/v1/messages/om_retry/resources/file_retry"):
333+
attempts++
334+
if attempts < 3 {
335+
return nil, fmt.Errorf("temporary network failure")
336+
}
337+
return shortcutRawResponse(200, payload, http.Header{"Content-Type": []string{"application/octet-stream"}}), nil
338+
default:
339+
return nil, fmt.Errorf("unexpected request: %s", req.URL.String())
340+
}
341+
}))
342+
343+
target := filepath.Join(t.TempDir(), "out.bin")
344+
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry", "file_retry", "file", target)
345+
if err != nil {
346+
t.Fatalf("downloadIMResourceToPath() error = %v", err)
347+
}
348+
if attempts != 3 {
349+
t.Fatalf("download attempts = %d, want 3", attempts)
350+
}
351+
if size != int64(len(payload)) {
352+
t.Fatalf("downloadIMResourceToPath() size = %d, want %d", size, len(payload))
353+
}
354+
}
355+
356+
func TestDownloadIMResourceToPathRetrySecondAttemptSuccess(t *testing.T) {
357+
attempts := 0
358+
payload := []byte("second retry success")
359+
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {
360+
switch {
361+
case strings.Contains(req.URL.Path, "tenant_access_token"):
362+
return shortcutJSONResponse(200, map[string]interface{}{
363+
"code": 0,
364+
"tenant_access_token": "tenant-token",
365+
"expire": 7200,
366+
}), nil
367+
case strings.Contains(req.URL.Path, "/open-apis/im/v1/messages/om_retry2/resources/file_retry2"):
368+
attempts++
369+
if attempts < 2 {
370+
return nil, fmt.Errorf("temporary network failure")
371+
}
372+
return shortcutRawResponse(200, payload, http.Header{"Content-Type": []string{"application/octet-stream"}}), nil
373+
default:
374+
return nil, fmt.Errorf("unexpected request: %s", req.URL.String())
375+
}
376+
}))
377+
378+
target := filepath.Join(t.TempDir(), "out.bin")
379+
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry2", "file_retry2", "file", target)
380+
if err != nil {
381+
t.Fatalf("downloadIMResourceToPath() error = %v", err)
382+
}
383+
if attempts != 2 {
384+
t.Fatalf("download attempts = %d, want 2", attempts)
385+
}
386+
if size != int64(len(payload)) {
387+
t.Fatalf("downloadIMResourceToPath() size = %d, want %d", size, len(payload))
388+
}
389+
}
390+
391+
func TestDownloadIMResourceToPathRetryContextCanceled(t *testing.T) {
392+
attempts := 0
393+
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {
394+
switch {
395+
case strings.Contains(req.URL.Path, "tenant_access_token"):
396+
return shortcutJSONResponse(200, map[string]interface{}{
397+
"code": 0,
398+
"tenant_access_token": "tenant-token",
399+
"expire": 7200,
400+
}), nil
401+
case strings.Contains(req.URL.Path, "/open-apis/im/v1/messages/om_cancel/resources/file_cancel"):
402+
attempts++
403+
return nil, fmt.Errorf("temporary network failure")
404+
default:
405+
return nil, fmt.Errorf("unexpected request: %s", req.URL.String())
406+
}
407+
}))
408+
409+
ctx, cancel := context.WithCancel(context.Background())
410+
// Cancel context immediately to trigger context error on first retry
411+
cancel()
412+
413+
target := filepath.Join(t.TempDir(), "out.bin")
414+
_, _, err := downloadIMResourceToPath(ctx, runtime, "om_cancel", "file_cancel", "file", target)
415+
if err != context.Canceled {
416+
t.Fatalf("downloadIMResourceToPath() error = %v, want context.Canceled", err)
417+
}
418+
// First attempt is made, then retry checks ctx.Err() and returns
419+
if attempts != 1 {
420+
t.Fatalf("download attempts = %d, want 1", attempts)
421+
}
422+
}
423+
424+
func TestDownloadIMResourceToPathRangeDownload(t *testing.T) {
425+
cases := []struct {
426+
name string
427+
payloadLen int64
428+
wantRanges []string
429+
}{
430+
{
431+
name: "single small chunk",
432+
payloadLen: 16,
433+
wantRanges: []string{"bytes=0-131071"},
434+
},
435+
{
436+
name: "exact probe chunk",
437+
payloadLen: probeChunkSize,
438+
wantRanges: []string{"bytes=0-131071"},
439+
},
440+
{
441+
name: "multiple chunks with tail",
442+
payloadLen: probeChunkSize + normalChunkSize + 1234,
443+
wantRanges: []string{
444+
"bytes=0-131071",
445+
fmt.Sprintf("bytes=%d-%d", probeChunkSize, probeChunkSize+normalChunkSize-1),
446+
fmt.Sprintf("bytes=%d-%d", probeChunkSize+normalChunkSize, probeChunkSize+normalChunkSize+1233),
447+
},
448+
},
449+
{
450+
name: "multiple chunks exact 8mb tail",
451+
payloadLen: probeChunkSize + 2*normalChunkSize,
452+
wantRanges: []string{
453+
"bytes=0-131071",
454+
fmt.Sprintf("bytes=%d-%d", probeChunkSize, probeChunkSize+normalChunkSize-1),
455+
fmt.Sprintf("bytes=%d-%d", probeChunkSize+normalChunkSize, probeChunkSize+2*normalChunkSize-1),
456+
},
457+
},
458+
}
459+
460+
for _, tt := range cases {
461+
t.Run(tt.name, func(t *testing.T) {
462+
payload := bytes.Repeat([]byte("range-download-"), int(tt.payloadLen/15)+1)
463+
payload = payload[:tt.payloadLen]
464+
465+
var gotRanges []string
466+
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {
467+
switch {
468+
case strings.Contains(req.URL.Path, "tenant_access_token"):
469+
return shortcutJSONResponse(200, map[string]interface{}{
470+
"code": 0,
471+
"tenant_access_token": "tenant-token",
472+
"expire": 7200,
473+
}), nil
474+
case strings.Contains(req.URL.Path, "/open-apis/im/v1/messages/om_range/resources/file_range"):
475+
rangeHeader := req.Header.Get("Range")
476+
gotRanges = append(gotRanges, rangeHeader)
477+
if req.Header.Get("Authorization") != "Bearer tenant-token" {
478+
return nil, fmt.Errorf("missing authorization header")
479+
}
480+
start, end, err := parseRangeHeader(rangeHeader, int64(len(payload)))
481+
if err != nil {
482+
return nil, err
483+
}
484+
return shortcutRawResponse(http.StatusPartialContent, payload[start:end+1], http.Header{
485+
"Content-Type": []string{"application/octet-stream"},
486+
"Content-Range": []string{fmt.Sprintf("bytes %d-%d/%d", start, end, len(payload))},
487+
}), nil
488+
default:
489+
return nil, fmt.Errorf("unexpected request: %s", req.URL.String())
490+
}
491+
}))
492+
493+
target := filepath.Join(t.TempDir(), "nested", "resource.bin")
494+
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_range", "file_range", "file", target)
495+
if err != nil {
496+
t.Fatalf("downloadIMResourceToPath() error = %v", err)
497+
}
498+
if size != int64(len(payload)) {
499+
t.Fatalf("downloadIMResourceToPath() size = %d, want %d", size, len(payload))
500+
}
501+
if !reflect.DeepEqual(gotRanges, tt.wantRanges) {
502+
t.Fatalf("Range requests = %#v, want %#v", gotRanges, tt.wantRanges)
503+
}
504+
505+
got, err := os.ReadFile(target)
506+
if err != nil {
507+
t.Fatalf("ReadFile() error = %v", err)
508+
}
509+
if md5.Sum(got) != md5.Sum(payload) {
510+
t.Fatalf("downloaded payload MD5 = %x, want %x", md5.Sum(got), md5.Sum(payload))
511+
}
512+
})
513+
}
514+
}
515+
516+
func TestDownloadIMResourceToPathInvalidContentRange(t *testing.T) {
517+
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {
518+
switch {
519+
case strings.Contains(req.URL.Path, "tenant_access_token"):
520+
return shortcutJSONResponse(200, map[string]interface{}{
521+
"code": 0,
522+
"tenant_access_token": "tenant-token",
523+
"expire": 7200,
524+
}), nil
525+
case strings.Contains(req.URL.Path, "/open-apis/im/v1/messages/om_bad/resources/file_bad"):
526+
return shortcutRawResponse(http.StatusPartialContent, []byte("bad"), http.Header{
527+
"Content-Type": []string{"application/octet-stream"},
528+
"Content-Range": []string{"bytes 0-2/not-a-number"},
529+
}), nil
530+
default:
531+
return nil, fmt.Errorf("unexpected request: %s", req.URL.String())
532+
}
533+
}))
534+
535+
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_bad", "file_bad", "file", filepath.Join(t.TempDir(), "out.bin"))
536+
if err == nil || !strings.Contains(err.Error(), "invalid Content-Range header") {
537+
t.Fatalf("downloadIMResourceToPath() error = %v", err)
538+
}
539+
}
540+
541+
func parseRangeHeader(header string, totalSize int64) (int64, int64, error) {
542+
if !strings.HasPrefix(header, "bytes=") {
543+
return 0, 0, fmt.Errorf("unexpected range header: %q", header)
544+
}
545+
parts := strings.SplitN(strings.TrimPrefix(header, "bytes="), "-", 2)
546+
if len(parts) != 2 {
547+
return 0, 0, fmt.Errorf("unexpected range header: %q", header)
548+
}
549+
550+
start, err := strconv.ParseInt(parts[0], 10, 64)
551+
if err != nil {
552+
return 0, 0, fmt.Errorf("parse start: %w", err)
553+
}
554+
end, err := strconv.ParseInt(parts[1], 10, 64)
555+
if err != nil {
556+
return 0, 0, fmt.Errorf("parse end: %w", err)
557+
}
558+
if start < 0 || end < start || start >= totalSize {
559+
return 0, 0, fmt.Errorf("invalid range bounds: %d-%d for size %d", start, end, totalSize)
560+
}
561+
if end >= totalSize {
562+
end = totalSize - 1
563+
}
564+
return start, end, nil
565+
}
566+
316567
func TestUploadImageToIMSuccess(t *testing.T) {
317568
var gotBody string
318569
runtime := newBotShortcutRuntime(t, shortcutRoundTripFunc(func(req *http.Request) (*http.Response, error) {

‎shortcuts/im/helpers_test.go‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -489,6 +489,43 @@ func TestDownloadIMResourceToPathHTTPClientError(t *testing.T) {
489489
}
490490
}
491491

492+
func TestParseTotalSize(t *testing.T) {
493+
tests := []struct {
494+
name string
495+
contentRange string
496+
want int64
497+
wantErr string
498+
}{
499+
{name: "normal", contentRange: "bytes 0-131071/104857600", want: 104857600},
500+
{name: "single probe chunk", contentRange: "bytes 0-131071/131072", want: 131072},
501+
{name: "single small chunk", contentRange: "bytes 0-15/16", want: 16},
502+
{name: "empty", contentRange: "", wantErr: "content-range is empty"},
503+
{name: "invalid prefix", contentRange: "items 0-15/16", wantErr: `unsupported content-range: "items 0-15/16"`},
504+
{name: "missing total", contentRange: "bytes 0-15/", wantErr: `unsupported content-range: "bytes 0-15/"`},
505+
{name: "wildcard", contentRange: "bytes */16", wantErr: `unsupported content-range: "bytes */16"`},
506+
{name: "unknown total size", contentRange: "bytes 0-99/*", wantErr: `unknown total size in content-range: "bytes 0-99/*"`},
507+
{name: "invalid total", contentRange: "bytes 0-15/not-a-number", wantErr: "parse total size:"},
508+
}
509+
510+
for _, tt := range tests {
511+
t.Run(tt.name, func(t *testing.T) {
512+
got, err := parseTotalSize(tt.contentRange)
513+
if tt.wantErr != "" {
514+
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
515+
t.Fatalf("parseTotalSize() error = %v, want substring %q", err, tt.wantErr)
516+
}
517+
return
518+
}
519+
if err != nil {
520+
t.Fatalf("parseTotalSize() unexpected error = %v", err)
521+
}
522+
if got != tt.want {
523+
t.Fatalf("parseTotalSize() = %d, want %d", got, tt.want)
524+
}
525+
})
526+
}
527+
}
528+
492529
func TestShortcuts(t *testing.T) {
493530
var commands []string
494531
for _, shortcut := range Shortcuts() {

0 commit comments

Comments
 (0)