zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-14 20:54:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: fix context (#18118)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-08-11 19:19:29 +08:00
committed by GitHub
parent 8bd5768ebc
commit c75edbfbe8
36 changed files with 492 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
internal/agent/component/streaming_test.go
View File

@@ -27,9 +27,10 @@ import (
func TestLLM_Stream_HappyPath(t *testing.T) {
stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "hello", Model: "echo"}}
withStubInvoker(t, stub)
ctx := t.Context()
c := NewLLMComponent(LLMParam{ModelID: "echo"})
ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
if err != nil {
t.Fatalf("Stream: %v", err)
}
@@ -61,9 +62,10 @@ func TestLLM_Stream_HappyPath(t *testing.T) {
func TestLLM_Stream_Error(t *testing.T) {
stub := &stubInvoker{err: context.DeadlineExceeded}
withStubInvoker(t, stub)
ctx := t.Context()
c := NewLLMComponent(LLMParam{ModelID: "echo"})
ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
if err != nil {
t.Fatalf("Stream: %v", err)
}
@@ -86,7 +88,7 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) {
withStubInvoker(t, stub)
c := NewLLMComponent(LLMParam{ModelID: "echo"})
ctx, cancel := context.WithCancel(context.Background())
ctx, cancel := context.WithCancel(t.Context())
cancel() // pre-cancel
ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
@@ -111,9 +113,10 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) {
func TestLLM_Stream_BufferDoesNotBlock(t *testing.T) {
stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
withStubInvoker(t, stub)
ctx := t.Context()
c := NewLLMComponent(LLMParam{ModelID: "echo"})
ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
if err != nil {
t.Fatalf("Stream: %v", err)
}