Prefer custom SVG when getting stat var node hierarchy path (#980)

datacommonsorg · Nov 9, 2022 · 7e518ba · 7e518ba
1 parent 77b63d7
commit 7e518ba
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/internal/server/v1/variable/ancestors.go b/internal/server/v1/variable/ancestors.go
@@ -16,6 +16,7 @@ package variable
 
 import (
 	"context"
+	"strings"
 
 	pb "github.com/datacommonsorg/mixer/internal/proto"
 	"github.com/datacommonsorg/mixer/internal/server/resource"
@@ -42,6 +43,13 @@ func Ancestors(
 	for {
 		if parents, ok := cache.ParentSvg[curr]; ok {
 			curr = parents[0]
+			for _, parent := range parents {
+				// Prefer parent from custom import group
+				if strings.HasPrefix(parent, "dc/g/Custom_") {
+					curr = parent
+					break
+				}
+			}
 			if curr == statvar.SvgRoot {
 				break
 			}

diff --git a/internal/server/v1/variable/golden/ancestors/custom.json b/internal/server/v1/variable/golden/ancestors/custom.json
@@ -0,0 +1,9 @@
+{
+  "ancestors": [
+    "dc/g/Custom_Atmosphere_Pollutant-CarbonMonoxide",
+    "dc/g/Custom_Atmosphere_Pollutant",
+    "dc/g/Custom_Atmosphere",
+    "dc/g/Custom_Environment",
+    "dc/g/Custom_Root"
+  ]
+}
diff --git a/internal/server/v1/variable/golden/variable_ancestors_test.go b/internal/server/v1/variable/golden/variable_ancestors_test.go
@@ -50,6 +50,10 @@ func TestVariableAncestors(t *testing.T) {
 				"Count_Person_FoodInsecure",
 				"memdb.json",
 			},
+			{
+				"Max_Concentration_AirPollutant_CO",
+				"custom.json",
+			},
 		} {
 			resp, err := mixer.VariableAncestors(ctx, &pb.VariableAncestorsRequest{
 				Node: c.node,
@@ -83,7 +87,7 @@ func TestVariableAncestors(t *testing.T) {
 
 	if err := test.TestDriver(
 		"VariableAncestors",
-		&test.TestOption{UseCache: true, UseMemdb: true},
+		&test.TestOption{UseCache: true, UseMemdb: true, UseCustomTable: true},
 		testSuite,
 	); err != nil {
 		t.Errorf("TestDriver() = %s", err)