Bug 1386902: Add URIUtils.getFormattedDomain. r=liuche draft
authorMichael Comella <michael.l.comella@gmail.com>
Thu, 10 Aug 2017 15:21:42 -0700
changeset 646216 eed9361fef02609776c772d21a004004a1d9d543
parent 646207 824d4f269c6323e1ad2bd8ebeb6496d60b8ba3e5
child 646217 02e1f1224854df4006f275a6877eb5cd6ecf8a78
push id74029
push usermichael.l.comella@gmail.com
push dateTue, 15 Aug 2017 00:19:33 +0000
reviewersliuche
bugs1386902
milestone57.0a1
Bug 1386902: Add URIUtils.getFormattedDomain. r=liuche This combines getBaseDomain & getHostSecondLevel domain and allows us to be more flexible. MozReview-Commit-ID: 7FdWsfZvGFt
mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/menu/BottomSheetContextMenu.java
mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/stream/HighlightItem.java
mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/topsites/TopSitesCard.java
mobile/android/base/java/org/mozilla/gecko/util/URIUtils.java
mobile/android/tests/background/junit4/src/org/mozilla/gecko/util/TestURIUtils.java
--- a/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/menu/BottomSheetContextMenu.java
+++ b/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/menu/BottomSheetContextMenu.java
@@ -25,16 +25,17 @@ import org.mozilla.gecko.icons.IconCallb
 import org.mozilla.gecko.icons.IconResponse;
 import org.mozilla.gecko.icons.Icons;
 import org.mozilla.gecko.util.StringUtils;
 import org.mozilla.gecko.util.URIUtils;
 import org.mozilla.gecko.widget.FaviconView;
 
 import java.lang.ref.WeakReference;
 import java.net.URI;
+import java.net.URISyntaxException;
 
 /* package-private */ class BottomSheetContextMenu
         extends ActivityStreamContextMenu {
 
 
     private final BottomSheetDialog bottomSheetDialog;
 
     private final NavigationView navigationView;
@@ -66,19 +67,27 @@ import java.net.URI;
 
         bottomSheetDialog.setContentView(content);
 
         final String pageTitle = item.getTitle();
         final String sheetPageTitle = !TextUtils.isEmpty(pageTitle) ? pageTitle : item.getUrl();
         ((TextView) content.findViewById(R.id.title)).setText(sheetPageTitle);
 
         final TextView pageDomainView = (TextView) content.findViewById(R.id.url);
-        final UpdatePageDomainAsyncTask updateDomainAsyncTask = new UpdatePageDomainAsyncTask(context, pageDomainView,
-                item.getUrl());
-        updateDomainAsyncTask.execute();
+        final URI itemURI;
+        try {
+            itemURI = new URI(item.getUrl());
+            final UpdatePageDomainAsyncTask updateDomainAsyncTask = new UpdatePageDomainAsyncTask(context, pageDomainView,
+                    itemURI);
+            updateDomainAsyncTask.execute();
+        } catch (final URISyntaxException e) {
+            // Invalid URI: not much processing we can do. Like the async task, the page title view sets itself to the
+            // url on error so we leave this field blank.
+            pageDomainView.setText("");
+        }
 
         // Copy layouted parameters from the Highlights / TopSites items to ensure consistency
         final FaviconView faviconView = (FaviconView) content.findViewById(R.id.icon);
         ViewGroup.LayoutParams layoutParams = faviconView.getLayoutParams();
         layoutParams.width = tilesWidth;
         layoutParams.height = tilesHeight;
         faviconView.setLayoutParams(layoutParams);
 
@@ -122,59 +131,40 @@ import java.net.URI;
         bottomSheetDialog.show();
     }
 
     public void dismiss() {
         bottomSheetDialog.dismiss();
     }
 
     /** Updates the given TextView's text to the page domain. */
-    private static class UpdatePageDomainAsyncTask extends AsyncTask<Void, Void, String> {
-        private final WeakReference<Context> contextWeakReference;
+    private static class UpdatePageDomainAsyncTask extends URIUtils.GetFormattedDomainAsyncTask {
         private final WeakReference<TextView> pageDomainViewWeakReference;
 
-        private final String uriString;
-        @Nullable private final URI uri;
-
-        private UpdatePageDomainAsyncTask(final Context context, final TextView pageDomainView, final String uriString) {
-            this.contextWeakReference = new WeakReference<>(context);
+        private UpdatePageDomainAsyncTask(final Context context, final TextView pageDomainView, final URI uri) {
+            super(context, uri, true, 0); // baseDomain.
             this.pageDomainViewWeakReference = new WeakReference<>(pageDomainView);
-
-            this.uriString = uriString;
-            this.uri = URIUtils.uriOrNull(uriString);
-        }
-
-        @Override
-        protected String doInBackground(final Void... params) {
-            final Context context = contextWeakReference.get();
-            if (context == null || uri == null) {
-                return null;
-            }
-
-            return URIUtils.getBaseDomain(context, uri);
         }
 
         @Override
         protected void onPostExecute(final String baseDomain) {
             super.onPostExecute(baseDomain);
 
             final TextView pageDomainView = pageDomainViewWeakReference.get();
             if (pageDomainView == null) {
                 return;
             }
 
             final String updateText;
             if (!TextUtils.isEmpty(baseDomain)) {
                 updateText = baseDomain;
 
-            // In the unlikely error case, we leave the field blank (null) rather than setting it to the url because
+            // In the unlikely error case, we leave the field blank rather than setting it to the url because
             // the page title view sets itself to the url on error.
-            } else if (uri != null) {
+            } else {
                 final String normalizedHost = StringUtils.stripCommonSubdomains(uri.getHost());
-                updateText = !TextUtils.isEmpty(normalizedHost) ? normalizedHost : null;
-            } else {
-                updateText = null;
+                updateText = !TextUtils.isEmpty(normalizedHost) ? normalizedHost : "";
             }
 
             pageDomainView.setText(updateText);
         }
     }
 }
--- a/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/stream/HighlightItem.java
+++ b/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/stream/HighlightItem.java
@@ -26,16 +26,18 @@ import org.mozilla.gecko.icons.IconRespo
 import org.mozilla.gecko.icons.Icons;
 import org.mozilla.gecko.util.DrawableUtil;
 import org.mozilla.gecko.util.TouchTargetUtil;
 import org.mozilla.gecko.util.URIUtils;
 import org.mozilla.gecko.util.ViewUtil;
 import org.mozilla.gecko.widget.FaviconView;
 
 import java.lang.ref.WeakReference;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.UUID;
 import java.util.concurrent.Future;
 
 public class HighlightItem extends StreamItem implements IconCallback {
     private static final String LOGTAG = "GeckoHighlightItem";
 
     public static final int LAYOUT_ID = R.layout.activity_stream_card_history_item;
     private static final double SIZE_RATIO = 0.75;
@@ -141,35 +143,44 @@ public class HighlightItem extends Strea
             default:
                 pageSourceView.setVisibility(View.INVISIBLE);
                 pageSourceIconView.setImageResource(0);
                 break;
         }
     }
 
     private void updatePageDomain() {
-        final UpdatePageDomainAsyncTask hostSLDTask = new UpdatePageDomainAsyncTask(itemView.getContext(),
-                highlight.getUrl(), pageDomainView);
-        hostSLDTask.execute();
+        final URI highlightURI;
+        try {
+            highlightURI = new URI(highlight.getUrl());
+        } catch (final URISyntaxException e) {
+            // If the URL is invalid, there's not much extra processing we can do on it.
+            pageDomainView.setText(highlight.getUrl());
+            return;
+        }
+
+        final UpdatePageDomainAsyncTask updatePageDomainTask = new UpdatePageDomainAsyncTask(itemView.getContext(),
+                highlightURI, pageDomainView);
+        updatePageDomainTask.execute();
     }
 
     @Override
     public void onIconResponse(IconResponse response) {
         pageIconView.updateImage(response);
     }
 
     /** Updates the text of the given view to the host second level domain. */
-    private static class UpdatePageDomainAsyncTask extends URIUtils.GetHostSecondLevelDomainAsyncTask {
+    private static class UpdatePageDomainAsyncTask extends URIUtils.GetFormattedDomainAsyncTask {
         private static final int VIEW_TAG_ID = R.id.page; // same as the view.
 
         private final WeakReference<TextView> pageDomainViewWeakReference;
         private final UUID viewTagAtStart;
 
-        UpdatePageDomainAsyncTask(final Context contextReference, final String uriString, final TextView pageDomainView) {
-            super(contextReference, uriString);
+        UpdatePageDomainAsyncTask(final Context contextReference, final URI uri, final TextView pageDomainView) {
+            super(contextReference, uri, false, 0); // hostSLD.
             this.pageDomainViewWeakReference = new WeakReference<>(pageDomainView);
 
             // See isTagSameAsStartTag for details.
             viewTagAtStart = UUID.randomUUID();
             pageDomainView.setTag(VIEW_TAG_ID, viewTagAtStart);
         }
 
         @Override
--- a/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/topsites/TopSitesCard.java
+++ b/mobile/android/base/java/org/mozilla/gecko/activitystream/homepanel/topsites/TopSitesCard.java
@@ -26,16 +26,18 @@ import org.mozilla.gecko.icons.IconRespo
 import org.mozilla.gecko.icons.Icons;
 import org.mozilla.gecko.util.DrawableUtil;
 import org.mozilla.gecko.util.TouchTargetUtil;
 import org.mozilla.gecko.util.URIUtils;
 import org.mozilla.gecko.util.ViewUtil;
 import org.mozilla.gecko.widget.FaviconView;
 
 import java.lang.ref.WeakReference;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.UUID;
 import java.util.concurrent.Future;
 
 /* package-local */ class TopSitesCard extends RecyclerView.ViewHolder
         implements IconCallback {
     private final FaviconView faviconView;
 
     private final TextView title;
@@ -106,58 +108,71 @@ import java.util.concurrent.Future;
             pinDrawable = null;
         }
         TextViewCompat.setCompoundDrawablesRelativeWithIntrinsicBounds(title, pinDrawable, null, null, null);
 
         final String provider = topSite.getMetadata().getProvider();
         if (!TextUtils.isEmpty(provider)) {
             title.setText(provider.toLowerCase());
         } else {
+            final URI topSiteURI;
+            try {
+                topSiteURI = new URI(topSite.getUrl());
+            } catch (final URISyntaxException e) {
+                // If this is not a valid URI, there is not much processing we can do on it.
+                // Also, see comment below regarding setCenteredText.
+                setTopSiteTitle(title, topSite.getUrl());
+                return;
+            }
+
             // Our AsyncTask calls setCenteredText(), which needs to have all drawable's in place to correctly
             // layout the text, so we need to wait with requesting the title until we've set our pin icon.
             final UpdateCardTitleAsyncTask titleAsyncTask = new UpdateCardTitleAsyncTask(itemView.getContext(),
-                    topSite.getUrl(), title);
+                    topSiteURI, title);
             titleAsyncTask.execute();
         }
     }
 
+    private static void setTopSiteTitle(final TextView textView, final String title) {
+        // We use consistent padding all around the title, and the top padding is never modified,
+        // so we can pass that in as the default padding:
+        ViewUtil.setCenteredText(textView, title, textView.getPaddingTop());
+    }
+
     @Override
     public void onIconResponse(IconResponse response) {
         faviconView.updateImage(response);
     }
 
     /** Updates the text of the given view to the page domain. */
-    private static class UpdateCardTitleAsyncTask extends URIUtils.GetHostSecondLevelDomainAsyncTask {
+    private static class UpdateCardTitleAsyncTask extends URIUtils.GetFormattedDomainAsyncTask {
         private static final int VIEW_TAG_ID = R.id.title; // same as the view.
 
         private final WeakReference<TextView> titleViewWeakReference;
         private final UUID viewTagAtStart;
 
-        UpdateCardTitleAsyncTask(final Context contextReference, final String uriString, final TextView titleView) {
-            super(contextReference, uriString);
+        UpdateCardTitleAsyncTask(final Context contextReference, final URI uri, final TextView titleView) {
+            super(contextReference, uri, false, 0); // hostSLD.
             this.titleViewWeakReference = new WeakReference<>(titleView);
 
             // See isTagSameAsStartTag for details.
             viewTagAtStart = UUID.randomUUID();
             titleView.setTag(VIEW_TAG_ID, viewTagAtStart);
         }
 
         @Override
         protected void onPostExecute(final String hostSLD) {
             super.onPostExecute(hostSLD);
             final TextView titleView = titleViewWeakReference.get();
             if (titleView == null || !isTagSameAsStartTag(titleView)) {
                 return;
             }
 
-            final String updateText = !TextUtils.isEmpty(hostSLD) ? hostSLD : uriString;
-
-            // We use consistent padding all around the title, and the top padding is never modified,
-            // so we can pass that in as the default padding:
-            ViewUtil.setCenteredText(titleView, updateText, titleView.getPaddingTop());
+            final String updateText = !TextUtils.isEmpty(hostSLD) ? hostSLD : uri.toString();
+            setTopSiteTitle(titleView, updateText);
         }
 
         /**
          * Returns true if the tag on the given view matches the tag from the constructor. We do this to ensure
          * the View we're making this request for hasn't been re-used by the time this request completes.
          */
         @UiThread
         private boolean isTagSameAsStartTag(final View viewToCheck) {
--- a/mobile/android/base/java/org/mozilla/gecko/util/URIUtils.java
+++ b/mobile/android/base/java/org/mozilla/gecko/util/URIUtils.java
@@ -1,21 +1,23 @@
 /* This Source Code Form is subject to the terms of the Mozilla Public
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 package org.mozilla.gecko.util;
 
 import android.content.Context;
 import android.os.AsyncTask;
+import android.support.annotation.IntRange;
 import android.support.annotation.NonNull;
 import android.support.annotation.Nullable;
 import android.support.annotation.WorkerThread;
+import android.text.TextUtils;
 import android.util.Log;
-import ch.boye.httpclientandroidlib.util.TextUtils;
+import ch.boye.httpclientandroidlib.conn.util.InetAddressUtils;
 import org.mozilla.gecko.util.publicsuffix.PublicSuffix;
 
 import java.lang.ref.WeakReference;
 import java.net.URI;
 import java.net.URISyntaxException;
 
 /** Utilities for operating on URLs. */
 public class URIUtils {
@@ -29,103 +31,128 @@ public class URIUtils {
         try {
             return new URI(uriString);
         } catch (final URISyntaxException e) {
             return null;
         }
     }
 
     /**
-     * Returns the second level domain (SLD) of a url. It removes any subdomain/TLD.
-     * e.g. https://m.foo.com/bar/baz?noo=abc#123  => foo
+     * Returns the domain for the given URI, formatted by the other available parameters.
+     *
+     * A public suffix is a top-level domain. For the input, "https://github.com", you can specify
+     * {@code shouldIncludePublicSuffix}:
+     * - true: "github.com"
+     * - false: "github"
      *
-     * The return value may still contain a public suffix (e.g. .com) if the suffix does not match any of our
-     * known values. If a host cannot be determined from the given url String, the empty String will be returned.
+     * The subdomain count is the number of subdomains you want to include; the domain will always be included. For
+     * the input, "https://m.blog.github.io/", excluding the public suffix and the subdomain count:
+     * - 0: "github"
+     * - 1: "blog.github.com"
+     * - 2: "m.blog.github.com"
+     *
+     * ipv4 & ipv6 urls will return the address directly.
      *
-     * This implementation is taken from Firefox for iOS:
-     *   https://github.com/mozilla-mobile/firefox-ios/blob/deb9736c905cdf06822ecc4a20152df7b342925d/Shared/Extensions/NSURLExtensions.swift#L152
+     * This implementation is influenced by Firefox iOS and can be used in place of some URI formatting functions:
+     * - hostSLD [1]: exclude publicSuffix, 0 subdomains
+     * - baseDomain [2]: include publicSuffix, 0 subdomains
+     *
+     * Expressing the method this way (instead of separate baseDomain & hostSLD methods) is more flexible if we want to
+     * change the subdomain count and works well with our {@link GetFormattedDomainAsyncTask}, which can take the
+     * same parameters we pass in here, rather than creating a new Task for each method.
      *
-     * @param uriString A url from which to extract the second level domain.
-     * @return The second level domain of the url or the empty String when the host cannot be determined.
+     * [1]: https://github.com/mozilla-mobile/firefox-ios/blob/deb9736c905cdf06822ecc4a20152df7b342925d/Shared/Extensions/NSURLExtensions.swift#L152
+     * [2]: https://github.com/mozilla-mobile/firefox-ios/blob/deb9736c905cdf06822ecc4a20152df7b342925d/Shared/Extensions/NSURLExtensions.swift#L205
+     *
+     * @param context the Activity context.
+     * @param uri the URI whose host we should format.
+     * @param shouldIncludePublicSuffix true if the public suffix should be included, false otherwise.
+     * @param subdomainCount The number of subdomains to include.
+     *
+     * @return the formatted domain, or the empty String if the host cannot be found.
      */
-    @WorkerThread // PublicSuffix methods can touch the disk.
-    public static String getHostSecondLevelDomain(@NonNull final Context context, @NonNull final String uriString)
-            throws URISyntaxException {
+    @NonNull
+    @WorkerThread // calls PublicSuffix methods.
+    public static String getFormattedDomain(@NonNull final Context context, @NonNull final URI uri,
+            final boolean shouldIncludePublicSuffix, @IntRange(from = 0) final int subdomainCount) {
         if (context == null) { throw new NullPointerException("Expected non-null Context argument"); }
-        if (uriString == null) { throw new NullPointerException("Expected non-null uri argument"); }
+        if (uri == null) { throw new NullPointerException("Expected non-null uri argument"); }
+        if (subdomainCount < 0) { throw new IllegalArgumentException("Expected subdomainCount >= 0."); }
 
-        final URI uri = new URI(uriString);
-        final String baseDomain = getBaseDomain(context, uri);
-        if (baseDomain == null) {
-            final String normalizedHost = StringUtils.stripCommonSubdomains(uri.getHost());
-            return !TextUtils.isEmpty(normalizedHost) ? normalizedHost : "";
+        final String host = uri.getHost();
+        if (TextUtils.isEmpty(host)) {
+            return ""; // There's no host so there's no domain to retrieve.
         }
 
-        return PublicSuffix.stripPublicSuffix(context, baseDomain);
-    }
-
-    /**
-     * Returns the base domain from a given hostname. The base domain name is defined as the public domain suffix
-     * with the base private domain attached to the front. For example, for the URL www.bbc.co.uk, the base domain
-     * would be bbc.co.uk. The base domain includes the public suffix (co.uk) + one level down (bbc).
-     *
-     * IPv4 & IPv6 urls are not supported and will return null.
-     *
-     * This implementation is taken from Firefox for iOS:
-     *   https://github.com/mozilla-mobile/firefox-ios/blob/deb9736c905cdf06822ecc4a20152df7b342925d/Shared/Extensions/NSURLExtensions.swift#L205
-     *
-     * @param uri The uri to find the base domain of
-     * @return The base domain string for the given host name, or null if not applicable.
-     */
-    @Nullable
-    @WorkerThread // PublicSuffix methods can touch the disk.
-    public static String getBaseDomain(@NonNull final Context context, final URI uri) {
-        final String host = uri.getHost();
-        if (isIPv6(uri) || TextUtils.isEmpty(host)) {
-            return null;
-        }
-
-        // If this is just a hostname and not a FQDN, use the entire hostname.
-        if (!host.contains(".")) {
+        if (InetAddressUtils.isIPv4Address(host) ||
+                isIPv6(uri) ||
+                !host.contains(".")) { // If this is just a hostname and not a FQDN, use the entire hostname.
             return host;
         }
 
-        final String publicSuffixWithDomain = PublicSuffix.getPublicSuffix(context, host, 1);
-        return !TextUtils.isEmpty(publicSuffixWithDomain) ? publicSuffixWithDomain : null;
+        final String domainStr = PublicSuffix.getPublicSuffix(context, host, subdomainCount + 1);
+        if (TextUtils.isEmpty(domainStr)) {
+            // There is no public suffix found so we assume the whole host is a domain.
+            return stripSubdomains(host, subdomainCount);
+        }
+
+        if (!shouldIncludePublicSuffix) {
+            // We could be slightly more efficient if we wrote a new algorithm rather than using PublicSuffix twice
+            // but I don't think it's worth the time and it'd complicate the code with more independent branches.
+            return PublicSuffix.stripPublicSuffix(context, domainStr);
+        }
+        return domainStr;
+    }
+
+    /** Strips any subdomains from the host over the given limit. */
+    private static String stripSubdomains(String host, final int desiredSubdomainCount) {
+        int includedSubdomainCount = 0;
+        for (int i = host.length() - 1; i >= 0; --i) {
+            if (host.charAt(i) == '.') {
+                if (includedSubdomainCount >= desiredSubdomainCount) {
+                    return host.substring(i + 1, host.length());
+                }
+
+                includedSubdomainCount += 1;
+            }
+        }
+
+        // There are fewer subdomains than the total we'll accept so return them all!
+        return host;
     }
 
     // impl via FFiOS: https://github.com/mozilla-mobile/firefox-ios/blob/deb9736c905cdf06822ecc4a20152df7b342925d/Shared/Extensions/NSURLExtensions.swift#L292
     private static boolean isIPv6(final URI uri) {
         final String host = uri.getHost();
         return !TextUtils.isEmpty(host) && host.contains(":");
     }
 
     /**
      * An async task that will take a URI formatted as a String and will retrieve
-     * {@link #getHostSecondLevelDomain(Context, String)}. To use this, extend the class and override
-     * {@link #onPostExecute(Object)}, where the secondLevelDomain, or the empty String if the host cannot determined,
+     * {@link #getFormattedDomain(Context, URI, boolean, int)}. To use this, extend the class and override
+     * {@link #onPostExecute(Object)}, where the formatted domain, or the empty String if the host cannot determined,
      * will be returned.
      */
-    public static abstract class GetHostSecondLevelDomainAsyncTask extends AsyncTask<Void, Void, String> {
+    public static abstract class GetFormattedDomainAsyncTask extends AsyncTask<Void, Void, String> {
         protected final WeakReference<Context> contextWeakReference;
-        protected final String uriString;
+        protected final URI uri;
+        protected final boolean shouldIncludePublicSuffix;
+        protected final int subdomainCount;
 
-        public GetHostSecondLevelDomainAsyncTask(final Context contextWeakReference, final String uriString) {
-            this.contextWeakReference = new WeakReference<>(contextWeakReference);
-            this.uriString = uriString;
+        public GetFormattedDomainAsyncTask(final Context context, final URI uri, final boolean shouldIncludePublicSuffix,
+                final int subdomainCount) {
+            this.contextWeakReference = new WeakReference<>(context);
+            this.uri = uri;
+            this.shouldIncludePublicSuffix = shouldIncludePublicSuffix;
+            this.subdomainCount = subdomainCount;
         }
 
         @Override
         protected String doInBackground(final Void... params) {
             final Context context = contextWeakReference.get();
             if (context == null) {
-                return null;
+                return "";
             }
 
-            try {
-                return URIUtils.getHostSecondLevelDomain(context, uriString);
-            } catch (final URISyntaxException e) {
-                Log.w(LOGTAG, "Unable to fetch second level domain."); // Don't log exception to avoid logging pii/urls.
-                return null;
-            }
+            return URIUtils.getFormattedDomain(context, uri, shouldIncludePublicSuffix, subdomainCount);
         }
     }
 }
--- a/mobile/android/tests/background/junit4/src/org/mozilla/gecko/util/TestURIUtils.java
+++ b/mobile/android/tests/background/junit4/src/org/mozilla/gecko/util/TestURIUtils.java
@@ -7,116 +7,168 @@ package org.mozilla.gecko.util;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mozilla.gecko.background.testhelpers.TestRunner;
 import org.robolectric.RuntimeEnvironment;
 
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.HashMap;
-import java.util.Map;
 
 @RunWith(TestRunner.class)
 public class TestURIUtils {
 
     private final String BUGZILLA_URL = "https://bugzilla.mozilla.org/enter_bug.cgi?format=guided#h=dupes%7CData%20%26%20BI%20Services%20Team%7C";
 
+    // --- getFormattedDomain, include PublicSuffix --- //
     @Test
-    public void testGetHostSecondLevelDomain() throws Exception {
-        assertGetHostSLD("https://www.example.com/index.html", "example");
-        assertGetHostSLD("https://m.foo.com/bar/baz?noo=abc#123", "foo");
-        assertGetHostSLD("https://user:pass@m.foo.com/bar/baz?noo=abc#123", "foo");
+    public void testGetFormattedDomainWithSuffix0Parts() {
+        final boolean includePublicSuffix = true;
+        final int subdomainCount = 0;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google.com");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "example.com");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "foo.com");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "foo.com");
     }
 
     @Test
-    public void testGetHostSecondLevelDomainURIHasNoHost() throws Exception {
-        assertGetHostSLD("file:///usr/bin", "");
-    }
-
-    @Test
-    public void testGetHostSecondLevelDomainIPv4() throws Exception {
-        assertGetHostSLD("http://192.168.1.1", "192.168.1.1");
+    public void testGetFormattedDomainWithSuffix1Parts() {
+        final boolean includePublicSuffix = true;
+        final int subdomainCount = 1;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google.com");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "www.example.com");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "blog.foo.com");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.foo.com");
     }
 
     @Test
-    public void testGetHostSecondLevelDomainIPv6() throws Exception {
-        assertGetHostSLD("http://[3ffe:1900:4545:3:200:f8ff:fe21:67cf]", "[3ffe:1900:4545:3:200:f8ff:fe21:67cf]");
-    }
-
-    @Test(expected = URISyntaxException.class)
-    public void testGetHostSecondLevelDomainNonURI() throws Exception {
-        URIUtils.getHostSecondLevelDomain(RuntimeEnvironment.application, "this  -is  -not-a-uri");
+    public void testGetFormattedDomainWithSuffix2Parts() {
+        final boolean includePublicSuffix = true;
+        final int subdomainCount = 2;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google.com");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "www.example.com");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.blog.foo.com");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.foo.com");
     }
 
-    @Test(expected = NullPointerException.class)
-    public void testGetHostSecondLevelDomainNullContextThrows() throws Exception {
-        URIUtils.getHostSecondLevelDomain(null, "http://google.com");
-    }
-
-    @Test(expected = NullPointerException.class)
-    public void testGetHostSecondLevelDomainNullURIThrows() throws Exception {
-        URIUtils.getHostSecondLevelDomain(RuntimeEnvironment.application, null);
-    }
-
-    // SLD = second level domain.
-    private void assertGetHostSLD(final String input, final String expected) throws Exception {
-        Assert.assertEquals("for input:" + input + "||", expected,
-                URIUtils.getHostSecondLevelDomain(RuntimeEnvironment.application, input));
+    // --- getFormattedDomain, exclude PublicSuffix --- //
+    @Test
+    public void testGetFormattedDomainNoSuffix0Parts() {
+        final boolean includePublicSuffix = false;
+        final int subdomainCount = 0;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "example");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "foo");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "foo");
     }
 
     @Test
-    public void testGetBaseDomainNormal() throws Exception {
-        assertGetBaseDomain("http://bbc.co.uk", "bbc.co.uk");
+    public void testGetFormattedDomainNoSuffix1Parts() {
+        final boolean includePublicSuffix = false;
+        final int subdomainCount = 1;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "www.example");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "blog.foo");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.foo");
+    }
+
+    @Test
+    public void testGetFormattedDomainNoSuffix2Parts() {
+        final boolean includePublicSuffix = false;
+        final int subdomainCount = 2;
+        assertGetFormattedDomain("https://google.com/search", includePublicSuffix, subdomainCount, "google");
+        assertGetFormattedDomain("https://www.example.com/index.html", includePublicSuffix, subdomainCount, "www.example");
+        assertGetFormattedDomain("https://m.blog.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.blog.foo");
+        assertGetFormattedDomain("https://user:pass@m.foo.com/bar/baz?noo=abc#123", includePublicSuffix, subdomainCount, "m.foo");
+    }
+
+    // --- getFormattedDomain, saving time by not splitting up these tests on public suffix param. --- //
+    @Test
+    public void testGetFormattedDomainTwoLevelPublicSuffix() throws Exception {
+        assertGetFormattedDomain("http://bbc.co.uk", false, 0, "bbc");
+        assertGetFormattedDomain("http://bbc.co.uk", true, 0, "bbc.co.uk");
     }
 
     @Test
-    public void testGetBaseDomainNormalWithAdditionalSubdomain() throws Exception {
-        assertGetBaseDomain("http://a.bbc.co.uk", "bbc.co.uk");
-        assertGetBaseDomain(BUGZILLA_URL, "mozilla.org");
+    public void testGetFormattedDomainNormalTwoLevelPublicSuffixWithSubdomain() throws Exception {
+        assertGetFormattedDomain("http://a.bbc.co.uk", false, 0, "bbc");
+        assertGetFormattedDomain("http://a.bbc.co.uk", true, 0, "bbc.co.uk");
+        assertGetFormattedDomain(BUGZILLA_URL, false, 0, "mozilla");
+        assertGetFormattedDomain(BUGZILLA_URL, true, 0, "mozilla.org");
+    }
+
+    @Test
+    public void testGetFormattedDomainWilcardDomain() throws Exception {
+        // TLD entry: *.kawasaki.jp
+        assertGetFormattedDomain("http://a.b.kawasaki.jp", false, 0, "a");
+        assertGetFormattedDomain("http://a.b.kawasaki.jp", true, 0, "a.b.kawasaki.jp");
     }
 
     @Test
-    public void testGetBaseDomainWilcardDomain() throws Exception {
+    public void testGetFormattedDomainWilcardDomainWithAdditionalSubdomain() throws Exception {
         // TLD entry: *.kawasaki.jp
-        assertGetBaseDomain("http://a.b.kawasaki.jp", "a.b.kawasaki.jp");
+        assertGetFormattedDomain("http://a.b.c.kawasaki.jp", false, 0, "b");
+        assertGetFormattedDomain("http://a.b.c.kawasaki.jp", true, 0, "b.c.kawasaki.jp");
     }
 
     @Test
-    public void testGetBaseDomainWilcardDomainWithAdditionalSubdomain() throws Exception {
-        // TLD entry: *.kawasaki.jp
-        assertGetBaseDomain("http://a.b.c.kawasaki.jp", "b.c.kawasaki.jp");
+    public void testGetFormattedDomainExceptionDomain() throws Exception {
+        // TLD entry: !city.kawasaki.jp
+        assertGetFormattedDomain("http://city.kawasaki.jp", false, 0, "city");
+        assertGetFormattedDomain("http://city.kawasaki.jp", true, 0, "city.kawasaki.jp");
     }
 
     @Test
-    public void testGetBaseDomainExceptionDomain() throws Exception {
+    public void testGetFormattedDomainExceptionDomainWithAdditionalSubdomain() throws Exception {
         // TLD entry: !city.kawasaki.jp
-        assertGetBaseDomain("http://city.kawasaki.jp", "city.kawasaki.jp");
+        assertGetFormattedDomain("http://a.city.kawasaki.jp", false, 0, "city");
+        assertGetFormattedDomain("http://a.city.kawasaki.jp", true, 0, "city.kawasaki.jp");
     }
 
     @Test
-    public void testGetBaseDomainExceptionDomainWithAdditionalSubdomain() throws Exception {
+    public void testGetFormattedDomainExceptionDomainBugzillaURL() throws Exception {
         // TLD entry: !city.kawasaki.jp
-        assertGetBaseDomain("http://a.city.kawasaki.jp", "city.kawasaki.jp");
+        assertGetFormattedDomain("http://a.city.kawasaki.jp", false, 0, "city");
+        assertGetFormattedDomain("http://a.city.kawasaki.jp", true, 0, "city.kawasaki.jp");
     }
 
     @Test
-    public void testGetBaseDomainExceptionDomainBugzillaURL() throws Exception {
-        // TLD entry: !city.kawasaki.jp
-        assertGetBaseDomain("http://a.city.kawasaki.jp", "city.kawasaki.jp");
+    public void testGetFormattedDomainURIHasNoHost() throws Exception {
+        assertGetFormattedDomain("file:///usr/bin", false, 0, "");
+        assertGetFormattedDomain("file:///usr/bin", true, 0, "");
+    }
+
+    @Test
+    public void testGetFormattedDomainIPv4() throws Exception {
+        assertGetFormattedDomain("http://192.168.1.1", false, 0, "192.168.1.1");
+        assertGetFormattedDomain("http://192.168.1.1", true, 0, "192.168.1.1");
     }
 
     @Test
-    public void testGetBaseDomainIPv4() throws Exception {
-        assertGetBaseDomain("http://192.168.1.1", null);
+    public void testGetFormattedDomainIPv6() throws Exception {
+        assertGetFormattedDomain("http://[3ffe:1900:4545:3:200:f8ff:fe21:67cf]", false, 0, "[3ffe:1900:4545:3:200:f8ff:fe21:67cf]");
+        assertGetFormattedDomain("http://[3ffe:1900:4545:3:200:f8ff:fe21:67cf]", true, 0, "[3ffe:1900:4545:3:200:f8ff:fe21:67cf]");
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testGetFormattedDomainNullContextThrows() throws Exception {
+        URIUtils.getFormattedDomain(null, new URI("http://google.com"), false, 0);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testGetFormattedDomainNullURIThrows() throws Exception {
+        URIUtils.getFormattedDomain(RuntimeEnvironment.application, null, false, 0);
     }
 
-    @Test
-    public void testGetBaseDomainIPv6() throws Exception {
-        assertGetBaseDomain("http://[3ffe:1900:4545:3:200:f8ff:fe21:67cf]", null);
-    }
+    private void assertGetFormattedDomain(final String uriString, final boolean includePublicSuffix,
+            final int subdomainCount, final String expected) {
+        final URI uri;
+        try {
+            uri = new URI(uriString);
+        } catch (final URISyntaxException e) {
+            throw new IllegalArgumentException("Invalid URI passed into test: " + uriString);
+        }
 
-    private void assertGetBaseDomain(final String input, final String expected) throws Exception {
-        Assert.assertEquals("for input:" + input + "||",
+        Assert.assertEquals("for input:" + uriString + "||",
                 expected,
-                URIUtils.getBaseDomain(RuntimeEnvironment.application, new URI(input)));
+               URIUtils.getFormattedDomain(RuntimeEnvironment.application, uri, includePublicSuffix, subdomainCount));
     }
 }
\ No newline at end of file