From ad54441d9ced79f0dbda235bdc0e8a45305a0608 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Fri, 5 Apr 2024 19:15:29 -0400
Subject: [PATCH] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 .../storage/integration/elasticsearch_test.go | 16 ++++++---------
 plugin/storage/integration/integration.go     | 20 ++++++++++++++-----
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/plugin/storage/integration/elasticsearch_test.go b/plugin/storage/integration/elasticsearch_test.go
index 45c2449f758..a50e572db78 100644
--- a/plugin/storage/integration/elasticsearch_test.go
+++ b/plugin/storage/integration/elasticsearch_test.go
@@ -167,20 +167,16 @@ func (s *ESStorageIntegration) initializeESFactory(t *testing.T, allTagsAsFields
 	require.NoError(t, command.ParseFlags(args))
 	f.InitFromViper(v, s.logger)
 	require.NoError(t, f.Initialize(metrics.NullFactory, s.logger))
-	t.Cleanup(func() {
-		require.NoError(t, f.Close())
-	})
+
+	// TODO ideally we need to close the factory once the test is finished
+	// but because esCleanup calls initialize() we get a panic later
+	// t.Cleanup(func() {
+	// 	require.NoError(t, f.Close())
+	// })
 	return f
 }
 
 func (s *ESStorageIntegration) initSpanstore(t *testing.T, allTagsAsFields bool) {
-	// client := s.getEsClient(t)
-	// opts := es.NewOptions(primaryNamespace, archiveNamespace)
-	// cfg := opts.Primary.Configuration
-	// cfg.IndexPrefix = indexPrefix
-	// cfg.Tags.AllAsFields = allTagsAsFields
-	// f, err := es.NewFactoryWithConfigTest(cfg, metrics.NullFactory, s.logger, client)
-	// require.NoError(t, err)
 	f := s.initializeESFactory(t, allTagsAsFields)
 	var err error
 	s.SpanWriter, err = f.CreateSpanWriter()
diff --git a/plugin/storage/integration/integration.go b/plugin/storage/integration/integration.go
index 506900e8ee3..d6edc3a8abf 100644
--- a/plugin/storage/integration/integration.go
+++ b/plugin/storage/integration/integration.go
@@ -460,12 +460,22 @@ func (s *StorageIntegration) testGetDependencies(t *testing.T) {
 
 	require.NoError(t, s.DependencyWriter.WriteDependencies(time.Now(), expected))
 	s.refresh(t)
-	actual, err := s.DependencyReader.GetDependencies(context.Background(), time.Now(), 5*time.Minute)
-	require.NoError(t, err)
-	sort.Slice(actual, func(i, j int) bool {
-		return actual[i].Parent < actual[j].Parent
+
+	var actual []model.DependencyLink
+	found := s.waitForCondition(t, func(t *testing.T) bool {
+		var err error
+		actual, err = s.DependencyReader.GetDependencies(context.Background(), time.Now(), 5*time.Minute)
+		require.NoError(t, err)
+		sort.Slice(actual, func(i, j int) bool {
+			return actual[i].Parent < actual[j].Parent
+		})
+		return assert.ObjectsAreEqualValues(expected, actual)
 	})
-	assert.EqualValues(t, expected, actual)
+
+	if !assert.True(t, found) {
+		t.Log("\t Expected:", expected)
+		t.Log("\t Actual  :", actual)
+	}
 }
 
 // === Sampling Store Integration Tests ===