From c1e799d9ca3b96614e3644adb930dcbcb359f5da Mon Sep 17 00:00:00 2001
From: Hanz Vora <hvora@google.com>
Date: Tue, 21 Jan 2025 17:13:11 +0000
Subject: [PATCH] Add handling for TPC to GetRegionFromZones

---
 pkg/common/utils.go      | 11 ++++++++---
 pkg/common/utils_test.go | 10 ++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pkg/common/utils.go b/pkg/common/utils.go
index a69d30453..03659f7fb 100644
--- a/pkg/common/utils.go
+++ b/pkg/common/utils.go
@@ -180,17 +180,22 @@ func NodeIDToZoneAndName(id string) (string, string, error) {
 }
 
 func GetRegionFromZones(zones []string) (string, error) {
+	const tpcPrefix = "u"
 	regions := sets.String{}
 	if len(zones) < 1 {
 		return "", fmt.Errorf("no zones specified")
 	}
 	for _, zone := range zones {
 		// Zone expected format {locale}-{region}-{zone}
+		// TPC zone expected format is u-{}-{}-{}, e.g. u-europe-central2-a. See go/tpc-region-naming.
 		splitZone := strings.Split(zone, "-")
-		if len(splitZone) != 3 {
-			return "", fmt.Errorf("zone in unexpected format, expected: {locale}-{region}-{zone}, got: %v", zone)
+		if len(splitZone) == 3 {
+			regions.Insert(strings.Join(splitZone[0:2], "-"))
+		} else if len(splitZone) == 4 && splitZone[0] == tpcPrefix {
+			regions.Insert(strings.Join(splitZone[0:3], "-"))
+		} else {
+			return "", fmt.Errorf("zone in unexpected format, expected: {locale}-{region}-{zone} or u-{locale}-{region}-{zone}, got: %v", zone)
 		}
-		regions.Insert(strings.Join(splitZone[0:2], "-"))
 	}
 	if regions.Len() != 1 {
 		return "", fmt.Errorf("multiple or no regions gotten from zones, got: %v", regions)
diff --git a/pkg/common/utils_test.go b/pkg/common/utils_test.go
index 0df0ed1d1..59e6f7ca8 100644
--- a/pkg/common/utils_test.go
+++ b/pkg/common/utils_test.go
@@ -293,6 +293,16 @@ func TestGetRegionFromZones(t *testing.T) {
 			zones:  []string{"blah/blooh"},
 			expErr: true,
 		},
+		{
+			name:      "tpc zone",
+			zones:     []string{"u-europe-central2-a"},
+			expRegion: "u-europe-central2",
+		},
+		{
+			name:   "malformed tpc zone",
+			zones:  []string{"us-europe-central2-a"},
+			expErr: true,
+		},
 	}
 	for _, tc := range testCases {
 		t.Logf("test case: %s", tc.name)