integration test sort of working

suprjinx · Jul 8, 2024 · 3cf1653 · 3cf1653
1 parent f5c3eff
commit 3cf1653
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 25 deletions.
diff --git 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
@@ -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
@@ -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
+}