From c1d27147c4fb7541805dffd2029e2c2f0433fb25 Mon Sep 17 00:00:00 2001
From: David Thompson <davthomp@redhat.com>
Date: Thu, 8 Dec 2022 12:04:36 -0500
Subject: [PATCH] Fix surround tag in empty document

- Fill in grammar when invoking surround tag in an empty document
- Fix NPE that prevented surround tags from working
  in empty documents that are linked to a grammar with
  file association

Closes #1395

Signed-off-by: David Thompson <davthomp@redhat.com>
---
 .../commands/SurroundWithCommand.java         | 10 +++++-
 .../java/org/eclipse/lemminx/XMLAssert.java   | 11 ++++--
 .../commands/SurroundWithTagsCommandTest.java | 36 +++++++++++++++++++
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/org.eclipse.lemminx/src/main/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithCommand.java b/org.eclipse.lemminx/src/main/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithCommand.java
index 1a9b7d5f8..cf080a397 100644
--- a/org.eclipse.lemminx/src/main/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithCommand.java
+++ b/org.eclipse.lemminx/src/main/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithCommand.java
@@ -35,7 +35,7 @@
 
 /**
  * XML Command "xml.refactor.surround.with" to support surround:
- * 
+ *
  * <ul>
  * <li>Surround with Tags (Wrap)</li>
  * <li>Surround with Comments</li>
@@ -180,6 +180,14 @@ private List<String> getTags(DOMNode node, String prefix, int offset) {
 		DOMElement parentElement = node.isElement() ? (DOMElement) node : node.getParentElement();
 		Collection<CMDocument> cmDocuments = parentElement != null ? contentModelManager.findCMDocument(parentElement)
 				: contentModelManager.findCMDocument(node.getOwnerDocument(), null);
+		if (parentElement == null) {
+			return cmDocuments.stream() //
+					.flatMap(cmDocument -> cmDocument.getElements().stream()) //
+					.map(decl -> decl.getName(prefix)) //
+					.distinct() //
+					.sorted() //
+					.collect(Collectors.toList());
+		}
 		for (CMDocument cmDocument : cmDocuments) {
 			CMElementDeclaration elementDeclaration = cmDocument.findCMElement(parentElement);
 			if (elementDeclaration != null) {
diff --git a/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/XMLAssert.java b/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/XMLAssert.java
index 815f76d72..6204c0ed8 100644
--- a/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/XMLAssert.java
+++ b/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/XMLAssert.java
@@ -1811,11 +1811,18 @@ public static SelectionRange sr(List<Range> ranges) {
 		selectionRange.setParent(sr(ranges.subList(1, ranges.size())));
 		return selectionRange;
 	}
-	
+
 	public static void assertSurroundWith(String xml, SurroundWithKind kind, boolean snippetsSupported,
+	String expected) throws BadLocationException, InterruptedException, ExecutionException  {
+		assertSurroundWith(xml, kind, snippetsSupported, (service) -> {}, "src/test/resources/test.xml", expected);
+	}
+
+	public static void assertSurroundWith(String xml, SurroundWithKind kind, boolean snippetsSupported, Consumer<XMLLanguageService> configuration, String uri,
 			String expected) throws BadLocationException, InterruptedException, ExecutionException {
 		MockXMLLanguageServer languageServer = new MockXMLLanguageServer();
 
+		configuration.accept(languageServer.getXMLLanguageService());
+
 		int rangeStart = xml.indexOf('|');
 		int rangeEnd = xml.lastIndexOf('|');
 		// remove '|'
@@ -1830,7 +1837,7 @@ public static void assertSurroundWith(String xml, SurroundWithKind kind, boolean
 		Position endPos = rangeStart == rangeEnd ? startPos : document.positionAt(rangeEnd - 1);
 		Range selection = new Range(startPos, endPos);
 
-		TextDocumentIdentifier xmlIdentifier = languageServer.didOpen("src/test/resources/test.xml", x.toString());
+		TextDocumentIdentifier xmlIdentifier = languageServer.didOpen(uri, x.toString());
 
 		// Execute surround with tags command
 		SurroundWithResponse response = (SurroundWithResponse) languageServer
diff --git a/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithTagsCommandTest.java b/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithTagsCommandTest.java
index b80c55731..bf0b53b5d 100644
--- a/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithTagsCommandTest.java
+++ b/org.eclipse.lemminx/src/test/java/org/eclipse/lemminx/extensions/contentmodel/commands/SurroundWithTagsCommandTest.java
@@ -13,8 +13,13 @@
 
 import static org.eclipse.lemminx.XMLAssert.assertSurroundWith;
 
+import java.util.function.Consumer;
+
 import org.eclipse.lemminx.extensions.contentmodel.BaseFileTempTest;
 import org.eclipse.lemminx.extensions.contentmodel.commands.SurroundWithCommand.SurroundWithKind;
+import org.eclipse.lemminx.extensions.contentmodel.model.ContentModelManager;
+import org.eclipse.lemminx.extensions.contentmodel.settings.XMLFileAssociation;
+import org.eclipse.lemminx.services.XMLLanguageService;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -145,4 +150,35 @@ public void surroundEmptySelectionInEmptyText() throws Exception {
 		assertSurroundWith(xml, SurroundWithKind.tags, true, expected);
 	}
 
+	@Test
+	public void surroundEmptySelectionInEmptyDocumentWithFileAssociation() throws Exception {
+		Consumer<XMLLanguageService> configuration = (service) -> {
+			service.initializeIfNeeded();
+			ContentModelManager cmManager = (ContentModelManager) service.getComponent(ContentModelManager.class);
+			cmManager.setRootURI("src/test/resources/xsd/");
+			cmManager.setFileAssociations(createXSDAssociationsNoNamespaceSchemaLocationLike(""));
+		};
+
+		String xml = "|";
+		String expected = "<${1|resources|}>$2</${1:resources}>$0";
+		assertSurroundWith(xml, SurroundWithKind.tags, true, configuration, "file:///test/resources.xml", expected);
+	}
+
+	@Test
+	public void surroundEmptySelectionInEmptyDocumentWithTwoSchema() throws Exception {
+		String xml = "<?xml-model href=\"relaxng/tei_all.rng\" ?>\r\n" + //
+				"<?xml-model href=\"relaxng/simple.rng\" ?>\r\n" + //
+				"|";
+		String expected = "<?xml-model href=\"relaxng/tei_all.rng\" ?>\r\n" + //
+				"<?xml-model href=\"relaxng/simple.rng\" ?>\r\n" + //
+				"<${1|TEI,rootelt,teiCorpus|}>$2</${1:TEI}>$0";
+		assertSurroundWith(xml, SurroundWithKind.tags, true, expected);
+	}
+
+	private static XMLFileAssociation[] createXSDAssociationsNoNamespaceSchemaLocationLike(String baseSystemId) {
+		XMLFileAssociation resources = new XMLFileAssociation();
+		resources.setPattern("**/*resources*.xml");
+		resources.setSystemId(baseSystemId + "resources.xsd");
+		return new XMLFileAssociation[] { resources };
+	}
 }