From a96ce3223ab7a4cfd8f7a4fe2f46cb6aa665b8ea Mon Sep 17 00:00:00 2001
From: Aditya Choudhary <caditya@google.com>
Date: Wed, 15 Nov 2023 11:02:37 +0000
Subject: [PATCH] Refactor metadata tool to support metadata generation for
 different rules.

Bug: 296873595
Test: Manual test (use go test inside tools/metadata/testdata)

Change-Id: I881fd76213ec78001f9e12ed2fbc860d1503a364
---
 tools/metadata/OWNERS                    |  4 ++
 tools/metadata/generator.go              | 56 ++++++++++++++----------
 tools/metadata/testdata/metadata_test.go |  4 +-
 3 files changed, 38 insertions(+), 26 deletions(-)
 create mode 100644 tools/metadata/OWNERS

diff --git a/tools/metadata/OWNERS b/tools/metadata/OWNERS
new file mode 100644
index 0000000000..03bcdf1c40
--- /dev/null
+++ b/tools/metadata/OWNERS
@@ -0,0 +1,4 @@
+dariofreni@google.com
+joeo@google.com
+ronish@google.com
+caditya@google.com
diff --git a/tools/metadata/generator.go b/tools/metadata/generator.go
index 8e82f7fc32..eb877553d3 100644
--- a/tools/metadata/generator.go
+++ b/tools/metadata/generator.go
@@ -73,7 +73,7 @@ func readFileToString(filePath string) string {
 	return string(data)
 }
 
-func processProtobuf(
+func processTestSpecProtobuf(
 	filePath string, ownershipMetadataMap *sync.Map, keyLocks *keyToLocksMap,
 	errCh chan error, wg *sync.WaitGroup,
 ) {
@@ -130,10 +130,11 @@ func processProtobuf(
 func main() {
 	inputFile := flag.String("inputFile", "", "Input file path")
 	outputFile := flag.String("outputFile", "", "Output file path")
+	rule := flag.String("rule", "", "Metadata rule (Hint: test_spec or code_metadata)")
 	flag.Parse()
 
-	if *inputFile == "" || *outputFile == "" {
-		fmt.Println("Usage: metadata -inputFile <input file path> -outputFile <output file path>")
+	if *inputFile == "" || *outputFile == "" || *rule == "" {
+		fmt.Println("Usage: metadata -rule <rule> -inputFile <input file path> -outputFile <output file path>")
 		os.Exit(1)
 	}
 
@@ -144,26 +145,33 @@ func main() {
 	errCh := make(chan error, len(filePaths))
 	var wg sync.WaitGroup
 
-	for _, filePath := range filePaths {
-		wg.Add(1)
-		go processProtobuf(filePath, ownershipMetadataMap, keyLocks, errCh, &wg)
+	switch *rule {
+	case "test_spec":
+		for _, filePath := range filePaths {
+			wg.Add(1)
+			go processTestSpecProtobuf(filePath, ownershipMetadataMap, keyLocks, errCh, &wg)
+		}
+
+		wg.Wait()
+		close(errCh)
+
+		for err := range errCh {
+			log.Fatal(err)
+		}
+
+		allKeys := getSortedKeys(ownershipMetadataMap)
+		var allMetadata []*test_spec_proto.TestSpec_OwnershipMetadata
+
+		for _, key := range allKeys {
+			value, _ := ownershipMetadataMap.Load(key)
+			metadataList := value.([]*test_spec_proto.TestSpec_OwnershipMetadata)
+			allMetadata = append(allMetadata, metadataList...)
+		}
+
+		writeOutput(*outputFile, allMetadata)
+		break
+	case "code_metadata":
+	default:
+		log.Fatalf("No specific processing implemented for rule '%s'.\n", *rule)
 	}
-
-	wg.Wait()
-	close(errCh)
-
-	for err := range errCh {
-		log.Fatal(err)
-	}
-
-	allKeys := getSortedKeys(ownershipMetadataMap)
-	var allMetadata []*test_spec_proto.TestSpec_OwnershipMetadata
-
-	for _, key := range allKeys {
-		value, _ := ownershipMetadataMap.Load(key)
-		metadataList := value.([]*test_spec_proto.TestSpec_OwnershipMetadata)
-		allMetadata = append(allMetadata, metadataList...)
-	}
-
-	writeOutput(*outputFile, allMetadata)
 }
diff --git a/tools/metadata/testdata/metadata_test.go b/tools/metadata/testdata/metadata_test.go
index 0cb80c3fd2..03c4f29ab2 100644
--- a/tools/metadata/testdata/metadata_test.go
+++ b/tools/metadata/testdata/metadata_test.go
@@ -10,7 +10,7 @@ import (
 
 func TestMetadata(t *testing.T) {
 	cmd := exec.Command(
-		"metadata", "-inputFile", "./inputFiles.txt", "-outputFile",
+		"metadata", "-rule", "test_spec", "-inputFile", "./inputFiles.txt", "-outputFile",
 		"./generatedOutputFile.txt",
 	)
 	stderr, err := cmd.CombinedOutput()
@@ -40,7 +40,7 @@ func TestMetadata(t *testing.T) {
 
 func TestMetadataNegativeCase(t *testing.T) {
 	cmd := exec.Command(
-		"metadata", "-inputFile", "./inputFilesNegativeCase.txt", "-outputFile",
+		"metadata", "-rule", "test_spec", "-inputFile", "./inputFilesNegativeCase.txt", "-outputFile",
 		"./generatedOutputFileNegativeCase.txt",
 	)
 	stderr, err := cmd.CombinedOutput()