Refactor: Main refactored, tests added for main (#52)

Author: 0xdade

email_validator/__init__.py

@@ -525,32 +525,31 @@ def main():
     import sys
     import json
 
+    def __utf8_input_shim(input_str):
+        if sys.version_info < (3,):
+            return input_str.decode("utf-8")
+        return input_str
+
+    def __utf8_output_shim(output_str):
+        if sys.version_info < (3,):
+            return unicode_class(output_str).encode("utf-8")
+        return output_str
+
     if len(sys.argv) == 1:
-        # Read lines for STDIN and validate the email address on each line.
-        allow_smtputf8 = True
         for line in sys.stdin:
+            email = __utf8_input_shim(line.strip())
             try:
-                email = line.strip()
-                if sys.version_info < (3,):
-                    email = email.decode("utf8")  # assume utf8 in input
-                validate_email(email, allow_smtputf8=allow_smtputf8)
+                validate_email(email)
             except EmailNotValidError as e:
-                print(email, e)
+                print(__utf8_output_shim("{} {}".format(email, e)))
     else:
         # Validate the email address passed on the command line.
-        email = sys.argv[1]
-        allow_smtputf8 = True
-        check_deliverability = True
-        if sys.version_info < (3,):
-            email = email.decode("utf8")  # assume utf8 in input
+        email = __utf8_input_shim(sys.argv[1])
         try:
-            result = validate_email(email, allow_smtputf8=allow_smtputf8, check_deliverability=check_deliverability)
+            result = validate_email(email)
             print(json.dumps(result.as_dict(), indent=2, sort_keys=True, ensure_ascii=False))
         except EmailNotValidError as e:
-            if sys.version_info < (3,):
-                print(unicode_class(e).encode("utf8"))
-            else:
-                print(e)
+            print(__utf8_output_shim(e))
 
 
 if __name__ == "__main__":

tests/test_main.py

@@ -2,6 +2,8 @@
 from email_validator import EmailSyntaxError, EmailUndeliverableError, \
                             validate_email, validate_email_deliverability, \
                             ValidatedEmail
+# Let's test main but rename it to be clear
+from email_validator import main as validator_main
 
 
 @pytest.mark.parametrize(
@@ -284,3 +286,61 @@ def test_deliverability_dns_timeout():
     assert response.get("unknown-deliverability") == "timeout"
     validate_email('test@gmail.com')
     del validate_email_deliverability.TEST_CHECK_TIMEOUT
+
+
+def test_main_single_good_input(monkeypatch, capsys):
+    import json
+    test_email = "test@example.com"
+    monkeypatch.setattr('sys.argv', ['email_validator', test_email])
+    validator_main()
+    stdout, _ = capsys.readouterr()
+    output = json.loads(str(stdout))
+    assert isinstance(output, dict)
+    assert validate_email(test_email).original_email == output["original_email"]
+
+
+def test_main_single_bad_input(monkeypatch, capsys):
+    bad_email = 'test@..com'
+    monkeypatch.setattr('sys.argv', ['email_validator', bad_email])
+    validator_main()
+    stdout, _ = capsys.readouterr()
+    assert stdout == 'An email address cannot have a period immediately after the @-sign.\n'
+
+
+def test_main_multi_input(monkeypatch, capsys):
+    import io
+    test_cases = ["test@example.com", "test2@example.com", "test@.com", "test3@.com"]
+    test_input = io.StringIO("\n".join(test_cases))
+    monkeypatch.setattr('sys.stdin', test_input)
+    monkeypatch.setattr('sys.argv', ['email_validator'])
+    validator_main()
+    stdout, _ = capsys.readouterr()
+    assert test_cases[0] not in stdout
+    assert test_cases[1] not in stdout
+    assert test_cases[2] in stdout
+    assert test_cases[3] in stdout
+
+
+def test_main_input_shim(monkeypatch, capsys):
+    import json
+    monkeypatch.setattr('sys.version_info', (2, 7))
+    test_email = b"test@example.com"
+    monkeypatch.setattr('sys.argv', ['email_validator', test_email])
+    validator_main()
+    stdout, _ = capsys.readouterr()
+    output = json.loads(str(stdout))
+    assert isinstance(output, dict)
+    assert validate_email(test_email).original_email == output["original_email"]
+
+
+def test_main_output_shim(monkeypatch, capsys):
+    monkeypatch.setattr('sys.version_info', (2, 7))
+    test_email = b"test@.com"
+    monkeypatch.setattr('sys.argv', ['email_validator', test_email])
+    validator_main()
+    stdout, _ = capsys.readouterr()
+
+    # This looks bad but it has to do with the way python 2.7 prints vs py3
+    # The \n is part of the print statement, not part of the string, which is what the b'...' is
+    # Since we're mocking py 2.7 here instead of actually using 2.7, this was the closest I could get
+    assert stdout == "b'An email address cannot have a period immediately after the @-sign.'\n"