Improved remove_password_from_url(), added tests

Author: erezsh

data_diff/utils.py

@@ -54,7 +54,18 @@ def number_to_human(n):
     return "{:.0f}{}".format(n / 10 ** (3 * millidx), millnames[millidx])
 
 
-def remove_password_from_url(url, replace_with="***"):
+def _join_if_any(sym, args):
+    args = list(args)
+    if not args:
+        return ''
+    return sym.join(str(a) for a in args if a)
+
+def remove_password_from_url(url: str, replace_with: str="***") -> str:
     parsed = urlparse(url)
-    replaced = parsed._replace(netloc="{}:{}@{}".format(parsed.username, replace_with, parsed.hostname or ""))
+    account = parsed.username or ''
+    if parsed.password:
+        account += ':' + replace_with
+    host = _join_if_any(":", filter(None, [parsed.hostname, parsed.port]))
+    netloc = _join_if_any("@", filter(None, [account, host]))
+    replaced = parsed._replace(netloc=netloc)
     return replaced.geturl()

tests/test_config.py

@@ -1,6 +1,7 @@
 import unittest
 
 from data_diff.config import apply_config_from_string, ConfigParseError
+from data_diff.utils import remove_password_from_url
 
 
 class TestConfig(unittest.TestCase):
@@ -41,3 +42,20 @@ def test_basic(self):
         res = apply_config_from_string(config, "pg_pg", {"update_column": "foo", "table2": "bar"})
         assert res["update_column"] == "foo"
         assert res["table2"] == "bar"
+
+    def test_remove_password(self):
+        replace_with = "*****"
+        urls = [
+            'd://host/',
+            'd://host:123/',
+            'd://user@host:123/',
+            'd://user:PASS@host:123/',
+            'd://:PASS@host:123/',
+            'd://:PASS@host:123/path',
+            'd://:PASS@host:123/path?whatever#blabla',
+        ]
+        for url in urls:
+            removed = remove_password_from_url(url, replace_with)
+            expected = url.replace('PASS', replace_with)
+            removed = remove_password_from_url(url, replace_with)
+            self.assertEqual(removed, expected)