From ce0143ca06c25cafdaf6b4051809fa05dbd89da6 Mon Sep 17 00:00:00 2001 From: Marvin Zhang Date: Thu, 13 Mar 2025 18:10:24 +0800 Subject: [PATCH] refactor: enhance health check function and add comprehensive test coverage - Updated GetHealthFn to return an error for better error handling and clarity. - Introduced a new test file for schedule management, covering various endpoints including creation, retrieval, updating, and deletion of schedules. - Added tests for task management, including task creation, retrieval, updating, and cancellation. - Implemented utility tests for filtering and response generation to ensure consistent API behavior. - Improved logging in the task scheduler service for better traceability. --- core/controllers/health.go | 6 +- core/controllers/schedule_test.go | 376 ++++++++++++++++++++++++++++++ core/controllers/spider_test.go | 5 +- core/controllers/task_test.go | 343 +++++++++++++++++++++++++++ core/controllers/utils_test.go | 178 ++++++++++++++ core/task/scheduler/service.go | 1 + 6 files changed, 906 insertions(+), 3 deletions(-) create mode 100644 core/controllers/schedule_test.go create mode 100644 core/controllers/task_test.go create mode 100644 core/controllers/utils_test.go diff --git a/core/controllers/health.go b/core/controllers/health.go index 8e19adb1..bf687a5e 100644 --- a/core/controllers/health.go +++ b/core/controllers/health.go @@ -1,12 +1,13 @@ package controllers import ( + "errors" "github.com/gin-gonic/gin" "net/http" ) -func GetHealthFn(healthFn func() bool) func(c *gin.Context) { - return func(c *gin.Context) { +func GetHealthFn(healthFn func() bool) func(c *gin.Context) error { + return func(c *gin.Context) (err error) { if healthFn() { c.Writer.Write([]byte("ok")) c.AbortWithStatus(http.StatusOK) @@ -14,5 +15,6 @@ func GetHealthFn(healthFn func() bool) func(c *gin.Context) { } c.Writer.Write([]byte("not ready")) c.AbortWithStatus(http.StatusServiceUnavailable) + return errors.New("not ready") } } diff --git a/core/controllers/schedule_test.go b/core/controllers/schedule_test.go new file mode 100644 index 00000000..9ea635f2 --- /dev/null +++ b/core/controllers/schedule_test.go @@ -0,0 +1,376 @@ +package controllers_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/crawlab-team/crawlab/core/controllers" + "github.com/crawlab-team/crawlab/core/middlewares" + "github.com/crawlab-team/crawlab/core/models/models" + "github.com/crawlab-team/crawlab/core/models/service" + "github.com/gin-gonic/gin" + "github.com/loopfz/gadgeto/tonic" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/bson/primitive" +) + +// Helper function to create a test schedule +func createTestSchedule(t *testing.T) (*models.Schedule, primitive.ObjectID) { + // First, create a test spider + spider := models.Spider{ + Name: "Test Spider for Schedule", + ColName: "test_spider_for_schedule", + } + spiderSvc := service.NewModelService[models.Spider]() + spiderId, err := spiderSvc.InsertOne(spider) + require.NoError(t, err) + require.False(t, spiderId.IsZero()) + + // Now create a schedule associated with the spider + schedule := &models.Schedule{ + Name: "Test Schedule", + SpiderId: spiderId, + Cron: "0 0 * * *", // Run daily at midnight + Cmd: "python main.py", + Param: "test param", + Priority: 5, + Enabled: false, // Disabled initially + } + + // Set timestamps + now := time.Now() + schedule.CreatedAt = now + schedule.UpdatedAt = now + + scheduleSvc := service.NewModelService[models.Schedule]() + scheduleId, err := scheduleSvc.InsertOne(*schedule) + require.NoError(t, err) + require.False(t, scheduleId.IsZero()) + + schedule.Id = scheduleId + return schedule, spiderId +} + +// Test PostSchedule endpoint +func TestPostSchedule(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test spider first + spider := models.Spider{ + Name: "Test Spider for Schedule Post", + ColName: "test_spider_for_schedule_post", + } + spiderSvc := service.NewModelService[models.Spider]() + spiderId, err := spiderSvc.InsertOne(spider) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/schedules", nil, tonic.Handler(controllers.PostSchedule, 200)) + + // Create payload + schedule := models.Schedule{ + Name: "Test Schedule for Post", + SpiderId: spiderId, + Cron: "0 0 * * *", + Cmd: "python main.py", + Param: "test schedule param", + Priority: 3, + Enabled: false, + } + + payload := controllers.PostScheduleParams{ + Data: schedule, + } + jsonValue, _ := json.Marshal(payload) + + // Create test request + req, err := http.NewRequest("POST", "/schedules", bytes.NewBuffer(jsonValue)) + req.Header.Set("Authorization", TestToken) + req.Header.Set("Content-Type", "application/json") + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.Response[models.Schedule] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.Equal(t, "ok", response.Status) + assert.False(t, response.Data.Id.IsZero()) + assert.Equal(t, schedule.Name, response.Data.Name) + assert.Equal(t, schedule.SpiderId, response.Data.SpiderId) + assert.Equal(t, schedule.Cron, response.Data.Cron) + assert.Equal(t, schedule.Enabled, response.Data.Enabled) +} + +// Test GetScheduleById endpoint +func TestGetScheduleById(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test schedule + schedule, _ := createTestSchedule(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.GET("/schedules/:id", nil, tonic.Handler(controllers.NewController[models.Schedule]().GetById, 200)) + + // Create test request + req, err := http.NewRequest("GET", "/schedules/"+schedule.Id.Hex(), nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.Response[models.Schedule] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.Equal(t, "ok", response.Status) + assert.Equal(t, schedule.Id, response.Data.Id) + assert.Equal(t, schedule.Name, response.Data.Name) + assert.Equal(t, schedule.Cron, response.Data.Cron) +} + +// Test PutScheduleById endpoint +func TestPutScheduleById(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test schedule + schedule, _ := createTestSchedule(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.PUT("/schedules/:id", nil, tonic.Handler(controllers.PutScheduleById, 200)) + + // Update the schedule + updatedSchedule := *schedule + updatedSchedule.Name = "Updated Schedule Name" + updatedSchedule.Cron = "*/5 * * * *" // Every 5 minutes + updatedSchedule.Enabled = true + + payload := controllers.PutScheduleByIdParams{ + Id: schedule.Id.Hex(), + Data: updatedSchedule, + } + jsonValue, _ := json.Marshal(payload) + + // Create test request + req, err := http.NewRequest("PUT", "/schedules/"+schedule.Id.Hex(), bytes.NewBuffer(jsonValue)) + req.Header.Set("Authorization", TestToken) + req.Header.Set("Content-Type", "application/json") + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.Response[models.Schedule] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.Equal(t, "ok", response.Status) + + // Fetch the updated schedule to verify changes + scheduleSvc := service.NewModelService[models.Schedule]() + updatedScheduleFromDB, err := scheduleSvc.GetById(schedule.Id) + require.NoError(t, err) + + assert.Equal(t, "Updated Schedule Name", updatedScheduleFromDB.Name) + assert.Equal(t, "*/5 * * * *", updatedScheduleFromDB.Cron) + assert.True(t, updatedScheduleFromDB.Enabled) +} + +// Test DeleteScheduleById endpoint +func TestDeleteScheduleById(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test schedule + schedule, _ := createTestSchedule(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.DELETE("/schedules/:id", nil, tonic.Handler(controllers.NewController[models.Schedule]().DeleteById, 200)) + + // Create test request + req, err := http.NewRequest("DELETE", "/schedules/"+schedule.Id.Hex(), nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + // Verify schedule is deleted from database + scheduleSvc := service.NewModelService[models.Schedule]() + _, err = scheduleSvc.GetById(schedule.Id) + assert.Error(t, err) // Should return error as the schedule is deleted +} + +// Test GetScheduleList endpoint +func TestGetScheduleList(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create several test schedules + schedule1, _ := createTestSchedule(t) + schedule2, _ := createTestSchedule(t) + schedule2.Name = "Second Test Schedule" + + scheduleSvc := service.NewModelService[models.Schedule]() + err := scheduleSvc.ReplaceById(schedule2.Id, *schedule2) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.GET("/schedules", nil, tonic.Handler(controllers.NewController[models.Schedule]().GetList, 200)) + + // Create test request + req, err := http.NewRequest("GET", "/schedules?page=1&size=10", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.ListResponse[models.Schedule] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.Equal(t, "ok", response.Status) + assert.Equal(t, 2, response.Total) // We created 2 schedules + assert.Equal(t, 2, len(response.Data)) + + // Verify both schedules are in the response + foundSchedule1 := false + foundSchedule2 := false + for _, schedule := range response.Data { + if schedule.Id == schedule1.Id { + foundSchedule1 = true + assert.Equal(t, schedule1.Name, schedule.Name) + } + if schedule.Id == schedule2.Id { + foundSchedule2 = true + assert.Equal(t, "Second Test Schedule", schedule.Name) + } + } + assert.True(t, foundSchedule1, "schedule1 should be in the response") + assert.True(t, foundSchedule2, "schedule2 should be in the response") +} + +// Test EnableSchedule endpoint +func TestEnableSchedule(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test schedule (disabled by default) + schedule, _ := createTestSchedule(t) + assert.False(t, schedule.Enabled) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/schedules/:id/enable", nil, tonic.Handler(controllers.PostScheduleEnable, 200)) + + // Create test request + req, err := http.NewRequest("POST", "/schedules/"+schedule.Id.Hex()+"/enable", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + // Verify schedule is now enabled + scheduleSvc := service.NewModelService[models.Schedule]() + updatedSchedule, err := scheduleSvc.GetById(schedule.Id) + require.NoError(t, err) + assert.True(t, updatedSchedule.Enabled) +} + +// Test DisableSchedule endpoint +func TestDisableSchedule(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test schedule and enable it + schedule, _ := createTestSchedule(t) + schedule.Enabled = true + + scheduleSvc := service.NewModelService[models.Schedule]() + err := scheduleSvc.ReplaceById(schedule.Id, *schedule) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/schedules/:id/disable", nil, tonic.Handler(controllers.PostScheduleDisable, 200)) + + // Create test request + req, err := http.NewRequest("POST", "/schedules/"+schedule.Id.Hex()+"/disable", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + // Verify schedule is now disabled + updatedSchedule, err := scheduleSvc.GetById(schedule.Id) + require.NoError(t, err) + assert.False(t, updatedSchedule.Enabled) +} diff --git a/core/controllers/spider_test.go b/core/controllers/spider_test.go index d4614222..fe2e3954 100644 --- a/core/controllers/spider_test.go +++ b/core/controllers/spider_test.go @@ -34,7 +34,10 @@ func TestCreateSpider(t *testing.T) { Name: "Test Spider", ColName: "test_spiders", } - jsonValue, _ := json.Marshal(payload) + requestParams := controllers.PostParams[models.Spider]{ + Data: payload, + } + jsonValue, _ := json.Marshal(requestParams) req, _ := http.NewRequest("POST", "/spiders", bytes.NewBuffer(jsonValue)) req.Header.Set("Authorization", TestToken) resp := httptest.NewRecorder() diff --git a/core/controllers/task_test.go b/core/controllers/task_test.go new file mode 100644 index 00000000..363f4d14 --- /dev/null +++ b/core/controllers/task_test.go @@ -0,0 +1,343 @@ +package controllers_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/crawlab-team/crawlab/core/constants" + "github.com/crawlab-team/crawlab/core/controllers" + "github.com/crawlab-team/crawlab/core/middlewares" + "github.com/crawlab-team/crawlab/core/models/models" + "github.com/crawlab-team/crawlab/core/models/service" + "github.com/gin-gonic/gin" + "github.com/loopfz/gadgeto/tonic" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/bson/primitive" +) + +// Helper function to create a test task +func createTestTask(t *testing.T) (task *models.Task, spiderId primitive.ObjectID) { + // First, create a test spider + spider := models.Spider{ + Name: "Test Spider for Task", + ColName: "test_spider_for_task", + } + spiderSvc := service.NewModelService[models.Spider]() + var err error + spiderId, err = spiderSvc.InsertOne(spider) + require.NoError(t, err) + require.False(t, spiderId.IsZero()) + + // Now create a task associated with the spider + task = &models.Task{ + SpiderId: spiderId, + Status: constants.TaskStatusPending, + Priority: 10, + Mode: constants.RunTypeAllNodes, + Param: "test param", + Cmd: "python main.py", + UserId: TestUserId, + } + + // Set timestamps + now := time.Now() + task.CreatedAt = now + task.UpdatedAt = now + + taskSvc := service.NewModelService[models.Task]() + taskId, err := taskSvc.InsertOne(*task) + require.NoError(t, err) + require.False(t, taskId.IsZero()) + + task.Id = taskId + return task, spiderId +} + +// Test GetTaskById endpoint +func TestGetTaskById(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test task + task, _ := createTestTask(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.GET("/tasks/:id", nil, tonic.Handler(controllers.GetTaskById, 200)) + + // Create test request + req, err := http.NewRequest("GET", "/tasks/"+task.Id.Hex(), nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.Response[models.Task] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.True(t, response.Status == "ok") + assert.Equal(t, task.Id, response.Data.Id) + assert.Equal(t, task.SpiderId, response.Data.SpiderId) + assert.Equal(t, task.Status, response.Data.Status) +} + +// Test GetTaskList endpoint +func TestGetTaskList(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create several test tasks + task1, _ := createTestTask(t) + task2, _ := createTestTask(t) + task2.Status = constants.TaskStatusRunning + + // Use ReplaceById instead of UpdateById with the model + taskSvc := service.NewModelService[models.Task]() + err := taskSvc.ReplaceById(task2.Id, *task2) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.GET("/tasks", nil, tonic.Handler(controllers.GetTaskList, 200)) + + // Create test request + req, err := http.NewRequest("GET", "/tasks?page=1&size=10", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.ListResponse[models.Task] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + assert.True(t, response.Status == "ok") + assert.Equal(t, 2, response.Total) // We created 2 tasks + assert.Equal(t, 2, len(response.Data)) + + // Verify both tasks (including task1) are in the response + foundTask1 := false + foundTask2 := false + for _, task := range response.Data { + if task.Id == task1.Id { + foundTask1 = true + assert.Equal(t, constants.TaskStatusPending, task.Status) + } + if task.Id == task2.Id { + foundTask2 = true + assert.Equal(t, constants.TaskStatusRunning, task.Status) + } + } + assert.True(t, foundTask1, "task1 should be in the response") + assert.True(t, foundTask2, "task2 should be in the response") +} + +// Test DeleteTaskById endpoint +func TestDeleteTaskById(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test task + task, _ := createTestTask(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.DELETE("/tasks/:id", nil, tonic.Handler(controllers.DeleteTaskById, 200)) + + // Create test request + req, err := http.NewRequest("DELETE", "/tasks/"+task.Id.Hex(), nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + // Verify task is deleted from database + taskSvc := service.NewModelService[models.Task]() + _, err = taskSvc.GetById(task.Id) + assert.Error(t, err) // Should return error as the task is deleted +} + +// Test PostTaskRun endpoint +func TestPostTaskRun(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test spider + spider := models.Spider{ + Name: "Test Spider for Run", + ColName: "test_spider_for_run", + } + spiderSvc := service.NewModelService[models.Spider]() + spiderId, err := spiderSvc.InsertOne(spider) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/tasks/run", nil, tonic.Handler(controllers.PostTaskRun, 200)) + + // Create payload + payload := controllers.PostTaskRunParams{ + SpiderId: spiderId.Hex(), + Mode: constants.RunTypeAllNodes, + Cmd: "python main.py", + Param: "test param", + Priority: 1, + } + jsonValue, _ := json.Marshal(payload) + + // Create test request + req, err := http.NewRequest("POST", "/tasks/run", bytes.NewBuffer(jsonValue)) + req.Header.Set("Authorization", TestToken) + req.Header.Set("Content-Type", "application/json") + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response - it may fail if the scheduler service is not properly initialized in test environment + // This is more of an integration test, so we'll check the status code but not the exact response + assert.Equal(t, http.StatusOK, resp.Code) +} + +// Test PostTaskCancel endpoint +func TestPostTaskCancel(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test task + task, _ := createTestTask(t) + + // Set status to running to make it cancellable + task.Status = constants.TaskStatusRunning + + // Use ReplaceById instead of UpdateById with the model + taskSvc := service.NewModelService[models.Task]() + err := taskSvc.ReplaceById(task.Id, *task) + require.NoError(t, err) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/tasks/:id/cancel", nil, tonic.Handler(controllers.PostTaskCancel, 200)) + + // Create payload + payload := controllers.PostTaskCancelParams{ + Force: true, + } + jsonValue, _ := json.Marshal(payload) + + // Create test request + req, err := http.NewRequest("POST", "/tasks/"+task.Id.Hex()+"/cancel", bytes.NewBuffer(jsonValue)) + req.Header.Set("Authorization", TestToken) + req.Header.Set("Content-Type", "application/json") + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response - it may fail if the scheduler service is not properly initialized in test environment + // This is more of an integration test, so we'll check the status code but not the exact response + assert.Equal(t, http.StatusOK, resp.Code) +} + +// Test PostTaskRestart endpoint +func TestPostTaskRestart(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test task + task, _ := createTestTask(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.POST("/tasks/:id/restart", nil, tonic.Handler(controllers.PostTaskRestart, 200)) + + // Create test request + req, err := http.NewRequest("POST", "/tasks/"+task.Id.Hex()+"/restart", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response - it may fail if the scheduler service is not properly initialized in test environment + // This is more of an integration test, so we'll check the status code but not the exact response + assert.Equal(t, http.StatusOK, resp.Code) +} + +// Test GetTaskLogs endpoint +func TestGetTaskLogs(t *testing.T) { + SetupTestDB() + defer CleanupTestDB() + + gin.SetMode(gin.TestMode) + + // Create a test task + task, _ := createTestTask(t) + + // Set up router + router := SetupRouter() + router.Use(middlewares.AuthorizationMiddleware()) + router.GET("/tasks/:id/logs", nil, tonic.Handler(controllers.GetTaskLogs, 200)) + + // Create test request + req, err := http.NewRequest("GET", "/tasks/"+task.Id.Hex()+"/logs?page=1&size=100", nil) + req.Header.Set("Authorization", TestToken) + require.NoError(t, err) + + // Execute request + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + // Verify response + assert.Equal(t, http.StatusOK, resp.Code) + + var response controllers.ListResponse[string] + err = json.Unmarshal(resp.Body.Bytes(), &response) + require.NoError(t, err) + + // Check status is ok - the logs might be empty since we didn't create any, + // but the endpoint should still function correctly + assert.Equal(t, "ok", response.Status) +} diff --git a/core/controllers/utils_test.go b/core/controllers/utils_test.go new file mode 100644 index 00000000..4ace445a --- /dev/null +++ b/core/controllers/utils_test.go @@ -0,0 +1,178 @@ +package controllers_test + +import ( + "testing" + + "github.com/crawlab-team/crawlab/core/controllers" + "github.com/crawlab-team/crawlab/core/models/models" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/bson" + "go.mongodb.org/mongo-driver/bson/primitive" +) + +func TestGetFilterFromConditionString(t *testing.T) { + // Simple condition with string value + condStr := `[{"key":"name","op":"eq","value":"test"}]` + filter, err := controllers.GetFilterFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, filter) + require.Len(t, filter.Conditions, 1) + assert.Equal(t, "name", filter.Conditions[0].Key) + assert.Equal(t, "eq", filter.Conditions[0].Op) + assert.Equal(t, "test", filter.Conditions[0].Value) + + // Multiple conditions with different types + condStr = `[{"key":"name","op":"eq","value":"test"},{"key":"priority","op":"gt","value":5}]` + filter, err = controllers.GetFilterFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, filter) + require.Len(t, filter.Conditions, 2) + assert.Equal(t, "name", filter.Conditions[0].Key) + assert.Equal(t, "eq", filter.Conditions[0].Op) + assert.Equal(t, "test", filter.Conditions[0].Value) + assert.Equal(t, "priority", filter.Conditions[1].Key) + assert.Equal(t, "gt", filter.Conditions[1].Op) + assert.Equal(t, float64(5), filter.Conditions[1].Value) // JSON parses numbers as float64 + + // Invalid JSON should return error + condStr = `[{"key":"name","op":"eq","value":"test"` + _, err = controllers.GetFilterFromConditionString(condStr) + assert.Error(t, err) +} + +func TestGetFilterQueryFromConditionString(t *testing.T) { + // Simple equality condition + condStr := `[{"key":"name","op":"eq","value":"test"}]` + query, err := controllers.GetFilterQueryFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, query) + expected := bson.M{"name": "test"} + assert.Equal(t, expected, query) + + // Greater than condition + condStr = `[{"key":"priority","op":"gt","value":5}]` + query, err = controllers.GetFilterQueryFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, query) + expected = bson.M{"priority": bson.M{"$gt": float64(5)}} + assert.Equal(t, expected, query) + + // Multiple conditions + condStr = `[{"key":"name","op":"eq","value":"test"},{"key":"priority","op":"gt","value":5}]` + query, err = controllers.GetFilterQueryFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, query) + expected = bson.M{"name": "test", "priority": bson.M{"$gt": float64(5)}} + assert.Equal(t, expected, query) + + // Contains operator + condStr = `[{"key":"name","op":"contains","value":"test"}]` + query, err = controllers.GetFilterQueryFromConditionString(condStr) + require.NoError(t, err) + require.NotNil(t, query) + expectedRegex := bson.M{"name": bson.M{"$regex": "test", "$options": "i"}} + assert.Equal(t, expectedRegex, query) + + // Invalid condition should return error + condStr = `[{"key":"name","op":"invalid_op","value":"test"}]` + _, err = controllers.GetFilterQueryFromConditionString(condStr) + assert.Error(t, err) +} + +func TestGetFilterQueryFromListParams(t *testing.T) { + // No conditions + params := &controllers.GetListParams{} + query, err := controllers.GetFilterQueryFromListParams(params) + require.NoError(t, err) + assert.Nil(t, query) + + // With conditions + params.Conditions = `[{"key":"name","op":"eq","value":"test"}]` + query, err = controllers.GetFilterQueryFromListParams(params) + require.NoError(t, err) + require.NotNil(t, query) + expected := bson.M{"name": "test"} + assert.Equal(t, expected, query) +} + +func TestGetUserFromContext(t *testing.T) { + // Empty context should return nil + c := &gin.Context{} + user := controllers.GetUserFromContext(c) + assert.Nil(t, user) + + // Context with non-user value should return nil + c = &gin.Context{} + c.Set("user", "not a user object") + user = controllers.GetUserFromContext(c) + assert.Nil(t, user) + + // Context with user should return the user + c = &gin.Context{} + expectedUser := &models.User{Username: "test_user"} + expectedUser.Id = primitive.NewObjectID() + c.Set("user", expectedUser) + user = controllers.GetUserFromContext(c) + assert.NotNil(t, user) + assert.Equal(t, expectedUser.Id, user.Id) + assert.Equal(t, expectedUser.Username, user.Username) +} + +func TestGetErrorResponse(t *testing.T) { + // Error response test + err := assert.AnError + resp, _ := controllers.GetErrorResponse[models.Task](err) + assert.Equal(t, err.Error(), resp.Error) + assert.Equal(t, models.Task{}, resp.Data) +} + +func TestGetDataResponse(t *testing.T) { + // Data response test + task := models.Task{ + Status: "running", + Cmd: "python main.py", + Param: "test param", + } + task.Id = primitive.NewObjectID() + + resp, err := controllers.GetDataResponse(task) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Status) + assert.Equal(t, task, resp.Data) + assert.Empty(t, resp.Error) +} + +func TestGetListResponse(t *testing.T) { + // List response test + tasks := []models.Task{ + { + Status: "running", + Cmd: "python main.py", + }, + { + Status: "pending", + Cmd: "python main.py", + }, + } + tasks[0].Id = primitive.NewObjectID() + tasks[1].Id = primitive.NewObjectID() + + total := 2 + resp, err := controllers.GetListResponse(tasks, total) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Status) + assert.Equal(t, tasks, resp.Data) + assert.Equal(t, total, resp.Total) + assert.Empty(t, resp.Error) +} + +func TestGetErrorListResponse(t *testing.T) { + // Error list response test + err := assert.AnError + resp, _ := controllers.GetErrorListResponse[models.Task](err) + assert.Equal(t, err.Error(), resp.Error) + assert.Nil(t, resp.Data) + assert.Equal(t, 0, resp.Total) +} diff --git a/core/task/scheduler/service.go b/core/task/scheduler/service.go index 53c552bd..85b7d90b 100644 --- a/core/task/scheduler/service.go +++ b/core/task/scheduler/service.go @@ -245,6 +245,7 @@ func newTaskSchedulerService() *Service { interval: 5 * time.Second, svr: server.GetGrpcServer(), handlerSvc: handler.GetTaskHandlerService(), + Logger: utils.NewLogger("TaskScheduler"), } }