adding support for custom suffices for sitemap generation

dfabulich · Sep 9, 2015 · be330d6 · be330d6
1 parent e0c7147
commit be330d6
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git 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 @@ public THIS fileNamePrefix(String fileNamePrefix) {
 		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 @@ public THIS gzip(boolean gzip) {
 	THIS getThis() {
 		return (THIS)this;
 	}
-}
+}
diff --git a/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java b/src/main/java/com/redfin/sitemapgenerator/SitemapGenerator.java
@@ -45,9 +45,15 @@ public SitemapGenerator(AbstractSitemapGeneratorOptions<?> options, ISitemapUrlR
 		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
@@ -4,6 +4,7 @@
 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 void testEmptyWrite() throws Exception {
 			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++) {