From d3aa6c2110290e428a21d34e2fae5b7055afe7ac Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 09:21:30 +0000 Subject: [PATCH] test: enhance test coverage across agent, db, utils, and CLI commands - Add end-to-end CLI behavioral tests in `cmd/cli_test.go` covering root, run, search, status, show, list, and generate-image commands. - Enhance agent integration tests in `internal/agent/client_test.go` and `internal/agent/researcher_test.go` to test streaming, polling, status, and image generation using realistic test servers. - Enhance database unit tests in `internal/db/db_test.go` for task queries, status updates, and boundary conditions, adding `ResetDBForTesting` for test isolation. - Expand path validation, sanitization, file saving, and env lookup tests in `internal/utils/fs_test.go`. Co-authored-by: rmedranollamas <45878745+rmedranollamas@users.noreply.github.com> --- cmd/cli_test.go | 254 ++++++++++++++++++++++++++++ internal/agent/client_test.go | 88 +++++++++- internal/agent/researcher_test.go | 268 +++++++++++++++++++++++++++++- internal/db/db.go | 8 + internal/db/db_test.go | 120 +++++++++++-- internal/utils/fs_test.go | 184 ++++++++++++++------ 6 files changed, 850 insertions(+), 72 deletions(-) create mode 100644 cmd/cli_test.go diff --git a/cmd/cli_test.go b/cmd/cli_test.go new file mode 100644 index 0000000..ed67532 --- /dev/null +++ b/cmd/cli_test.go @@ -0,0 +1,254 @@ +package cmd + +import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/google/research-cli/internal/config" + "github.com/google/research-cli/internal/db" + "github.com/spf13/cobra" +) + +func setupCmdTestEnv(t *testing.T) (*httptest.Server, string, string) { + t.Helper() + db.ResetDBForTesting() + + workspace := t.TempDir() + dbPath := filepath.Join(t.TempDir(), "history.db") + + oldWorkspace := config.WorkspaceDir + oldDBPath := config.DbPath + oldBaseURL := config.GeminiApiBaseUrl + + config.WorkspaceDir = workspace + config.DbPath = dbPath + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "/v1alpha/interactions") && r.Method == "POST" { + if strings.Contains(r.URL.RawQuery, "alt=sse") { + w.WriteHeader(http.StatusOK) + w.Write([]byte("data: {\"interaction\":{\"id\":\"cmd-inter-1\"},\"delta\":{\"type\":\"text\",\"text\":\"CLI report output\"}}\n")) + return + } + // Image generation or JSON POST + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"outputs":[{"type":"image","data":"aGVsbG8="}]}`)) + return + } + if strings.Contains(r.URL.Path, "/v1alpha/interactions/status-id") { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"status":"COMPLETED","outputs":[{"text":"status report text"}]}`)) + return + } + w.WriteHeader(http.StatusOK) + })) + + config.GeminiApiBaseUrl = ts.URL + t.Setenv(config.GeminiApiKeyVar, "test-api-key") + + // Ensure PersistentPreRun uses our test baseURL + cobraPreRun := rootCmd.PersistentPreRun + rootCmd.PersistentPreRun = func(cmd *cobra.Command, args []string) { + if cobraPreRun != nil { + cobraPreRun(cmd, args) + } + config.GeminiApiBaseUrl = ts.URL + } + + // Make current working directory equal to workspace for relative path saving in tests + oldWd, err := os.Getwd() + if err == nil { + _ = os.Chdir(workspace) + } + + t.Cleanup(func() { + ts.Close() + db.ResetDBForTesting() + config.WorkspaceDir = oldWorkspace + config.DbPath = oldDBPath + config.GeminiApiBaseUrl = oldBaseURL + rootCmd.PersistentPreRun = cobraPreRun + if oldWd != "" { + _ = os.Chdir(oldWd) + } + }) + + return ts, workspace, dbPath +} + +func TestRootCmdVersion(t *testing.T) { + buf := new(bytes.Buffer) + rootCmd.SetOut(buf) + rootCmd.SetErr(buf) + rootCmd.SetArgs([]string{"--version"}) + + oldVersion := version + version = "1.2.3" + defer func() { version = oldVersion }() + + err := rootCmd.Execute() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestRootCmdDefaultQuery(t *testing.T) { + _, workspace, _ := setupCmdTestEnv(t) + _ = workspace + + rootCmd.SetArgs([]string{"quantum computing"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("rootCmd default query error: %v", err) + } +} + +func TestRunCmd(t *testing.T) { + _, workspace, _ := setupCmdTestEnv(t) + outPath := "run_report.md" + + t.Run("basic run command with output flag", func(t *testing.T) { + rootCmd.SetArgs([]string{"run", "deep learning query", "-o", outPath, "--force", "--plan", "--vis", "-v"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("run command error: %v", err) + } + saved, err := os.ReadFile(filepath.Join(workspace, outPath)) + if err != nil { + t.Fatalf("failed to read saved report: %v", err) + } + if string(saved) != "CLI report output" { + t.Fatalf("saved report = %q, want 'CLI report output'", string(saved)) + } + }) + + t.Run("run command missing API key", func(t *testing.T) { + t.Setenv(config.GeminiApiKeyVar, "") + rootCmd.SetArgs([]string{"run", "query without key"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error when API key is missing, got nil") + } + }) +} + +func TestSearchCmd(t *testing.T) { + _, workspace, _ := setupCmdTestEnv(t) + outPath := "search_report.md" + + t.Run("search command success with output", func(t *testing.T) { + rootCmd.SetArgs([]string{"search", "fast search query", "-o", outPath, "-f"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("search command error: %v", err) + } + saved, err := os.ReadFile(filepath.Join(workspace, outPath)) + if err != nil { + t.Fatalf("failed to read saved search report: %v", err) + } + if string(saved) != "CLI report output" { + t.Fatalf("saved report = %q, want 'CLI report output'", string(saved)) + } + }) +} + +func TestStatusCmd(t *testing.T) { + setupCmdTestEnv(t) + + rootCmd.SetArgs([]string{"status", "status-id"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("status command error: %v", err) + } +} + +func TestShowCmd(t *testing.T) { + _, workspace, _ := setupCmdTestEnv(t) + + reportText := "DB saved report" + taskID, err := db.SaveTask("show query", "model-x", nil, nil) + if err != nil { + t.Fatal(err) + } + if err := db.UpdateTask(taskID, "COMPLETED", &reportText, nil); err != nil { + t.Fatal(err) + } + + t.Run("show existing task", func(t *testing.T) { + outPath := "show_report.md" + rootCmd.SetArgs([]string{"show", fmt.Sprintf("%d", taskID), "-o", outPath, "-f"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("show command error: %v", err) + } + saved, err := os.ReadFile(filepath.Join(workspace, outPath)) + if err != nil { + t.Fatal(err) + } + if string(saved) != reportText { + t.Fatalf("saved report = %q, want %q", string(saved), reportText) + } + }) + + t.Run("show non-existent task", func(t *testing.T) { + rootCmd.SetArgs([]string{"show", "999999"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for non-existent task, got nil") + } + if !strings.Contains(err.Error(), "task not found") { + t.Fatalf("unexpected error message: %v", err) + } + }) + + t.Run("show invalid task ID", func(t *testing.T) { + rootCmd.SetArgs([]string{"show", "invalid-id"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for invalid task ID, got nil") + } + if !strings.Contains(err.Error(), "invalid task ID") { + t.Fatalf("unexpected error message: %v", err) + } + }) +} + +func TestListCmd(t *testing.T) { + setupCmdTestEnv(t) + + _, err := db.SaveTask("list query 1", "model-1", nil, nil) + if err != nil { + t.Fatal(err) + } + + rootCmd.SetArgs([]string{"list", "-n", "5"}) + err = rootCmd.Execute() + if err != nil { + t.Fatalf("list command error: %v", err) + } +} + +func TestGenerateImageCmd(t *testing.T) { + _, workspace, _ := setupCmdTestEnv(t) + outPath := "test_generated.png" + + rootCmd.SetArgs([]string{"generate-image", "a futuristic city", "-o", outPath, "-f"}) + err := rootCmd.Execute() + if err != nil { + t.Fatalf("generate-image command error: %v", err) + } + + got, err := os.ReadFile(filepath.Join(workspace, outPath)) + if err != nil { + t.Fatal(err) + } + if string(got) != "hello" { + t.Fatalf("generated image content = %q, want 'hello'", string(got)) + } +} diff --git a/internal/agent/client_test.go b/internal/agent/client_test.go index 22a0b81..77f43ae 100644 --- a/internal/agent/client_test.go +++ b/internal/agent/client_test.go @@ -1,6 +1,91 @@ package agent -import "testing" +import ( + "path/filepath" + "strings" + "testing" + + "github.com/google/research-cli/internal/config" +) + +func TestNewResearchAgent(t *testing.T) { + t.Run("valid secure baseURL", func(t *testing.T) { + agent, err := NewResearchAgent("test-api-key", "https://api.example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if agent == nil { + t.Fatal("expected agent, got nil") + } + if agent.GetClient() == nil { + t.Fatal("expected non-nil genai.Client from GetClient()") + } + }) + + t.Run("valid loopback http baseURL", func(t *testing.T) { + agent, err := NewResearchAgent("test-api-key", "http://127.0.0.1:8080") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if agent == nil { + t.Fatal("expected agent, got nil") + } + }) + + t.Run("insecure remote http baseURL rejected", func(t *testing.T) { + agent, err := NewResearchAgent("test-api-key", "http://api.example.com") + if err == nil { + t.Fatal("expected error for insecure http baseURL, got nil") + } + if agent != nil { + t.Fatal("expected nil agent on error") + } + expectedSubstr := "insecure baseURL" + if !strings.Contains(err.Error(), expectedSubstr) { + t.Errorf("error %q does not contain %q", err.Error(), expectedSubstr) + } + }) + + t.Run("empty baseURL uses default", func(t *testing.T) { + agent, err := NewResearchAgent("test-api-key", "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if agent == nil { + t.Fatal("expected agent, got nil") + } + }) +} + +func TestUploadFilesEmpty(t *testing.T) { + agent := &ResearchAgent{} + uris, err := agent.UploadFiles(t.Context(), []string{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(uris) != 0 { + t.Fatalf("expected empty uris, got %v", uris) + } +} + +func TestUploadFileNonExistent(t *testing.T) { + workspace := t.TempDir() + oldWorkspace := config.WorkspaceDir + config.WorkspaceDir = workspace + t.Cleanup(func() { + config.WorkspaceDir = oldWorkspace + }) + + agent := &ResearchAgent{} + nonExistentPath := filepath.Join("sub", "nonexistent.txt") + _, err := agent.uploadFile(t.Context(), nonExistentPath) + if err == nil { + t.Fatal("expected error for missing file, got nil") + } + if !strings.Contains(err.Error(), "file not found") { + t.Fatalf("error %q does not contain 'file not found'", err.Error()) + } +} func TestIsSecureOrLoopbackBaseURL(t *testing.T) { tests := []struct { @@ -16,6 +101,7 @@ func TestIsSecureOrLoopbackBaseURL(t *testing.T) { {name: "http ipv4 prefix", url: "http://127.0.0.1.evil.test", want: false}, {name: "http remote", url: "http://api.example.test", want: false}, {name: "unsupported scheme", url: "ftp://localhost", want: false}, + {name: "invalid url", url: "http://[::1", want: false}, } for _, tt := range tests { diff --git a/internal/agent/researcher_test.go b/internal/agent/researcher_test.go index 84cdb13..785e4b7 100644 --- a/internal/agent/researcher_test.go +++ b/internal/agent/researcher_test.go @@ -1,16 +1,32 @@ package agent import ( + "encoding/base64" + "fmt" "io" "net/http" + "net/http/httptest" + "os" "path/filepath" "strings" "testing" "time" "github.com/google/research-cli/internal/config" + "github.com/google/research-cli/internal/db" ) +func resetDBForAgentTest(t *testing.T) { + t.Helper() + db.ResetDBForTesting() + oldDBPath := config.DbPath + config.DbPath = filepath.Join(t.TempDir(), "history.db") + t.Cleanup(func() { + db.ResetDBForTesting() + config.DbPath = oldDBPath + }) +} + type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { @@ -82,6 +98,7 @@ func TestGetToolsIncludesRequestedToolsAndMCPServers(t *testing.T) { } func TestProcessSSELineAppendsContentAndDelta(t *testing.T) { + resetDBForAgentTest(t) line := `data: {"content":{"parts":[{"text":"hello " }]},"delta":{"type":"text","text":"world"}}` var interactionID string var reportParts []string @@ -93,13 +110,50 @@ func TestProcessSSELineAppendsContentAndDelta(t *testing.T) { } } -func TestStreamInteractionReadsSSE(t *testing.T) { - oldDBPath := config.DbPath - config.DbPath = filepath.Join(t.TempDir(), "history.db") +func TestProcessSSELineImageDelta(t *testing.T) { + resetDBForAgentTest(t) + workspace := t.TempDir() + oldWorkspace := config.WorkspaceDir + config.WorkspaceDir = workspace t.Cleanup(func() { - config.DbPath = oldDBPath + config.WorkspaceDir = oldWorkspace }) + pngBytes := []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a} + encoded := base64.StdEncoding.EncodeToString(pngBytes) + line := fmt.Sprintf(`data: {"delta":{"type":"image","data":"%s"}}`, encoded) + + var interactionID string + var reportParts []string + + processSSELine(line, &interactionID, &reportParts, 123, true) + + files, err := os.ReadDir(workspace) + if err != nil { + t.Fatal(err) + } + if len(files) == 0 { + t.Fatal("expected image file to be saved, found none") + } + if !strings.HasPrefix(files[0].Name(), "research_task_123_") { + t.Fatalf("unexpected filename: %s", files[0].Name()) + } +} + +func TestProcessSSELineInvalidJSON(t *testing.T) { + resetDBForAgentTest(t) + line := `data: {invalid json` + var interactionID string + var reportParts []string + processSSELine(line, &interactionID, &reportParts, 1, false) + if len(reportParts) != 0 { + t.Fatalf("expected 0 report parts for invalid json, got %d", len(reportParts)) + } +} + +func TestStreamInteractionReadsSSE(t *testing.T) { + resetDBForAgentTest(t) + client := &http.Client{ Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { if req.URL.String() != "https://proxy.example.test/v1alpha/interactions?alt=sse" { @@ -128,3 +182,209 @@ func TestStreamInteractionReadsSSE(t *testing.T) { t.Fatalf("report = %q, want done", report) } } + +func TestStreamInteractionHTTPError(t *testing.T) { + resetDBForAgentTest(t) + + client := &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Body: io.NopCloser(strings.NewReader("bad request")), + Header: make(http.Header), + Request: req, + }, nil + }), + } + + a := &ResearchAgent{ + apiKey: "test-key", + baseURL: "https://proxy.example.test", + httpClient: client, + } + + _, err := a.streamInteraction(t.Context(), 123, InteractionRequest{Input: "query"}, false) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "status 400") { + t.Fatalf("unexpected error message: %v", err) + } +} + +func TestRunResearchAndRunSearch(t *testing.T) { + resetDBForAgentTest(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("data: {\"interaction\":{\"id\":\"inter-123\"},\"delta\":{\"type\":\"text\",\"text\":\"result report\"}}\n")) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + t.Run("RunResearch success", func(t *testing.T) { + report, err := a.RunResearch(t.Context(), "research query", "model-1", "parent-1", []string{"https://example.com"}, []string{"file://uri1"}, true, "high", true, true, false) + if err != nil { + t.Fatalf("RunResearch error: %v", err) + } + if report != "result report" { + t.Fatalf("report = %q, want 'result report'", report) + } + }) + + t.Run("RunSearch success", func(t *testing.T) { + report, err := a.RunSearch(t.Context(), "search query", "model-2", "", false) + if err != nil { + t.Fatalf("RunSearch error: %v", err) + } + if report != "result report" { + t.Fatalf("report = %q, want 'result report'", report) + } + }) +} + +func TestPollInteractionSuccess(t *testing.T) { + resetDBForAgentTest(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"status":"COMPLETED","outputs":[{"text":"polled output"}]}`)) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + report, err := a.GetStatus(t.Context(), "inter-999") + if err != nil { + t.Fatalf("GetStatus error: %v", err) + } + if report != "polled output" { + t.Fatalf("report = %q, want 'polled output'", report) + } +} + +func TestPollInteractionFailed(t *testing.T) { + resetDBForAgentTest(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"status":"FAILED","error":"task execution error"}`)) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + _, err = a.GetStatus(t.Context(), "inter-failed") + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "FAILED: task execution error") { + t.Fatalf("unexpected error message: %v", err) + } +} + +func TestGenerateImageSuccess(t *testing.T) { + resetDBForAgentTest(t) + workspace := t.TempDir() + oldWorkspace := config.WorkspaceDir + config.WorkspaceDir = workspace + t.Cleanup(func() { + config.WorkspaceDir = oldWorkspace + }) + + imgData := []byte("fake-image-bytes") + b64Data := base64.StdEncoding.EncodeToString(imgData) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(fmt.Sprintf(`{"outputs":[{"type":"image","data":"%s"}]}`, b64Data))) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + outPath := filepath.Join(workspace, "gen.png") + err = a.GenerateImage(t.Context(), "draw a cat", outPath, "model-img", true) + if err != nil { + t.Fatalf("GenerateImage error: %v", err) + } + + got, err := os.ReadFile(outPath) + if err != nil { + t.Fatal(err) + } + if string(got) != string(imgData) { + t.Fatalf("saved image content = %q, want %q", string(got), string(imgData)) + } +} + +func TestGenerateImageNoOutput(t *testing.T) { + resetDBForAgentTest(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"outputs":[]}`)) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + err = a.GenerateImage(t.Context(), "prompt", "out.png", "model-img", true) + if err == nil { + t.Fatal("expected error when no image output, got nil") + } + if !strings.Contains(err.Error(), "no image was generated") { + t.Fatalf("unexpected error message: %v", err) + } +} + +func TestStreamInteractionFallbackPolling(t *testing.T) { + resetDBForAgentTest(t) + + // Pre-create task in DB + taskID, err := db.SaveTask("query", "model", nil, nil) + if err != nil { + t.Fatal(err) + } + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "interactions/poll-id") { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"status":"COMPLETED","response":{"text":"polled fallback text"}}`)) + return + } + // SSE stream returns interaction ID without report parts + w.WriteHeader(http.StatusOK) + w.Write([]byte("data: {\"interaction\":{\"id\":\"poll-id\"}}\n")) + })) + t.Cleanup(ts.Close) + + a, err := NewResearchAgent("test-api-key", ts.URL) + if err != nil { + t.Fatal(err) + } + + report, err := a.streamInteraction(t.Context(), taskID, InteractionRequest{Input: "query"}, false) + if err != nil { + t.Fatalf("streamInteraction fallback error: %v", err) + } + if report != "polled fallback text" { + t.Fatalf("report = %q, want 'polled fallback text'", report) + } +} diff --git a/internal/db/db.go b/internal/db/db.go index 10563c9..e301dd8 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -27,6 +27,14 @@ type Task struct { CreatedAt string } +func ResetDBForTesting() { + if db != nil { + _ = db.Close() + } + db = nil + dbOnce = sync.Once{} +} + func GetDB() (*sql.DB, error) { var err error dbOnce.Do(func() { diff --git a/internal/db/db_test.go b/internal/db/db_test.go index 1bf37b1..5a03dbe 100644 --- a/internal/db/db_test.go +++ b/internal/db/db_test.go @@ -4,7 +4,6 @@ import ( "os" "path/filepath" "strings" - "sync" "testing" "github.com/google/research-cli/internal/config" @@ -12,21 +11,11 @@ import ( func resetTestDB(t *testing.T, path string) { t.Helper() - if db != nil { - if err := db.Close(); err != nil { - t.Fatal(err) - } - } - db = nil - dbOnce = sync.Once{} + ResetDBForTesting() oldPath := config.DbPath config.DbPath = path t.Cleanup(func() { - if db != nil { - _ = db.Close() - } - db = nil - dbOnce = sync.Once{} + ResetDBForTesting() config.DbPath = oldPath }) } @@ -36,7 +25,8 @@ func TestTaskLifecycle(t *testing.T) { resetTestDB(t, dbPath) parentID := "parent-1" - taskID, err := SaveTask("query", "model", nil, &parentID) + interactionIDInit := "interaction-init" + taskID, err := SaveTask("query", "model", &interactionIDInit, &parentID) if err != nil { t.Fatal(err) } @@ -67,6 +57,85 @@ func TestTaskLifecycle(t *testing.T) { } } +func TestUpdateTaskWithoutInteractionID(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "history.db") + resetTestDB(t, dbPath) + + taskID, err := SaveTask("query without interaction", "model-b", nil, nil) + if err != nil { + t.Fatal(err) + } + + report := "failed report" + if err := UpdateTask(taskID, "FAILED", &report, nil); err != nil { + t.Fatal(err) + } + + task, err := GetTask(taskID) + if err != nil { + t.Fatal(err) + } + if task == nil { + t.Fatal("GetTask returned nil") + } + if task.Status != "FAILED" || task.Report.String != report { + t.Fatalf("GetTask() = %+v", task) + } +} + +func TestGetTaskNotFound(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "history.db") + resetTestDB(t, dbPath) + + if _, err := GetDB(); err != nil { + t.Fatal(err) + } + + task, err := GetTask(999999) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if task != nil { + t.Fatalf("expected nil for non-existent task, got %+v", task) + } +} + +func TestGetRecentTasksLimitsAndOrder(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "history.db") + resetTestDB(t, dbPath) + + for i := 1; i <= 5; i++ { + q := "query " + string(rune('0'+i)) + if _, err := SaveTask(q, "model", nil, nil); err != nil { + t.Fatal(err) + } + } + + tasks, err := GetRecentTasks(3) + if err != nil { + t.Fatal(err) + } + if len(tasks) != 3 { + t.Fatalf("expected 3 tasks, got %d", len(tasks)) + } + + tasksAll, err := GetRecentTasks(10) + if err != nil { + t.Fatal(err) + } + if len(tasksAll) != 5 { + t.Fatalf("expected 5 tasks, got %d", len(tasksAll)) + } + + tasksZero, err := GetRecentTasks(0) + if err != nil { + t.Fatal(err) + } + if len(tasksZero) != 0 { + t.Fatalf("expected 0 tasks, got %d", len(tasksZero)) + } +} + func TestGetDBCreatesPrivateDatabaseFile(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "nested", "history.db") resetTestDB(t, dbPath) @@ -84,12 +153,33 @@ func TestGetDBCreatesPrivateDatabaseFile(t *testing.T) { } } +func TestGetDBExistingDirectoryPermissions(t *testing.T) { + tmpDir := t.TempDir() + dbDir := filepath.Join(tmpDir, "existing_dir") + if err := os.Mkdir(dbDir, 0755); err != nil { + t.Fatal(err) + } + dbPath := filepath.Join(dbDir, "history.db") + resetTestDB(t, dbPath) + + if _, err := GetDB(); err != nil { + t.Fatal(err) + } + + dirInfo, err := os.Stat(dbDir) + if err != nil { + t.Fatal(err) + } + if got := dirInfo.Mode().Perm(); got != 0700 { + t.Fatalf("dbDir mode = %o, want 0700", got) + } +} + func TestGetDBRejectsSymlinks(t *testing.T) { tmpDir := t.TempDir() realDBPath := filepath.Join(tmpDir, "real.db") symlinkPath := filepath.Join(tmpDir, "symlink.db") - // Create a symlink pointing to a location where the DB will be created if err := os.Symlink(realDBPath, symlinkPath); err != nil { t.Fatal(err) } diff --git a/internal/utils/fs_test.go b/internal/utils/fs_test.go index ba8df61..c15fca6 100644 --- a/internal/utils/fs_test.go +++ b/internal/utils/fs_test.go @@ -9,59 +9,63 @@ import ( "github.com/google/research-cli/internal/config" ) -func TestSaveToFileRejectsSymlinkOverwrite(t *testing.T) { +func TestValidatePath(t *testing.T) { workspace := t.TempDir() - outside := t.TempDir() oldWorkspace := config.WorkspaceDir config.WorkspaceDir = workspace t.Cleanup(func() { config.WorkspaceDir = oldWorkspace }) - target := filepath.Join(outside, "target.txt") - if err := os.WriteFile(target, []byte("original"), 0644); err != nil { - t.Fatal(err) - } - - link := filepath.Join(workspace, "output.txt") - if err := os.Symlink(target, link); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - err := SaveToFile([]byte("new"), link, true) - if err == nil { - t.Fatal("SaveToFile succeeded; want symlink overwrite rejection") - } - - data, err := os.ReadFile(target) - if err != nil { - t.Fatal(err) - } - if string(data) != "original" { - t.Fatalf("target was overwritten: got %q", data) - } -} + t.Run("empty path", func(t *testing.T) { + _, err := ValidatePath("") + if err == nil { + t.Fatal("expected error for empty path, got nil") + } + }) -func TestSaveToFileRejectsExistingFileWithoutForce(t *testing.T) { - workspace := t.TempDir() - oldWorkspace := config.WorkspaceDir - config.WorkspaceDir = workspace - t.Cleanup(func() { - config.WorkspaceDir = oldWorkspace + t.Run("relative path inside workspace", func(t *testing.T) { + relPath := filepath.Join("sub", "file.txt") + want := filepath.Join(workspace, relPath) + got, err := ValidatePath(relPath) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != want { + t.Fatalf("ValidatePath() = %q, want %q", got, want) + } + }) + + t.Run("absolute path inside workspace", func(t *testing.T) { + absPath := filepath.Join(workspace, "sub", "file.txt") + got, err := ValidatePath(absPath) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != absPath { + t.Fatalf("ValidatePath() = %q, want %q", got, absPath) + } }) - path := filepath.Join(workspace, "output.txt") - if err := os.WriteFile(path, []byte("original"), 0644); err != nil { - t.Fatal(err) - } + t.Run("path traversal rejected", func(t *testing.T) { + if _, err := ValidatePath(filepath.Join("..", "outside.txt")); err == nil { + t.Fatal("ValidatePath accepted parent traversal") + } + }) - err := SaveToFile([]byte("new"), path, false) - if err == nil || !strings.Contains(err.Error(), "already exists") { - t.Fatalf("SaveToFile error = %v, want already exists", err) - } + t.Run("dot dot prefixed name inside workspace allowed", func(t *testing.T) { + want := filepath.Join(workspace, "..cache", "out.txt") + got, err := ValidatePath(filepath.Join(".", "..cache", "out.txt")) + if err != nil { + t.Fatalf("ValidatePath returned error for valid workspace path: %v", err) + } + if got != want { + t.Fatalf("ValidatePath() = %q, want %q", got, want) + } + }) } -func TestValidatePathAllowsDotDotPrefixedNamesInsideWorkspace(t *testing.T) { +func TestSanitizePath(t *testing.T) { workspace := t.TempDir() oldWorkspace := config.WorkspaceDir config.WorkspaceDir = workspace @@ -69,17 +73,30 @@ func TestValidatePathAllowsDotDotPrefixedNamesInsideWorkspace(t *testing.T) { config.WorkspaceDir = oldWorkspace }) - want := filepath.Join(workspace, "..cache", "out.txt") - got, err := ValidatePath(filepath.Join(".", "..cache", "out.txt")) - if err != nil { - t.Fatalf("ValidatePath returned error for valid workspace path: %v", err) - } - if got != want { - t.Fatalf("ValidatePath() = %q, want %q", got, want) - } + t.Run("empty path", func(t *testing.T) { + if got := SanitizePath(""); got != "" { + t.Fatalf("SanitizePath(\"\") = %q, want empty string", got) + } + }) + + t.Run("path inside workspace", func(t *testing.T) { + fullPath := filepath.Join(workspace, "dir", "file.txt") + want := filepath.Join("dir", "file.txt") + if got := SanitizePath(fullPath); got != want { + t.Fatalf("SanitizePath(%q) = %q, want %q", fullPath, got, want) + } + }) + + t.Run("path outside workspace returns base name", func(t *testing.T) { + outsidePath := filepath.Join(t.TempDir(), "outside", "file.txt") + want := "file.txt" + if got := SanitizePath(outsidePath); got != want { + t.Fatalf("SanitizePath(%q) = %q, want %q", outsidePath, got, want) + } + }) } -func TestValidatePathRejectsParentTraversal(t *testing.T) { +func TestSaveToFile(t *testing.T) { workspace := t.TempDir() oldWorkspace := config.WorkspaceDir config.WorkspaceDir = workspace @@ -87,9 +104,72 @@ func TestValidatePathRejectsParentTraversal(t *testing.T) { config.WorkspaceDir = oldWorkspace }) - if _, err := ValidatePath(filepath.Join("..", "outside.txt")); err == nil { - t.Fatal("ValidatePath accepted parent traversal") - } + t.Run("save new file", func(t *testing.T) { + path := filepath.Join(workspace, "new.txt") + content := []byte("hello world") + if err := SaveToFile(content, path, false); err != nil { + t.Fatalf("unexpected error: %v", err) + } + got, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if string(got) != string(content) { + t.Fatalf("file content = %q, want %q", string(got), string(content)) + } + }) + + t.Run("overwrite existing file with force=true", func(t *testing.T) { + path := filepath.Join(workspace, "force.txt") + if err := os.WriteFile(path, []byte("old"), 0644); err != nil { + t.Fatal(err) + } + content := []byte("new content") + if err := SaveToFile(content, path, true); err != nil { + t.Fatalf("unexpected error: %v", err) + } + got, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if string(got) != string(content) { + t.Fatalf("file content = %q, want %q", string(got), string(content)) + } + }) + + t.Run("reject existing file without force", func(t *testing.T) { + path := filepath.Join(workspace, "existing.txt") + if err := os.WriteFile(path, []byte("original"), 0644); err != nil { + t.Fatal(err) + } + err := SaveToFile([]byte("new"), path, false) + if err == nil || !strings.Contains(err.Error(), "already exists") { + t.Fatalf("SaveToFile error = %v, want already exists", err) + } + }) + + t.Run("reject symlink overwrite", func(t *testing.T) { + outside := t.TempDir() + target := filepath.Join(outside, "target.txt") + if err := os.WriteFile(target, []byte("original"), 0644); err != nil { + t.Fatal(err) + } + link := filepath.Join(workspace, "output.txt") + if err := os.Symlink(target, link); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + err := SaveToFile([]byte("new"), link, true) + if err == nil { + t.Fatal("SaveToFile succeeded; want symlink overwrite rejection") + } + data, err := os.ReadFile(target) + if err != nil { + t.Fatal(err) + } + if string(data) != "original" { + t.Fatalf("target was overwritten: got %q", string(data)) + } + }) } func TestGetApiKey(t *testing.T) {