test: add unit test for no attributes selection

goto · Jul 25, 2024 · 3134bfd · 3134bfd
1 parent 4a7b474
commit 3134bfd
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/core/asset/service.go b/core/asset/service.go
@@ -192,8 +192,7 @@ func (s *Service) GetLineage(ctx context.Context, urn string, query LineageQuery
 
 	if !query.WithAttributes {
 		return Lineage{
-			Edges:     edges,
-			NodeAttrs: make(map[string]NodeAttributes),
+			Edges: edges,
 		}, nil
 	}
 

diff --git a/core/asset/service_test.go b/core/asset/service_test.go
@@ -891,6 +891,28 @@ func TestService_GetLineage(t *testing.T) {
 			},
 			Err: nil,
 		},
+		{
+			Description: `should return lineage with edges and without node attributes`,
+			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: false,
+			},
+			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: false}).Return(asset.LineageGraph{
+					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
+					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
+					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
+				}, nil)
+			},
+			Expected: asset.Lineage{
+				Edges: []asset.LineageEdge{
+					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
+					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
+					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
+				},
+			},
+			Err: nil,
+		},
 		{
 			Description: `should return lineage with edges and node attributes`,
 			ID:          assetID,