Implement with_tasks option to expand the task information (#3343)

* Implement with_tasks option to expand the task information

* Added tests

* format

Author: chkeita

src/ApiService/ApiService/Functions/Jobs.cs

@@ -135,10 +135,14 @@ private async Task<HttpResponseData> Get(HttpRequestData req) {
 
             static JobTaskInfo TaskToJobTaskInfo(Task t) => new(t.TaskId, t.Config.Task.Type, t.State);
 
-            // TODO: search.WithTasks is not checked in Python code?
-
-            var taskInfo = await _context.TaskOperations.SearchStates(jobId).Select(TaskToJobTaskInfo).ToListAsync();
-            return await RequestHandling.Ok(req, JobResponse.ForJob(job, taskInfo));
+            var tasks = _context.TaskOperations.SearchStates(jobId);
+            if (search.WithTasks ?? false) {
+                var ts = await tasks.ToListAsync();
+                return await RequestHandling.Ok(req, JobResponse.ForJob(job, ts));
+            } else {
+                var taskInfo = await tasks.Select(TaskToJobTaskInfo).ToListAsync();
+                return await RequestHandling.Ok(req, JobResponse.ForJob(job, taskInfo));
+            }
         }
 
         var jobs = await _context.JobOperations.SearchState(states: search.State ?? Enumerable.Empty<JobState>()).ToListAsync();

src/ApiService/ApiService/OneFuzzTypes/Model.cs

@@ -282,7 +282,8 @@ public record Task(
     ISecret<Authentication>? Auth = null,
     DateTimeOffset? Heartbeat = null,
     DateTimeOffset? EndTime = null,
-    StoredUserInfo? UserInfo = null) : StatefulEntityBase<TaskState>(State) {
+    StoredUserInfo? UserInfo = null) : StatefulEntityBase<TaskState>(State), IJobTaskInfo {
+    public TaskType Type => Config.Task.Type;
 }
 
 public record TaskEvent(
@@ -898,11 +899,19 @@ public JobConfig Truncate(int maxLength) {
     }
 }
 
+[JsonDerivedType(typeof(Task), typeDiscriminator: "Task")]
+[JsonDerivedType(typeof(JobTaskInfo), typeDiscriminator: "JobTaskInfo")]
+public interface IJobTaskInfo {
+    Guid TaskId { get; }
+    TaskType Type { get; }
+    TaskState State { get; }
+}
+
 public record JobTaskInfo(
     Guid TaskId,
     TaskType Type,
     TaskState State
-);
+) : IJobTaskInfo;
 
 public record Job(
     [PartitionKey][RowKey] Guid JobId,

src/ApiService/ApiService/OneFuzzTypes/Responses.cs

@@ -96,13 +96,13 @@ public record JobResponse(
     JobConfig Config,
     string? Error,
     DateTimeOffset? EndTime,
-    List<JobTaskInfo>? TaskInfo,
+    IEnumerable<IJobTaskInfo>? TaskInfo,
     StoredUserInfo? UserInfo,
     [property: JsonPropertyName("Timestamp")] // must retain capital T for backcompat
     DateTimeOffset? Timestamp
 // not including UserInfo from Job model
 ) : BaseResponse() {
-    public static JobResponse ForJob(Job j, List<JobTaskInfo>? taskInfo)
+    public static JobResponse ForJob(Job j, IEnumerable<IJobTaskInfo>? taskInfo)
         => new(
             JobId: j.JobId,
             State: j.State,

src/ApiService/IntegrationTests/JobsTests.cs

@@ -175,4 +175,55 @@ public async Async.Task Post_CreatesJob_AndContainer() {
         var metadata = Assert.Single(container.Value);
         Assert.Equal(new KeyValuePair<string, string>("container_type", "logs"), metadata);
     }
+
+
+    [Fact]
+    public async Async.Task Get_CanFindSpecificJobWithTaskInfo() {
+
+        var taskConfig = new TaskConfig(_jobId, new List<Guid>(), new TaskDetails(TaskType.Coverage, 60));
+        var task = new Task(_jobId, Guid.NewGuid(), TaskState.Running, Os.Windows, taskConfig);
+
+        await Context.InsertAll(
+            new Job(_jobId, JobState.Stopped, _config, null), task);
+
+        var func = new Jobs(Context, LoggerProvider.CreateLogger<Jobs>());
+
+        var ctx = new TestFunctionContext();
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", new JobSearch(JobId: _jobId, WithTasks: false)), ctx);
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        var response = BodyAs<JobResponse>(result);
+        Assert.Equal(_jobId, response.JobId);
+        Assert.NotNull(response.TaskInfo);
+        var returnedTasks = response.TaskInfo.OfType<JobTaskInfo>().ToList();
+        Assert.NotEmpty(returnedTasks);
+        Assert.Equal(task.TaskId, returnedTasks[0].TaskId);
+        Assert.Equal(task.State, returnedTasks[0].State);
+        Assert.Equal(task.Config.Task.Type, returnedTasks[0].Type);
+    }
+
+    [Fact]
+    public async Async.Task Get_CanFindSpecificJobWithFullTask() {
+        var taskConfig = new TaskConfig(_jobId, new List<Guid>(), new TaskDetails(TaskType.Coverage, 60));
+        var task = new Task(_jobId, Guid.NewGuid(), TaskState.Running, Os.Windows, taskConfig);
+
+        await Context.InsertAll(
+            new Job(_jobId, JobState.Stopped, _config, null), task);
+
+        var func = new Jobs(Context, LoggerProvider.CreateLogger<Jobs>());
+
+        var ctx = new TestFunctionContext();
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", new JobSearch(JobId: _jobId, WithTasks: true)), ctx);
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        var response = BodyAs<JobResponse>(result);
+        Assert.Equal(_jobId, response.JobId);
+        Assert.NotNull(response.TaskInfo);
+        var returnedTasks = response.TaskInfo.OfType<Task>().ToList();
+        Assert.NotEmpty(returnedTasks);
+        Assert.Equal(task.TaskId, returnedTasks[0].TaskId);
+        Assert.Equal(task.State, returnedTasks[0].State);
+        Assert.Equal(task.Config.Task.Type, returnedTasks[0].Type);
+
+    }
 }

src/cli/onefuzz/api.py

@@ -1308,14 +1308,16 @@ def delete(self, job_id: UUID_EXPANSION) -> models.Job:
             "DELETE", models.Job, data=requests.JobGet(job_id=job_id_expanded)
         )
 
-    def get(self, job_id: UUID_EXPANSION) -> models.Job:
+    def get(self, job_id: UUID_EXPANSION, with_tasks: bool = False) -> models.Job:
         """Get information about a specific job"""
         job_id_expanded = self._disambiguate_uuid(
             "job_id", job_id, lambda: [str(x.job_id) for x in self.list()]
         )
         self.logger.debug("get job: %s", job_id_expanded)
         job = self._req_model(
-            "GET", models.Job, data=requests.JobGet(job_id=job_id_expanded)
+            "GET",
+            models.Job,
+            data=requests.JobGet(job_id=job_id_expanded, with_tasks=with_tasks),
         )
         return job
 

src/pytypes/onefuzztypes/models.py

@@ -463,17 +463,6 @@ class JobTaskInfo(BaseModel):
     state: TaskState
 
 
-class Job(BaseModel):
-    timestamp: Optional[datetime] = Field(alias="Timestamp")
-    job_id: UUID = Field(default_factory=uuid4)
-    state: JobState = Field(default=JobState.init)
-    config: JobConfig
-    error: Optional[str]
-    end_time: Optional[datetime] = None
-    task_info: Optional[List[JobTaskInfo]]
-    user_info: Optional[UserInfo]
-
-
 class TaskHeartbeatEntry(BaseModel):
     task_id: UUID
     job_id: Optional[UUID]
@@ -757,6 +746,17 @@ class Task(BaseModel):
     user_info: Optional[UserInfo]
 
 
+class Job(BaseModel):
+    timestamp: Optional[datetime] = Field(alias="Timestamp")
+    job_id: UUID = Field(default_factory=uuid4)
+    state: JobState = Field(default=JobState.init)
+    config: JobConfig
+    error: Optional[str]
+    end_time: Optional[datetime] = None
+    task_info: Optional[List[Union[Task, JobTaskInfo]]]
+    user_info: Optional[UserInfo]
+
+
 class NetworkConfig(BaseModel):
     address_space: str = Field(default="10.0.0.0/8")
     subnet: str = Field(default="10.0.0.0/16")

src/pytypes/onefuzztypes/requests.py

@@ -39,13 +39,13 @@ class BaseRequest(BaseModel):
 
 class JobGet(BaseRequest):
     job_id: UUID
+    with_tasks: Optional[bool]
 
 
 class JobSearch(BaseRequest):
     job_id: Optional[UUID]
     state: Optional[List[JobState]]
     task_state: Optional[List[TaskState]]
-    with_tasks: Optional[bool]
 
 
 class NotificationCreate(BaseRequest, NotificationConfig):