From bc5515d678a187e6370f08bd303e5671ab994fa6 Mon Sep 17 00:00:00 2001
From: Matthias Kurz <m.kurz@irregular.at>
Date: Sun, 27 Jan 2019 15:02:38 +0100
Subject: [PATCH] Add support to write sitemaps index as string

---
 .../sitemapgenerator/SitemapGenerator.java    | 20 ++++++--
 .../SitemapIndexGenerator.java                | 51 ++++++++++++-------
 .../SitemapIndexGeneratorTest.java            |  4 ++
 3 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java b/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
index 99b5242..43eb193 100644
--- a/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
+++ b/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
@@ -218,9 +218,14 @@ abstract class SitemapGenerator<U extends ISitemapUrl, THIS extends SitemapGener
 	 * The sitemap index is written to {baseDir}/sitemap_index.xml
 	 */
 	public File writeSitemapsWithIndex() {
-		if (!finished) throw new RuntimeException("Sitemaps not generated yet; call write() first");
-		File outFile = new File(baseDir, "sitemap_index.xml");
-		return writeSitemapsWithIndex(outFile);
+		return writeSitemapsWithIndex(new File(baseDir, "sitemap_index.xml"));
+	}
+
+	/**
+	 * After you've called {@link #write()}, call this to generate a sitemap index of all sitemaps you generated.
+	 */
+	public String writeSitemapsWithIndexAsString() {
+		return prepareSitemapIndexGenerator(null).writeAsString();
 	}
 
 	/**
@@ -229,11 +234,16 @@ abstract class SitemapGenerator<U extends ISitemapUrl, THIS extends SitemapGener
 	 * @param outFile the destination file of the sitemap index.
 	 */
 	public File writeSitemapsWithIndex(File outFile) {
+		prepareSitemapIndexGenerator(outFile).write();
+		return outFile;
+	}
+
+	private SitemapIndexGenerator prepareSitemapIndexGenerator(File outFile) {
 		if (!finished) throw new RuntimeException("Sitemaps not generated yet; call write() first");
 		SitemapIndexGenerator sig;
 		sig = new SitemapIndexGenerator.Options(baseUrl, outFile).dateFormat(dateFormat).autoValidate(autoValidate).build();
-		sig.addUrls(fileNamePrefix, fileNameSuffix, mapCount).write();
-		return outFile;
+		sig.addUrls(fileNamePrefix, fileNameSuffix, mapCount);
+		return sig;
 	}
 	
 	private void writeSiteMap() throws IOException {
diff --git a/src/main/java/com/redfin/sitemapgenerator/SitemapIndexGenerator.java b/src/main/java/com/redfin/sitemapgenerator/SitemapIndexGenerator.java
index 513fb19..6e576b2 100644
--- a/src/main/java/com/redfin/sitemapgenerator/SitemapIndexGenerator.java
+++ b/src/main/java/com/redfin/sitemapgenerator/SitemapIndexGenerator.java
@@ -221,18 +221,21 @@ public class SitemapIndexGenerator {
 	
 	/** Writes out the sitemap index */
 	public void write() {
+		try {
+			// TODO gzip? is that legal for a sitemap index?
+			write(new FileWriter(outFile));
+		} catch (IOException e) {
+			throw new RuntimeException("Problem writing sitemap index file " + outFile, e);
+		}
+	}
+
+	private void write(OutputStreamWriter out) {
 		if (!allowEmptyIndex && urls.isEmpty()) throw new RuntimeException("No URLs added, sitemap index would be empty; you must add some URLs with addUrls");
 		try {
-			FileWriter out = null;
 			try {
-				// TODO gzip? is that legal for a sitemap index?
-				out = new FileWriter(outFile);
 				writeSiteMap(out);
 				out.flush();
-
 				if (autoValidate) SitemapValidator.validateSitemapIndex(outFile);
-			} catch (IOException e) {
-				throw new RuntimeException("Problem writing sitemap index file " + outFile, e);
 			} catch (SAXException e) {
 				throw new RuntimeException("Problem validating sitemap index file (bug?)", e);
 			} finally {
@@ -246,26 +249,38 @@ public class SitemapIndexGenerator {
 
 	}
 	
-	private void writeSiteMap(OutputStreamWriter out) throws IOException {
-		out.write("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"); 
-		out.write("<sitemapindex xmlns=\"http://www.sitemaps.org/schemas/sitemap/0.9\">\n");
+	public String writeAsString() {
+		StringBuilder sb = new StringBuilder();
+		writeAsString(sb);
+		return sb.toString();
+	}
+
+	private void writeAsString(StringBuilder sb) {
+		sb.append("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"); 
+		sb.append("<sitemapindex xmlns=\"http://www.sitemaps.org/schemas/sitemap/0.9\">\n");
 		for (SitemapIndexUrl url : urls) {
-			out.write("  <sitemap>\n");
-			out.write("    <loc>");
-			out.write(UrlUtils.escapeXml(url.url.toString()));
-			out.write("</loc>\n");
+			sb.append("  <sitemap>\n");
+			sb.append("    <loc>");
+			sb.append(UrlUtils.escapeXml(url.url.toString()));
+			sb.append("</loc>\n");
 			Date lastMod = url.lastMod;
 			
 			if (lastMod == null) lastMod = defaultLastMod;
 			
 			if (lastMod != null) {
-				out.write("    <lastmod>");
-				out.write(dateFormat.format(lastMod));
-				out.write("</lastmod>\n");
+				sb.append("    <lastmod>");
+				sb.append(dateFormat.format(lastMod));
+				sb.append("</lastmod>\n");
 			}
-			out.write("  </sitemap>\n");
+			sb.append("  </sitemap>\n");
 		}
-		out.write("</sitemapindex>");
+		sb.append("</sitemapindex>");
+	}
+
+	private void writeSiteMap(OutputStreamWriter out) throws IOException {
+		StringBuilder sb = new StringBuilder();
+		writeAsString(sb);
+		out.write(sb.toString());
 	}
 
 }
diff --git a/src/test/java/com/redfin/sitemapgenerator/SitemapIndexGeneratorTest.java b/src/test/java/com/redfin/sitemapgenerator/SitemapIndexGeneratorTest.java
index 70ab13e..ca0f1d9 100644
--- a/src/test/java/com/redfin/sitemapgenerator/SitemapIndexGeneratorTest.java
+++ b/src/test/java/com/redfin/sitemapgenerator/SitemapIndexGeneratorTest.java
@@ -95,6 +95,7 @@ public class SitemapIndexGeneratorTest extends TestCase {
 				"</sitemapindex>";
 		String actual = TestUtil.slurpFileAndDelete(outFile);
 		assertEquals(expected, actual);
+		assertEquals(expected, sig.writeAsString());
 	}
 	
 	public void testMaxUrls() throws Exception {
@@ -107,6 +108,7 @@ public class SitemapIndexGeneratorTest extends TestCase {
 		sig.write();
 		String actual = TestUtil.slurpFileAndDelete(outFile);
 		assertEquals(INDEX, actual);
+		assertEquals(INDEX, sig.writeAsString());
 	}
 	
 	public void testOneUrl() throws Exception {
@@ -123,6 +125,7 @@ public class SitemapIndexGeneratorTest extends TestCase {
 				"  </sitemap>\n" + 
 				"</sitemapindex>";
 		assertEquals(expected, actual);
+		assertEquals(expected, sig.writeAsString());
 	}
 	
 	public void testAddByPrefix() throws MalformedURLException {
@@ -132,6 +135,7 @@ public class SitemapIndexGeneratorTest extends TestCase {
 		sig.write();
 		String actual = TestUtil.slurpFileAndDelete(outFile);
 		assertEquals(INDEX, actual);
+		assertEquals(INDEX, sig.writeAsString());
 	}
 	
 }