From be330d6ec32ed50ce39c2041c233e7b9856235ee Mon Sep 17 00:00:00 2001
From: Ankita Nellimarla <anellimarla@grubhub.com>
Date: Tue, 8 Sep 2015 23:34:43 -0700
Subject: [PATCH] adding support for custom suffices for sitemap generation

---
 .../AbstractSitemapGeneratorOptions.java      |  9 ++++++++-
 .../sitemapgenerator/SitemapGenerator.java    | 10 ++++++++--
 .../SitemapGeneratorTest.java                 | 19 ++++++++++++++++++-
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/redfin/sitemapgenerator/AbstractSitemapGeneratorOptions.java b/src/main/java/com/redfin/sitemapgenerator/AbstractSitemapGeneratorOptions.java
index ccbd1d0..502fa98 100644
--- a/src/main/java/com/redfin/sitemapgenerator/AbstractSitemapGeneratorOptions.java
+++ b/src/main/java/com/redfin/sitemapgenerator/AbstractSitemapGeneratorOptions.java
@@ -10,6 +10,7 @@ abstract class AbstractSitemapGeneratorOptions<THIS extends AbstractSitemapGener
 	URL baseUrl;
 	String fileNamePrefix = "sitemap";
 	boolean allowMultipleSitemaps = true;
+	String suffixStringPattern; // this will store some type of string pattern suitable per needs.
 	W3CDateFormat dateFormat;
 	int maxUrls = SitemapGenerator.MAX_URLS_PER_SITEMAP;
 	boolean autoValidate = false;
@@ -31,6 +32,12 @@ abstract class AbstractSitemapGeneratorOptions<THIS extends AbstractSitemapGener
 		this.fileNamePrefix = fileNamePrefix;
 		return getThis();
 	}
+
+	public THIS suffixStringPattern(String pattern) {
+		this.suffixStringPattern = pattern;
+		return getThis();
+	}
+
 	/** When more than the maximum number of URLs are passed in, should we split into multiple sitemaps automatically, or just throw an exception? */
 	public THIS allowMultipleSitemaps(boolean allowMultipleSitemaps) {
 		this.allowMultipleSitemaps = allowMultipleSitemaps;
@@ -70,4 +77,4 @@ abstract class AbstractSitemapGeneratorOptions<THIS extends AbstractSitemapGener
 	THIS getThis() {
 		return (THIS)this;
 	}
-}
\ No newline at end of file
+}
diff --git a/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java b/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
index 221b546..2952b83 100644
--- a/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
+++ b/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
@@ -45,9 +45,15 @@ abstract class SitemapGenerator<U extends ISitemapUrl, THIS extends SitemapGener
 		autoValidate = options.autoValidate;
 		gzip = options.gzip;
 		this.renderer = renderer;
-		fileNameSuffix = gzip ? ".xml.gz" : ".xml";
+
+		if(options.suffixStringPattern != null && !options.suffixStringPattern.isEmpty()) {
+			fileNameSuffix = gzip ? options.suffixStringPattern + ".xml.gz" : options.suffixStringPattern + ".xml";
+		}
+		else {
+			fileNameSuffix = gzip ? ".xml.gz" : ".xml";
+		}
 	}
-	
+
 	/** Add one URL of the appropriate type to this sitemap.
 	 * If we have reached the maximum number of URLs, we'll throw an exception if {@link #allowMultipleSitemaps} is false,
 	 * or else write out one sitemap immediately.
diff --git a/src/test/java/com/redfin/sitemapgenerator/SitemapGeneratorTest.java b/src/test/java/com/redfin/sitemapgenerator/SitemapGeneratorTest.java
index 749f839..cea3bd2 100644
--- a/src/test/java/com/redfin/sitemapgenerator/SitemapGeneratorTest.java
+++ b/src/test/java/com/redfin/sitemapgenerator/SitemapGeneratorTest.java
@@ -4,6 +4,7 @@ import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStreamReader;
+import java.net.MalformedURLException;
 import java.util.Date;
 import java.util.List;
 import java.util.zip.GZIPInputStream;
@@ -192,7 +193,23 @@ public class SitemapGeneratorTest extends TestCase {
 			fail("Empty write is not allowed");
 		} catch (RuntimeException e) {}
 	}
-	
+
+	public void testSuffixPresent() throws MalformedURLException {
+		wsg = WebSitemapGenerator.builder("http://www.example.com", dir).suffixStringPattern("01").build();
+        wsg.addUrl("http://www.example.com/url1");
+        wsg.addUrl("http://www.example.com/url2");
+		List<File> files = wsg.write();
+		assertEquals("Sitemap has a suffix now", "sitemap01.xml", files.get(0).getName());
+	}
+
+    public void testNullSuffixPassed() throws MalformedURLException {
+        wsg = WebSitemapGenerator.builder("http://www.example.com", dir).suffixStringPattern("").build();
+        wsg.addUrl("http://www.example.com/url1");
+        wsg.addUrl("http://www.example.com/url2");
+        List<File> files = wsg.write();
+        assertEquals("Sitemap has a suffix now", "sitemap.xml", files.get(0).getName());
+    }
+
 	public void testTooManyUrls() throws Exception {
 		wsg = WebSitemapGenerator.builder("http://www.example.com", dir).allowMultipleSitemaps(false).build();
 		for (int i = 0; i < SitemapGenerator.MAX_URLS_PER_SITEMAP; i++) {