internal/adminapi/task_retry_handler_test.go

Rama main · réplica de lectura del repositorio original.

package adminapi

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/yeomyeonggeori/blueclaw/internal/connectors"
	"github.com/yeomyeonggeori/blueclaw/internal/task"
)

type retryTaskRunFunc func(context.Context, string) (task.TaskRun, error)

func (function retryTaskRunFunc) RetryTaskRun(contextValue context.Context, taskRunID string) (task.TaskRun, error) {
	return function(contextValue, taskRunID)
}

func TestTaskMonitorHandlerRetriesVisibleFailedTaskRun(t *testing.T) {
	taskRunService := task.NewTaskRunService(task.NewTaskEventService())
	taskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "retry this")
	if _, errorValue := taskRunService.FailTaskRun(taskRun.TaskRunID, "failed"); errorValue != nil {
		t.Fatal(errorValue)
	}
	var retriedTaskRunID string
	childTaskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "retry this")
	handler := TaskMonitorHandler{
		TaskRunService: taskRunService,
		RetryTaskRun: retryTaskRunFunc(func(_ context.Context, taskRunID string) (task.TaskRun, error) {
			retriedTaskRunID = taskRunID
			return childTaskRun, nil
		}),
	}
	request := httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+taskRun.TaskRunID+`","viewerIsAdmin":true}`))
	response := httptest.NewRecorder()
	handler.HandleRetryTaskRun(response, request)
	var responseBody map[string]string
	if json.NewDecoder(response.Body).Decode(&responseBody) != nil {
		t.Fatal("retry response was not JSON")
	}
	if response.Code != http.StatusAccepted || retriedTaskRunID != taskRun.TaskRunID || responseBody["taskRunID"] != childTaskRun.TaskRunID || responseBody["status"] != string(childTaskRun.Status) {
		t.Fatalf("retry answered %d and called %q", response.Code, retriedTaskRunID)
	}
}

func TestTaskMonitorHandlerRetryVisibilityIsScopedToViewer(t *testing.T) {
	taskRunService := task.NewTaskRunService(task.NewTaskEventService())
	aliceTaskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "alice retry")
	bobTaskRun := taskRunService.CreateTaskRun("person-2", "conversation-2", "bob retry")
	for _, taskRun := range []task.TaskRun{aliceTaskRun, bobTaskRun} {
		if _, errorValue := taskRunService.FailTaskRun(taskRun.TaskRunID, "failed"); errorValue != nil {
			t.Fatal(errorValue)
		}
	}
	invocations := 0
	handler := TaskMonitorHandler{
		TaskRunService:  taskRunService,
		IdentityService: newTestIdentityService(),
		RetryTaskRun: retryTaskRunFunc(func(_ context.Context, taskRunID string) (task.TaskRun, error) {
			invocations++
			retriedTaskRun, _ := taskRunService.FindTaskRun(taskRunID)
			return retriedTaskRun, nil
		}),
	}
	otherViewerResponse := httptest.NewRecorder()
	handler.HandleRetryTaskRun(otherViewerResponse, httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+bobTaskRun.TaskRunID+`","viewerEmail":"alice@example.com"}`)))
	if otherViewerResponse.Code != http.StatusNotFound || invocations != 0 {
		t.Fatalf("cross viewer retry answered %d and invoked runtime %d times", otherViewerResponse.Code, invocations)
	}
	ownViewerResponse := httptest.NewRecorder()
	handler.HandleRetryTaskRun(ownViewerResponse, httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+aliceTaskRun.TaskRunID+`","viewerEmail":"alice@example.com"}`)))
	if ownViewerResponse.Code != http.StatusAccepted || invocations != 1 {
		t.Fatalf("own viewer retry answered %d and invoked runtime %d times", ownViewerResponse.Code, invocations)
	}
}

func TestTaskMonitorHandlerRetryFailsClosedWhenViewerIdentityIsUnavailable(t *testing.T) {
	invocations := 0
	handler := TaskMonitorHandler{
		TaskRunService: task.NewTaskRunService(task.NewTaskEventService()),
		RetryTaskRun: retryTaskRunFunc(func(_ context.Context, _ string) (task.TaskRun, error) {
			invocations++
			return task.TaskRun{}, nil
		}),
	}
	taskRun := handler.TaskRunService.CreateTaskRun("person-1", "conversation-1", "retry")
	if _, errorValue := handler.TaskRunService.FailTaskRun(taskRun.TaskRunID, "failed"); errorValue != nil {
		t.Fatal(errorValue)
	}
	response := httptest.NewRecorder()
	handler.HandleRetryTaskRun(response, httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+taskRun.TaskRunID+`","viewerEmail":"alice@example.com"}`)))
	if response.Code != http.StatusServiceUnavailable || invocations != 0 {
		t.Fatalf("unresolved viewer answered %d and invoked runtime %d times", response.Code, invocations)
	}
}

func TestTaskMonitorHandlerRetryMapsSourceDeletedRaceToNotFound(t *testing.T) {
	taskRunService := task.NewTaskRunService(task.NewTaskEventService())
	taskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "retry")
	if _, errorValue := taskRunService.FailTaskRun(taskRun.TaskRunID, "failed"); errorValue != nil {
		t.Fatal(errorValue)
	}
	handler := TaskMonitorHandler{
		TaskRunService: taskRunService,
		RetryTaskRun: retryTaskRunFunc(func(_ context.Context, _ string) (task.TaskRun, error) {
			return task.TaskRun{}, task.ErrTaskRunNotFound
		}),
	}
	response := httptest.NewRecorder()
	handler.HandleRetryTaskRun(response, httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+taskRun.TaskRunID+`","viewerIsAdmin":true}`)))
	if response.Code != http.StatusNotFound {
		t.Fatalf("deleted source race answered %d, want 404", response.Code)
	}
}

func TestTaskMonitorHandlerRejectsMissingOrNonFailedTaskRun(t *testing.T) {
	taskRunService := task.NewTaskRunService(task.NewTaskEventService())
	runningTaskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "running")
	handler := TaskMonitorHandler{TaskRunService: taskRunService, RetryTaskRun: retryTaskRunFunc(func(_ context.Context, _ string) (task.TaskRun, error) { return task.TaskRun{}, nil })}
	for _, body := range []string{`{}`, `{"taskRunID":"` + runningTaskRun.TaskRunID + `","viewerIsAdmin":true}`} {
		response := httptest.NewRecorder()
		handler.HandleRetryTaskRun(response, httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(body)))
		wantedStatus := http.StatusBadRequest
		if body != `{}` {
			wantedStatus = http.StatusConflict
		}
		if response.Code != wantedStatus {
			t.Fatalf("body %s answered %d, want %d", body, response.Code, wantedStatus)
		}
	}
}

func TestTaskMonitorHandlerMapsRetryRuntimeErrors(t *testing.T) {
	for _, testCase := range []struct {
		name       string
		errorValue error
		status     int
	}{
		{name: "conflict", errorValue: connectors.ErrTaskRetryConflict, status: http.StatusConflict},
		{name: "unavailable", errorValue: connectors.ErrTaskRetryUnavailable, status: http.StatusServiceUnavailable},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			taskRunService := task.NewTaskRunService(task.NewTaskEventService())
			taskRun := taskRunService.CreateTaskRun("person-1", "conversation-1", "retry this")
			if _, errorValue := taskRunService.FailTaskRun(taskRun.TaskRunID, "failed"); errorValue != nil {
				t.Fatal(errorValue)
			}
			handler := TaskMonitorHandler{TaskRunService: taskRunService, RetryTaskRun: retryTaskRunFunc(func(_ context.Context, _ string) (task.TaskRun, error) { return task.TaskRun{}, testCase.errorValue })}
			request := httptest.NewRequest(http.MethodPost, "/admin/api/run/retry", strings.NewReader(`{"taskRunID":"`+taskRun.TaskRunID+`","viewerIsAdmin":true}`))
			response := httptest.NewRecorder()
			handler.HandleRetryTaskRun(response, request)
			if response.Code != testCase.status {
				t.Fatalf("retry answered %d, want %d", response.Code, testCase.status)
			}
		})
	}
}