From b6048587b9b7d909a92df1447078bafe2c0fd5d5 Mon Sep 17 00:00:00 2001
From: Wagner Maciel <wagnermaciel@google.com>
Date: Tue, 8 Mar 2022 19:40:47 +0000
Subject: [PATCH] feat(material/schematics): add fn for adding attrs to
 templates (#24550)

* create and unit test addAttribute fn
---
 .../rules/tree-traversal.spec.ts              | 36 ++++++++++++++++++-
 .../mdc-migration/rules/tree-traversal.ts     | 21 +++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.spec.ts b/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.spec.ts
index 30251ddee52b..af894e78854d 100644
--- a/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.spec.ts
+++ b/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.spec.ts
@@ -1,4 +1,10 @@
-import {visitElements, parseTemplate, replaceStartTag, replaceEndTag} from './tree-traversal';
+import {
+  addAttribute,
+  visitElements,
+  parseTemplate,
+  replaceStartTag,
+  replaceEndTag,
+} from './tree-traversal';
 
 function runTagNameDuplicationTest(html: string, result: string): void {
   visitElements(
@@ -13,6 +19,13 @@ function runTagNameDuplicationTest(html: string, result: string): void {
   expect(html).toBe(result);
 }
 
+function runAddAttributeTest(html: string, result: string): void {
+  visitElements(parseTemplate(html).nodes, undefined, node => {
+    html = addAttribute(html, node, 'attr', 'val');
+  });
+  expect(html).toBe(result);
+}
+
 describe('#visitElements', () => {
   describe('tag name replacements', () => {
     it('should handle basic cases', async () => {
@@ -76,4 +89,25 @@ describe('#visitElements', () => {
       );
     });
   });
+
+  describe('add attribute tests', () => {
+    it('should handle single element', async () => {
+      runAddAttributeTest('<a></a>', '<a attr="val"></a>');
+    });
+
+    it('should handle multiple unnested', async () => {
+      runAddAttributeTest('<a></a><b></b>', '<a attr="val"></a><b attr="val"></b>');
+    });
+
+    it('should handle multiple nested', async () => {
+      runAddAttributeTest('<a><b></b></a>', '<a attr="val"><b attr="val"></b></a>');
+    });
+
+    it('should handle multiple nested and unnested', async () => {
+      runAddAttributeTest(
+        '<a><b></b><c></c></a>',
+        '<a attr="val"><b attr="val"></b><c attr="val"></c></a>',
+      );
+    });
+  });
 });
diff --git a/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.ts b/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.ts
index 00cb9c3f38fe..8dc98a2a828d 100644
--- a/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.ts
+++ b/src/material/schematics/ng-generate/mdc-migration/rules/tree-traversal.ts
@@ -80,6 +80,27 @@ export function replaceEndTag(html: string, node: compiler.TmplAstElement, tag:
   return replaceAt(html, node.endSourceSpan.start.offset + 2, node.name, tag);
 }
 
+/**
+ * Appends an attribute to the given node of the template html.
+ *
+ * @param html The template html to be updated.
+ * @param node The node to be updated.
+ * @param name The name of the attribute.
+ * @param value The value of the attribute.
+ * @returns The updated template html.
+ */
+export function addAttribute(
+  html: string,
+  node: compiler.TmplAstElement,
+  name: string,
+  value: string,
+): string {
+  const index = node.startSourceSpan.start.offset + node.name.length + 1;
+  const prefix = html.slice(0, index);
+  const suffix = html.slice(index);
+  return prefix + ` ${name}="${value}"` + suffix;
+}
+
 /**
  * Replaces a substring of a given string starting at some offset index.
  *