Function TestPostWorkspaceAgentTaskSnapshot

coderd/aitasks_test.go:2201–2467

(t *testing.T)

}

func TestPostWorkspaceAgentTaskSnapshot(t *testing.T) {
	t.Parallel()

	// Shared coderd with mock clock for all tests.
	clock := quartz.NewMock(t)
	ownerClient, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
		Clock: clock,
	})
	owner := coderdtest.CreateFirstUser(t, ownerClient)

	createTaskWorkspace := func(t *testing.T, agentToken string) (taskID uuid.UUID, workspaceID uuid.UUID) {
		t.Helper()
		workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
			OrganizationID: owner.OrganizationID,
			OwnerID:        owner.UserID,
		}).WithTask(database.TaskTable{
			Prompt: "test prompt",
		}, &proto.App{
			Slug: "task-app",
			Url:  "http://localhost:8080",
		}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
			agents[0].Auth = &proto.Agent_Token{Token: agentToken}
			return agents
		}).Do()
		return workspaceBuild.Task.ID, workspaceBuild.Workspace.ID
	}

	makePayload := func(t *testing.T, content string) []byte {
		t.Helper()
		data := agentapisdk.GetMessagesResponse{
			Messages: []agentapisdk.Message{
				{Id: 0, Role: "agent", Content: content, Time: time.Now()},
			},
		}
		b, err := json.Marshal(data)
		require.NoError(t, err)
		return b
	}

	makeRequest := func(t *testing.T, taskID uuid.UUID, agentToken string, payload []byte, format string) *http.Response {
		t.Helper()
		ctx := testutil.Context(t, testutil.WaitShort)

		url := ownerClient.URL.JoinPath("/api/v2/workspaceagents/me/tasks", taskID.String(), "log-snapshot").String()
		if format != "" {
			url += "?format=" + format
		}

		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
		require.NoError(t, err)
		req.Header.Set(codersdk.SessionTokenHeader, agentToken)
		res, err := http.DefaultClient.Do(req)
		require.NoError(t, err)
		return res
	}

	unmarshalSnapshot := func(t *testing.T, snapshotJSON json.RawMessage) agentapisdk.GetMessagesResponse {
		t.Helper()