Function TestTrialLogs

master/internal/api_trials_intg_test.go:976–1068

(t *testing.T)

}

func TestTrialLogs(t *testing.T) {
	api, curUser, ctx := setupAPITest(t, nil)
	trial, task0 := createTestTrial(t, api, curUser)

	task1 := &model.Task{
		TaskType:   model.TaskTypeTrial,
		LogVersion: model.TaskLogVersion1,
		StartTime:  task0.StartTime.Add(time.Second),
		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
	}
	require.NoError(t, db.AddTask(ctx, task1))

	task2 := &model.Task{
		TaskType:   model.TaskTypeTrial,
		LogVersion: model.TaskLogVersion1,
		StartTime:  task1.StartTime.Add(time.Second),
		TaskID:     trialTaskID(trial.ExperimentID, model.NewRequestID(rand.Reader)),
	}
	require.NoError(t, db.AddTask(ctx, task2))

	_, err := db.Bun().NewInsert().Model(&[]model.RunTaskID{
		{RunID: trial.ID, TaskID: task1.TaskID},
		{RunID: trial.ID, TaskID: task2.TaskID},
	}).Exec(ctx)
	require.NoError(t, err)

	var expected []string
	tasks := []model.TaskID{task0.TaskID, task1.TaskID, task2.TaskID}
	var taskLogs []*model.TaskLog
	for i, taskID := range tasks {
		for j := 0; j < 9; j++ {
			log := fmt.Sprintf("%d-%d\n", i, j)
			taskLogs = append(taskLogs, &model.TaskLog{TaskID: string(taskID), Log: log})
			expected = append(expected, log)
		}
	}
	require.NoError(t, api.m.db.AddTaskLogs(taskLogs))

	stream := &mockStream[*apiv1.TrialLogsResponse]{ctx: ctx}

	err = api.TrialLogs(&apiv1.TrialLogsRequest{
		TrialId: int32(trial.ID),
	}, stream)
	require.NoError(t, err)

	require.Len(t, stream.data, len(expected))
	for i, expected := range expected {
		require.Equal(t, expected, *stream.data[i].Log)
	}

	// Retry with follow.
	newStream := &mockStream[*apiv1.TrialLogsResponse]{ctx: ctx}
	done := make(chan error, 1)
	go func() {
		done <- api.TrialLogs(&apiv1.TrialLogsRequest{
			TrialId: int32(trial.ID),
			Follow:  true,
		}, newStream)