diff --git a/pkg/api/aim/api/response/run.go b/pkg/api/aim/api/response/run.go index 6e577c1e4..130236a12 100644 --- a/pkg/api/aim/api/response/run.go +++ b/pkg/api/aim/api/response/run.go @@ -489,6 +489,7 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, totalRuns int64, var ( runID string runData fiber.Map + traces []fiber.Map cur int64 ) reportProgress := func() error { @@ -506,7 +507,7 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, totalRuns int64, } addImage := func(img models.Artifact) { if runData == nil { - imagesPerStep := result.StepImageCount(img.RunID, int(img.Step)) + imagesPerStep := result.StepImageCount(img.RunID, 0) runData = fiber.Map{ "ranges": fiber.Map{ "record_range_total": []int{0, result.TotalSteps(img.RunID)}, @@ -517,15 +518,16 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, totalRuns int64, "params": fiber.Map{ "images_per_step": imagesPerStep, }, - "traces": []fiber.Map{}, } + traces = []fiber.Map{} } - runData["traces"] = append(runData["traces"].([]fiber.Map), fiber.Map{}) + traces = append(traces, fiber.Map{}) } flushImages := func() error { if runID == "" { return nil } + runData["traces"] = traces if err := encoding.EncodeTree(w, fiber.Map{ runID: runData, }); err != nil { @@ -550,6 +552,7 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, totalRuns int64, } runID = image.RunID runData = nil + traces = nil } addImage(image) diff --git a/tests/integration/golang/aim/run/search_artifacts_test.go b/tests/integration/golang/aim/run/search_artifacts_test.go index d190b30cc..6ba48752a 100644 --- a/tests/integration/golang/aim/run/search_artifacts_test.go +++ b/tests/integration/golang/aim/run/search_artifacts_test.go @@ -54,24 +54,16 @@ func (s *SearchArtifactsTestSuite) Test_Ok() { LifecycleStage: models.LifecycleStageActive, }) s.Require().Nil(err) - _, err = s.MetricFixtures.CreateMetric(context.Background(), &models.Metric{ - Key: "TestMetric1", - Value: 1.1, - Timestamp: 123456789, - Step: 1, - IsNan: false, - RunID: run1.ID, - Iter: 1, - }) - s.Require().Nil(err) - metric1Run1, err := s.MetricFixtures.CreateLatestMetric(context.Background(), &models.LatestMetric{ - Key: "TestMetric1", - Value: 1.1, - Timestamp: 123456789, - Step: 1, - IsNan: false, - RunID: run1.ID, - LastIter: 1, + _, err = s.ArtifactFixtures.CreateArtifact(context.Background(), &models.Artifact{ + RunID: run1.ID, + BlobURI: "path/filename.png", + Step: 1, + Iter: 1, + Index: 1, + Caption: "caption1", + Format: "png", + Width: 100, + Height: 100, }) s.Require().Nil(err) @@ -84,9 +76,7 @@ func (s *SearchArtifactsTestSuite) Test_Ok() { { name: "SearchArtifact", request: request.SearchArtifactsRequest{}, - metrics: []*models.LatestMetric{ - metric1Run1, - }, + metrics: []*models.LatestMetric{}, }, } for _, tt := range tests { @@ -101,7 +91,7 @@ func (s *SearchArtifactsTestSuite) Test_Ok() { helpers.ResponseTypeBuffer, ).WithResponse( resp, - ).DoRequest("/runs/search/images"), + ).DoRequest("/runs/search/image"), ) decodedData, err := encoding.NewDecoder(resp).Decode() diff --git a/tests/integration/golang/fixtures/artifact.go b/tests/integration/golang/fixtures/artifact.go index 66296a1fe..4b802b0e9 100644 --- a/tests/integration/golang/fixtures/artifact.go +++ b/tests/integration/golang/fixtures/artifact.go @@ -31,3 +31,11 @@ func (f ArtifactFixtures) GetArtifactByRunID(ctx context.Context, runID string) } return &artifact, nil } + +// CreateArtifact creates new test Artifact. +func (f ArtifactFixtures) CreateArtifact(ctx context.Context, artifact *models.Artifact) (*models.Artifact, error) { + if err := f.baseFixtures.db.WithContext(ctx).Create(artifact).Error; err != nil { + return nil, eris.Wrap(err, "error creating artifact") + } + return artifact, nil +}