zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-15 05:04:27 +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

37
internal/agent/sandbox/e2b_test.go
View File

@@ -79,7 +79,7 @@ func TestE2BProvider_Initialize_MissingCreds(t *testing.T) {
t.Setenv(k, "")
}
p := newE2BProviderFromEnv()
err := p.Initialize(context.Background())
err := p.Initialize(t.Context())
if err == nil {
t.Fatalf("Initialize with no creds: got nil error, want one")
}
@@ -98,7 +98,7 @@ func TestE2BProvider_Initialize_WithAPIKey(t *testing.T) {
t.Skip("E2B_API_KEY not set — skipping network-dependent init check")
}
p := newE2BProviderFromEnv()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
defer cancel()
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
@@ -117,18 +117,19 @@ func TestE2BProvider_AllOps_BeforeInit(t *testing.T) {
t.Parallel()
p := newE2BProviderFromEnv()
// Do NOT call Initialize.
ctx := t.Context()
inst := &SandboxInstance{InstanceID: "x", Provider: ProviderE2B}
if _, err := p.CreateInstance(context.Background(), "python"); err == nil {
if _, err := p.CreateInstance(ctx, "python"); err == nil {
t.Errorf("CreateInstance before init: got nil error, want one")
}
if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil {
if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil {
t.Errorf("ExecuteCode before init: got nil error, want one")
}
if err := p.DestroyInstance(context.Background(), inst); err == nil {
if err := p.DestroyInstance(ctx, inst); err == nil {
t.Errorf("DestroyInstance before init: got nil error, want one")
}
if err := p.HealthCheck(context.Background()); err == nil {
if err := p.HealthCheck(ctx); err == nil {
t.Errorf("HealthCheck before init: got nil error, want one")
}
}
@@ -140,6 +141,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
// — this lets us test the input-validation paths without
// hitting the e2b control plane.
p.initialized = true
ctx := t.Context()
cases := []struct {
name string
@@ -149,7 +151,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "empty instance id",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: ""}, "x", "python", 5, nil)
return err
},
@@ -158,7 +160,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "nil instance",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
nil, "x", "python", 5, nil)
return err
},
@@ -167,7 +169,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "unsupported language",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil)
return err
},
@@ -176,7 +178,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "timeout too small",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "python", 0, nil)
return err
},
@@ -185,7 +187,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "timeout too large",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "python", 1000, nil)
return err
},
@@ -209,7 +211,8 @@ func TestE2BProvider_CreateInstance_UnsupportedLanguage(t *testing.T) {
t.Parallel()
p := newE2BProviderFromEnv()
p.initialized = true
if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil {
ctx := t.Context()
if _, err := p.CreateInstance(ctx, "ruby"); err == nil {
t.Errorf("CreateInstance(ruby): got nil error, want one")
}
}
@@ -218,10 +221,11 @@ func TestE2BProvider_DestroyInstance_EmptyID(t *testing.T) {
t.Parallel()
p := newE2BProviderFromEnv()
p.initialized = true
if err := p.DestroyInstance(context.Background(), &SandboxInstance{InstanceID: ""}); err == nil {
ctx := t.Context()
if err := p.DestroyInstance(ctx, &SandboxInstance{InstanceID: ""}); err == nil {
t.Errorf("DestroyInstance(empty id): got nil error, want one")
}
if err := p.DestroyInstance(context.Background(), nil); err == nil {
if err := p.DestroyInstance(ctx, nil); err == nil {
t.Errorf("DestroyInstance(nil): got nil error, want one")
}
}
@@ -270,7 +274,7 @@ func TestE2BProvider_FullE2E_SkipWithoutKey(t *testing.T) {
t.Skip("E2B_API_KEY not set — skipping full E2E test (real network call)")
}
p := newE2BProviderFromEnv()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute)
defer cancel()
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
@@ -318,7 +322,8 @@ func TestE2BProvider_AccessTokenFallback(t *testing.T) {
// Initialize should NOT fail with "E2B_API_KEY or
// E2B_ACCESS_TOKEN is required". The error we'd see is the
// SDK's auth error, which is what we want.
err := p.Initialize(context.Background())
ctx := t.Context()
err := p.Initialize(ctx)
if err == nil {
t.Skip("Initialize succeeded — env-var fallback accepted; skipping further checks")
}

69
internal/agent/sandbox/local_test.go
View File

@@ -17,7 +17,6 @@
package sandbox
import (
"context"
"encoding/base64"
"os"
"path/filepath"
@@ -43,7 +42,8 @@ func newLocalForTest(t *testing.T) *LocalProvider {
maxArtifactBytes: 10 << 20,
instances: map[string]string{},
}
if err := p.Initialize(context.Background()); err != nil {
ctx := t.Context()
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
return p
@@ -108,7 +108,8 @@ func TestLocal_Initialize_CreatesWorkDir(t *testing.T) {
}
t.Setenv("LOCAL_WORK_DIR", workDir)
p := newLocalProviderFromEnv()
if err := p.Initialize(context.Background()); err != nil {
ctx := t.Context()
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
info, err := os.Stat(workDir)
@@ -122,7 +123,8 @@ func TestLocal_Initialize_CreatesWorkDir(t *testing.T) {
func TestLocal_CreateInstance_CreatesArtifactsDir(t *testing.T) {
p := newLocalForTest(t)
inst, err := p.CreateInstance(context.Background(), "python")
ctx := t.Context()
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
@@ -141,7 +143,8 @@ func TestLocal_CreateInstance_CreatesArtifactsDir(t *testing.T) {
func TestLocal_CreateInstance_RejectsBadLanguage(t *testing.T) {
p := newLocalForTest(t)
if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil {
ctx := t.Context()
if _, err := p.CreateInstance(ctx, "ruby"); err == nil {
t.Errorf("CreateInstance(ruby): got nil error, want one")
}
}
@@ -149,17 +152,18 @@ func TestLocal_CreateInstance_RejectsBadLanguage(t *testing.T) {
func TestLocal_AllOps_BeforeInit(t *testing.T) {
t.Parallel()
p := &LocalProvider{}
ctx := t.Context()
inst := &SandboxInstance{InstanceID: "x", Provider: ProviderLocal}
if _, err := p.CreateInstance(context.Background(), "python"); err == nil {
if _, err := p.CreateInstance(ctx, "python"); err == nil {
t.Errorf("CreateInstance before init: got nil error, want one")
}
if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil {
if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil {
t.Errorf("ExecuteCode before init: got nil error, want one")
}
if err := p.DestroyInstance(context.Background(), inst); err == nil {
if err := p.DestroyInstance(ctx, inst); err == nil {
t.Errorf("DestroyInstance before init: got nil error, want one")
}
if err := p.HealthCheck(context.Background()); err == nil {
if err := p.HealthCheck(ctx); err == nil {
t.Errorf("HealthCheck before init: got nil error, want one")
}
}
@@ -173,15 +177,16 @@ func TestLocal_ExecuteCode_Python_RoundTrip(t *testing.T) {
t.Skip("python3 not on PATH — skipping local subprocess test")
}
p := newLocalForTest(t)
ctx := t.Context()
p.pythonBin = pythonPath
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
defer p.DestroyInstance(context.Background(), inst)
defer p.DestroyInstance(ctx, inst)
code := "def main(): return {'value': 7, 'type': 'json'}"
result, err := p.ExecuteCode(context.Background(), inst, code, "python", 10, nil)
result, err := p.ExecuteCode(ctx, inst, code, "python", 10, nil)
if err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
@@ -198,7 +203,7 @@ func TestLocal_ExecuteCode_Python_RoundTrip(t *testing.T) {
func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) {
p := newLocalForTest(t)
p.initialized = true
ctx := t.Context()
cases := []struct {
name string
fn func() error
@@ -207,7 +212,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "empty instance id",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: ""}, "x", "python", 5, nil)
return err
},
@@ -216,7 +221,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "unsupported language",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil)
return err
},
@@ -225,7 +230,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "timeout too small",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "python", 0, nil)
return err
},
@@ -247,50 +252,53 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) {
func TestLocal_DestroyInstance_RemovesDir(t *testing.T) {
p := newLocalForTest(t)
inst, err := p.CreateInstance(context.Background(), "python")
ctx := t.Context()
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
dir := filepath.Join(p.workDir, inst.InstanceID)
if _, err := os.Stat(dir); err != nil {
if _, err = os.Stat(dir); err != nil {
t.Fatalf("instance dir not created: %v", err)
}
if err := p.DestroyInstance(context.Background(), inst); err != nil {
if err = p.DestroyInstance(ctx, inst); err != nil {
t.Errorf("DestroyInstance: %v", err)
}
if _, err := os.Stat(dir); !os.IsNotExist(err) {
if _, err = os.Stat(dir); !os.IsNotExist(err) {
t.Errorf("instance dir still exists after destroy: %v", err)
}
// Idempotent: second call should be a no-op.
if err := p.DestroyInstance(context.Background(), inst); err != nil {
if err = p.DestroyInstance(ctx, inst); err != nil {
t.Errorf("DestroyInstance (idempotent): %v", err)
}
}
func TestLocal_HealthCheck(t *testing.T) {
p := newLocalForTest(t)
if err := p.HealthCheck(context.Background()); err != nil {
ctx := t.Context()
if err := p.HealthCheck(ctx); err != nil {
t.Errorf("HealthCheck: %v", err)
}
// Removing the work dir should make HealthCheck fail.
// Removing the work dir should make health check fail.
if err := os.RemoveAll(p.workDir); err != nil {
t.Fatalf("remove work dir: %v", err)
}
if err := p.HealthCheck(context.Background()); err == nil {
if err := p.HealthCheck(ctx); err == nil {
t.Errorf("HealthCheck after remove: got nil error, want one")
}
}
func TestLocal_CollectArtifacts_RejectsBadExtension(t *testing.T) {
p := newLocalForTest(t)
inst, err := p.CreateInstance(context.Background(), "python")
ctx := t.Context()
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
defer p.DestroyInstance(context.Background(), inst)
defer p.DestroyInstance(ctx, inst)
// Drop an unsupported extension into the artifacts dir.
artDir := filepath.Join(p.workDir, inst.InstanceID, "artifacts")
if err := os.WriteFile(filepath.Join(artDir, "evil.exe"), []byte("x"), 0o600); err != nil {
if err = os.WriteFile(filepath.Join(artDir, "evil.exe"), []byte("x"), 0o600); err != nil {
t.Fatalf("write artifact: %v", err)
}
_, err = p.collectArtifacts(p.workDir + "/" + inst.InstanceID)
@@ -304,13 +312,14 @@ func TestLocal_CollectArtifacts_RejectsBadExtension(t *testing.T) {
func TestLocal_CollectArtifacts_AllowsCSVRoundTrip(t *testing.T) {
p := newLocalForTest(t)
inst, err := p.CreateInstance(context.Background(), "python")
ctx := t.Context()
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
defer p.DestroyInstance(context.Background(), inst)
defer p.DestroyInstance(ctx, inst)
artDir := filepath.Join(p.workDir, inst.InstanceID, "artifacts")
if err := os.WriteFile(filepath.Join(artDir, "out.csv"), []byte("a,b\n1,2\n"), 0o600); err != nil {
if err = os.WriteFile(filepath.Join(artDir, "out.csv"), []byte("a,b\n1,2\n"), 0o600); err != nil {
t.Fatalf("write artifact: %v", err)
}
artifacts, err := p.collectArtifacts(p.workDir + "/" + inst.InstanceID)

6
internal/agent/sandbox/manager_client_test.go
View File

@@ -37,7 +37,8 @@ func TestManagerClient_MapsStructuredResultToSandboxResponse(t *testing.T) {
mgr.SetProvider(managerClientStubProvider{})
client := &ManagerClient{manager: mgr}
resp, err := client.ExecuteCode(context.Background(), agenttool.SandboxRequest{
ctx := t.Context()
resp, err := client.ExecuteCode(ctx, agenttool.SandboxRequest{
Lang: "python",
Script: "def main(): return 16",
})
@@ -57,7 +58,8 @@ func TestManagerClient_MapsLegacyResultKeyToSandboxResponse(t *testing.T) {
mgr.SetProvider(managerClientResultKeyProvider{})
client := &ManagerClient{manager: mgr}
resp, err := client.ExecuteCode(context.Background(), agenttool.SandboxRequest{
ctx := t.Context()
resp, err := client.ExecuteCode(ctx, agenttool.SandboxRequest{
Lang: "python",
Script: "def main(): return 16",
})

34
internal/agent/sandbox/manager_test.go
View File

@@ -129,8 +129,9 @@ func TestAliyun_Initialize_MissingCreds(t *testing.T) {
for _, k := range []string{"AGENTRUN_ACCESS_KEY_ID", "AGENTRUN_ACCESS_KEY_SECRET", "AGENTRUN_ACCOUNT_ID"} {
t.Setenv(k, "")
}
ctx := t.Context()
p := newAliyunProviderFromEnv()
if err := p.Initialize(context.Background()); err == nil {
if err := p.Initialize(ctx); err == nil {
t.Errorf("Initialize with missing creds: got nil error, want one")
}
}
@@ -153,19 +154,19 @@ func TestSelfManaged_EndToEnd_FullLoop(t *testing.T) {
}
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
if err := p.Initialize(context.Background()); err != nil {
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
if inst.Provider != ProviderSelfManaged {
t.Errorf("provider = %q, want %q", inst.Provider, ProviderSelfManaged)
}
result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 5, nil)
result, err := p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 5, nil)
if err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
@@ -178,7 +179,7 @@ func TestSelfManaged_EndToEnd_FullLoop(t *testing.T) {
if result.ExitCode != 0 {
t.Errorf("exit_code = %d, want 0", result.ExitCode)
}
if err := p.DestroyInstance(context.Background(), inst); err != nil {
if err = p.DestroyInstance(ctx, inst); err != nil {
t.Errorf("DestroyInstance: %v", err)
}
}
@@ -425,6 +426,7 @@ func TestLoadFromSettingsWithReader_HappyPath(t *testing.T) {
w.WriteHeader(http.StatusNotFound)
}))
defer srv.Close()
ctx := t.Context()
// Drive the mock server by setting the SANDBOX_EXECUTOR_MANAGER_URL
// env var; then have the settings config return a matching
@@ -443,7 +445,7 @@ func TestLoadFromSettingsWithReader_HappyPath(t *testing.T) {
},
}
m := &ProviderManager{}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader: %v", err)
}
if !m.IsConfigured() {
@@ -485,10 +487,11 @@ func TestLoadFromSettingsWithReader_EmptyFallback(t *testing.T) {
t.Setenv("SANDBOX_PROVIDER_TYPE", "")
t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL)
t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s")
ctx := t.Context()
r := &fakeSettingsReader{rows: map[string][]entity.SystemSettings{}}
m := &ProviderManager{}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader: %v", err)
}
if !m.IsConfigured() {
@@ -515,10 +518,11 @@ func TestLoadFromSettingsWithReader_DAOErrorFallback(t *testing.T) {
t.Setenv("SANDBOX_PROVIDER_TYPE", "")
t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL)
t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s")
ctx := t.Context()
r := &fakeSettingsReader{fakeErr: errors.New("db is down")}
m := &ProviderManager{}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader (DAO error fallback): %v", err)
}
if got := m.Provider().ProviderType(); got != ProviderSelfManaged {
@@ -542,6 +546,7 @@ func TestLoadFromSettingsWithReader_MalformedJSONFallback(t *testing.T) {
t.Setenv("SANDBOX_PROVIDER_TYPE", "")
t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL)
t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s")
ctx := t.Context()
r := &fakeSettingsReader{
rows: map[string][]entity.SystemSettings{
@@ -550,7 +555,7 @@ func TestLoadFromSettingsWithReader_MalformedJSONFallback(t *testing.T) {
},
}
m := &ProviderManager{}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader (malformed JSON fallback): %v", err)
}
sm, ok := m.Provider().(*SelfManagedProvider)
@@ -582,6 +587,7 @@ func TestLoadFromSettingsWithReader_UnknownProviderType(t *testing.T) {
t.Setenv("SANDBOX_PROVIDER_TYPE", "")
t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL)
t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s")
ctx := t.Context()
r := &fakeSettingsReader{
rows: map[string][]entity.SystemSettings{
@@ -589,7 +595,7 @@ func TestLoadFromSettingsWithReader_UnknownProviderType(t *testing.T) {
},
}
m := &ProviderManager{}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader (unknown type fallback): %v", err)
}
// Falls back to env-driven self_managed, NOT the unknown type.
@@ -607,13 +613,14 @@ func TestLoadFromSettingsWithReader_AlreadyLoaded_NoOp(t *testing.T) {
m := &ProviderManager{}
m.SetProvider(newSelfManagedProviderFromEnv())
original := m.Provider()
ctx := t.Context()
r := &fakeSettingsReader{
rows: map[string][]entity.SystemSettings{
"sandbox.provider_type": {{Name: "sandbox.provider_type", Value: "local"}},
},
}
if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("LoadFromSettingsWithReader: %v", err)
}
if m.Provider() != original {
@@ -636,6 +643,7 @@ func TestReloadFromSettingsWithReader(t *testing.T) {
w.WriteHeader(http.StatusNotFound)
}))
defer srv.Close()
ctx := t.Context()
r := &fakeSettingsReader{
rows: map[string][]entity.SystemSettings{
@@ -647,7 +655,7 @@ func TestReloadFromSettingsWithReader(t *testing.T) {
},
}
m := &ProviderManager{}
if err := m.ReloadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil {
if err := m.ReloadFromSettingsWithReader(ctx, dao.DB, r); err != nil {
t.Fatalf("ReloadFromSettingsWithReader: %v", err)
}
if got := m.Provider().ProviderType(); got != ProviderSelfManaged {

67
internal/agent/sandbox/self_managed_test.go
View File

@@ -17,7 +17,6 @@
package sandbox
import (
"context"
"encoding/base64"
"encoding/json"
"io"
@@ -50,9 +49,11 @@ func TestSelfManaged_HealthCheck_OK(t *testing.T) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"status":"ok"}`))
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
if err := p.HealthCheck(context.Background()); err != nil {
if err := p.HealthCheck(ctx); err != nil {
t.Fatalf("HealthCheck: %v", err)
}
}
@@ -63,9 +64,10 @@ func TestSelfManaged_HealthCheck_Fail(t *testing.T) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
if err := p.HealthCheck(context.Background()); err == nil {
if err := p.HealthCheck(ctx); err == nil {
t.Errorf("HealthCheck on 500: got nil error, want one")
}
}
@@ -86,9 +88,10 @@ func TestSelfManaged_Initialize(t *testing.T) {
w.WriteHeader(http.StatusNotFound)
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
if err := p.Initialize(context.Background()); err != nil {
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
if !p.isInitialized() {
@@ -103,9 +106,10 @@ func TestSelfManaged_Initialize_HealthFails(t *testing.T) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
if err := p.Initialize(context.Background()); err == nil {
if err := p.Initialize(ctx); err == nil {
t.Errorf("Initialize on 500 healthz: got nil error, want one")
}
}
@@ -114,7 +118,8 @@ func TestSelfManaged_CreateInstance(t *testing.T) {
t.Parallel()
p := newSelfManagedForTest("http://example.invalid:9999")
p.initialized = true // bypass probe for unit testing
inst, err := p.CreateInstance(context.Background(), "python")
ctx := t.Context()
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
@@ -133,7 +138,8 @@ func TestSelfManaged_CreateInstance_UnsupportedLanguage(t *testing.T) {
t.Parallel()
p := newSelfManagedForTest("http://example.invalid:9999")
p.initialized = true
if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil {
ctx := t.Context()
if _, err := p.CreateInstance(ctx, "ruby"); err == nil {
t.Errorf("CreateInstance(ruby): got nil error, want one")
}
}
@@ -153,14 +159,15 @@ func TestSelfManaged_ExecuteCode(t *testing.T) {
})
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1+1", "python", 10, nil)
result, err := p.ExecuteCode(ctx, inst, "def main(): return 1+1", "python", 10, nil)
if err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
@@ -209,14 +216,15 @@ func TestSelfManaged_ExecuteCode_JSWrapped(t *testing.T) {
handleRun(t, w, r, "ok", "")
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
inst, err := p.CreateInstance(context.Background(), "nodejs")
inst, err := p.CreateInstance(ctx, "nodejs")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
_, err = p.ExecuteCode(context.Background(), inst, "async function main() {}", "javascript", 5, nil)
_, err = p.ExecuteCode(ctx, inst, "async function main() {}", "javascript", 5, nil)
if err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
@@ -262,14 +270,15 @@ func TestSelfManaged_ExecuteCode_PrefersHTTPResultField(t *testing.T) {
}`))
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 16", "python", 10, nil)
result, err := p.ExecuteCode(ctx, inst, "def main(): return 16", "python", 10, nil)
if err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
@@ -292,11 +301,12 @@ func TestSelfManaged_ExecuteCode_Non200(t *testing.T) {
_, _ = w.Write([]byte("bad code"))
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
inst, _ := p.CreateInstance(context.Background(), "python")
_, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil)
inst, _ := p.CreateInstance(ctx, "python")
_, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil)
if err == nil {
t.Errorf("ExecuteCode on 400: got nil error, want one")
}
@@ -307,10 +317,11 @@ func TestSelfManaged_ExecuteCode_Non200(t *testing.T) {
func TestSelfManaged_ExecuteCode_NotInitialized(t *testing.T) {
t.Parallel()
ctx := t.Context()
p := newSelfManagedForTest("http://example.invalid:9999")
// do NOT set initialized
inst := &SandboxInstance{InstanceID: "x"}
_, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil)
_, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil)
if err == nil {
t.Errorf("ExecuteCode on uninitialized: got nil error, want one")
}
@@ -318,10 +329,11 @@ func TestSelfManaged_ExecuteCode_NotInitialized(t *testing.T) {
func TestSelfManaged_ExecuteCode_UnsupportedLanguage(t *testing.T) {
t.Parallel()
ctx := t.Context()
p := newSelfManagedForTest("http://example.invalid:9999")
p.initialized = true
inst, _ := p.CreateInstance(context.Background(), "python")
_, err := p.ExecuteCode(context.Background(), inst, "x", "ruby", 5, nil)
inst, _ := p.CreateInstance(ctx, "python")
_, err := p.ExecuteCode(ctx, inst, "x", "ruby", 5, nil)
if err == nil {
t.Errorf("ExecuteCode(ruby): got nil error, want one")
}
@@ -329,9 +341,10 @@ func TestSelfManaged_ExecuteCode_UnsupportedLanguage(t *testing.T) {
func TestSelfManaged_DestroyInstance_Noop(t *testing.T) {
t.Parallel()
ctx := t.Context()
p := newSelfManagedForTest("http://example.invalid:9999")
p.initialized = true
if err := p.DestroyInstance(context.Background(), &SandboxInstance{InstanceID: "x"}); err != nil {
if err := p.DestroyInstance(ctx, &SandboxInstance{InstanceID: "x"}); err != nil {
t.Errorf("DestroyInstance: %v", err)
}
}
@@ -378,7 +391,7 @@ func TestNewSelfManagedProviderFromEnv_BaseImages(t *testing.T) {
t.Errorf("nodejs baseImage = (%q, %v); want (\"\", true)", got, ok)
}
// Case 3: only python set. nodejs slot must be empty.
// Case 3: only python set. Node.js slot must be empty.
t.Setenv("SANDBOX_BASE_PYTHON_IMAGE", "only-python:latest")
t.Setenv("SANDBOX_BASE_NODEJS_IMAGE", "")
p3 := newSelfManagedProviderFromEnv()
@@ -402,6 +415,7 @@ func TestSelfManaged_ExecuteCode_PassesBaseImage(t *testing.T) {
handleRun(t, w, r, "ok", "")
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
@@ -409,15 +423,15 @@ func TestSelfManaged_ExecuteCode_PassesBaseImage(t *testing.T) {
"python": "custom-python:v1",
"nodejs": "",
}
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
if _, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 10, nil); err != nil {
if _, err = p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 10, nil); err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
var payload map[string]any
if err := json.Unmarshal(capturedBody, &payload); err != nil {
if err = json.Unmarshal(capturedBody, &payload); err != nil {
t.Fatalf("decode: %v (raw=%s)", err, capturedBody)
}
if got := payload["base_image"]; got != "custom-python:v1" {
@@ -438,6 +452,7 @@ func TestSelfManaged_ExecuteCode_OmitsEmptyBaseImage(t *testing.T) {
handleRun(t, w, r, "ok", "")
}))
defer srv.Close()
ctx := t.Context()
p := newSelfManagedForTest(srv.URL)
p.initialized = true
@@ -445,15 +460,15 @@ func TestSelfManaged_ExecuteCode_OmitsEmptyBaseImage(t *testing.T) {
"python": "", // operator did not override
"nodejs": "",
}
inst, err := p.CreateInstance(context.Background(), "python")
inst, err := p.CreateInstance(ctx, "python")
if err != nil {
t.Fatalf("CreateInstance: %v", err)
}
if _, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 10, nil); err != nil {
if _, err = p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 10, nil); err != nil {
t.Fatalf("ExecuteCode: %v", err)
}
var payload map[string]any
if err := json.Unmarshal(capturedBody, &payload); err != nil {
if err = json.Unmarshal(capturedBody, &payload); err != nil {
t.Fatalf("decode: %v (raw=%s)", err, capturedBody)
}
if _, present := payload["base_image"]; present {

39
internal/agent/sandbox/ssh_test.go
View File

@@ -102,11 +102,12 @@ func TestSSH_PrivateKeyInline(t *testing.T) {
}
func TestSSH_Initialize_MissingHost(t *testing.T) {
ctx := t.Context()
t.Setenv("SSH_HOST", "")
t.Setenv("SSH_USERNAME", "u")
t.Setenv("SSH_PASSWORD", "p")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err == nil {
if err := p.Initialize(ctx); err == nil {
t.Errorf("Initialize with empty host: got nil error, want one")
} else if !strings.Contains(err.Error(), "SSH_HOST") {
t.Errorf("err = %v, want to mention SSH_HOST", err)
@@ -114,23 +115,25 @@ func TestSSH_Initialize_MissingHost(t *testing.T) {
}
func TestSSH_Initialize_MissingUsername(t *testing.T) {
ctx := t.Context()
t.Setenv("SSH_HOST", "h")
t.Setenv("SSH_USERNAME", "")
t.Setenv("SSH_PASSWORD", "p")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err == nil {
if err := p.Initialize(ctx); err == nil {
t.Errorf("Initialize with empty username: got nil error, want one")
}
}
func TestSSH_Initialize_MissingAuth(t *testing.T) {
ctx := t.Context()
t.Setenv("SSH_HOST", "h")
t.Setenv("SSH_USERNAME", "u")
t.Setenv("SSH_PASSWORD", "")
t.Setenv("SSH_PRIVATE_KEY", "")
t.Setenv("SSH_PRIVATE_KEY_PATH", "")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err == nil {
if err := p.Initialize(ctx); err == nil {
t.Errorf("Initialize with no auth: got nil error, want one")
} else if !strings.Contains(err.Error(), "SSH_PASSWORD") {
t.Errorf("err = %v, want to mention SSH_PASSWORD", err)
@@ -139,31 +142,33 @@ func TestSSH_Initialize_MissingAuth(t *testing.T) {
func TestSSH_AllOps_BeforeInit(t *testing.T) {
t.Parallel()
ctx := t.Context()
p := &SSHProvider{}
inst := &SandboxInstance{InstanceID: "x", Provider: ProviderSSH}
if _, err := p.CreateInstance(context.Background(), "python"); err == nil {
if _, err := p.CreateInstance(ctx, "python"); err == nil {
t.Errorf("CreateInstance before init: got nil error, want one")
}
if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil {
if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil {
t.Errorf("ExecuteCode before init: got nil error, want one")
}
if err := p.DestroyInstance(context.Background(), inst); err == nil {
if err := p.DestroyInstance(ctx, inst); err == nil {
t.Errorf("DestroyInstance before init: got nil error, want one")
}
if err := p.HealthCheck(context.Background()); err == nil {
if err := p.HealthCheck(ctx); err == nil {
t.Errorf("HealthCheck before init: got nil error, want one")
}
}
func TestSSH_CreateInstance_RejectsBadLanguage(t *testing.T) {
ctx := t.Context()
t.Setenv("SSH_HOST", "h")
t.Setenv("SSH_USERNAME", "u")
t.Setenv("SSH_PASSWORD", "p")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err != nil {
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil {
if _, err := p.CreateInstance(ctx, "ruby"); err == nil {
t.Errorf("CreateInstance(ruby): got nil error, want one")
}
}
@@ -172,6 +177,7 @@ func TestSSH_CreateInstance_RejectsBadLanguage(t *testing.T) {
// a clear error when the host is unreachable. We bind then close
// an ephemeral listener to obtain a guaranteed-closed port.
func TestSSH_Dial_ConnectionRefused(t *testing.T) {
ctx := t.Context()
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen for ephemeral port: %v", err)
@@ -185,22 +191,23 @@ func TestSSH_Dial_ConnectionRefused(t *testing.T) {
t.Setenv("SSH_PASSWORD", "p")
t.Setenv("SSH_TIMEOUT", "2")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err != nil {
if err = p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
newCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
if _, err := p.dial(ctx); err == nil {
if _, err = p.dial(newCtx); err == nil {
t.Errorf("dial: got nil error, want one")
}
}
func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) {
ctx := t.Context()
t.Setenv("SSH_HOST", "h")
t.Setenv("SSH_USERNAME", "u")
t.Setenv("SSH_PASSWORD", "p")
p := newSSHProviderFromEnv()
if err := p.Initialize(context.Background()); err != nil {
if err := p.Initialize(ctx); err != nil {
t.Fatalf("Initialize: %v", err)
}
@@ -212,7 +219,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "empty instance id",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: ""}, "x", "python", 5, nil)
return err
},
@@ -221,7 +228,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "unsupported language",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil)
return err
},
@@ -230,7 +237,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) {
{
name: "unknown instance id",
fn: func() error {
_, err := p.ExecuteCode(context.Background(),
_, err := p.ExecuteCode(ctx,
&SandboxInstance{InstanceID: "nope"}, "x", "python", 5, nil)
return err
},