From 3134bfd656ff192765d49fa00c4fe3af95d5c255 Mon Sep 17 00:00:00 2001
From: Anwar Hidayat <15167551+irainia@users.noreply.github.com>
Date: Thu, 25 Jul 2024 17:43:36 +0700
Subject: [PATCH] test: add unit test for no attributes selection

---
 core/asset/service.go      |  3 +--
 core/asset/service_test.go | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/core/asset/service.go b/core/asset/service.go
index c4052a0c..a12ac61c 100644
--- 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
index 16ca2d7f..db01ed0d 100644
--- 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,