Function TestTrialLogFields

master/internal/api_trials_intg_test.go:1070–1173

(t *testing.T)

}

func TestTrialLogFields(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)

	expectedContainerIDs := make(map[string]bool)
	tasks := []model.TaskID{task0.TaskID, task1.TaskID, task2.TaskID}
	var taskLogs []*model.TaskLog
	for i, taskID := range tasks {
		containerID := fmt.Sprintf("id-%d", i)
		taskLogs = append(taskLogs, &model.TaskLog{
			TaskID:      string(taskID),
			Log:         "test log",
			ContainerID: ptrs.Ptr(containerID),
		})
		expectedContainerIDs[containerID] = true
	}
	require.NoError(t, api.m.db.AddTaskLogs(taskLogs))

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

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

	actualContainerIDs := make(map[string]bool)
	for _, s := range stream.getData() {
		for _, containerID := range s.ContainerIds {
			actualContainerIDs[containerID] = true
		}
	}
	require.Equal(t, expectedContainerIDs, actualContainerIDs)

	// Retry with follow.
	newStream := &mockStream[*apiv1.TrialLogsFieldsResponse]{ctx: ctx}
	done := make(chan error, 1)