70 lines
2.3 KiB
Go
70 lines
2.3 KiB
Go
|
package dify
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestStreamingPublishesAtNodeAndStopsInterruptedRun(t *testing.T) {
|
||
|
|
for _, interrupt := range []bool{false, true} {
|
||
|
|
t.Run(fmt.Sprint(interrupt), func(t *testing.T) {
|
||
|
|
var stops atomic.Int32
|
||
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
if r.Header.Get("Authorization") != "Bearer private-key" {
|
||
|
|
t.Error("missing authorization")
|
||
|
|
}
|
||
|
|
if r.URL.Path == "/workflows/tasks/task-1/stop" {
|
||
|
|
stops.Add(1)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
||
|
|
fmt.Fprintln(w, `data: {"event":"workflow_started","task_id":"task-1","workflow_run_id":"run-1"}`)
|
||
|
|
fmt.Fprintln(w, `data: {"event":"node_finished","workflow_run_id":"run-1","data":{"node_id":"validate-book","status":"succeeded"}}`)
|
||
|
|
fmt.Fprintln(w, `data: {"event":"workflow_finished","data":{"id":"run-1","status":"succeeded","outputs":{"chapter_count":7}}}`)
|
||
|
|
}))
|
||
|
|
defer srv.Close()
|
||
|
|
c, _ := NewClient(srv.URL, "private-key")
|
||
|
|
published := false
|
||
|
|
sentinel := errors.New("authorization revoked")
|
||
|
|
run, err := c.Run(context.Background(), map[string]any{"context_id": "opaque"}, func(e Event) error {
|
||
|
|
published = e.Data.NodeID == "validate-book"
|
||
|
|
if interrupt {
|
||
|
|
return sentinel
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
})
|
||
|
|
if !published || run.ID != "run-1" {
|
||
|
|
t.Fatalf("missing early publication or run ID: %+v", run)
|
||
|
|
}
|
||
|
|
if interrupt {
|
||
|
|
if !errors.Is(err, sentinel) || stops.Load() != 1 {
|
||
|
|
t.Fatalf("err=%v stops=%d", err, stops.Load())
|
||
|
|
}
|
||
|
|
} else if err != nil || stops.Load() != 0 || string(run.Outputs["chapter_count"]) != "7" {
|
||
|
|
t.Fatalf("run=%+v err=%v", run, err)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTruncatedStreamCannotReportSuccess(t *testing.T) {
|
||
|
|
var stops atomic.Int32
|
||
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
if r.URL.Path == "/workflows/tasks/task-2/stop" {
|
||
|
|
stops.Add(1)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
fmt.Fprintln(w, `data: {"event":"workflow_started","task_id":"task-2","workflow_run_id":"run-2"}`)
|
||
|
|
}))
|
||
|
|
defer srv.Close()
|
||
|
|
c, _ := NewClient(srv.URL, "private-key")
|
||
|
|
if _, err := c.Run(context.Background(), nil, nil); !errors.Is(err, ErrUnavailable) || stops.Load() != 1 {
|
||
|
|
t.Fatalf("err=%v stops=%d", err, stops.Load())
|
||
|
|
}
|
||
|
|
}
|