From cb527f1e3922e96b7c2c8dd1d5d607944ecab134 Mon Sep 17 00:00:00 2001
From: Saleem Latif <saleem_ee@hotmail.com>
Date: Mon, 17 Oct 2016 18:13:03 +0500
Subject: [PATCH] Fix NameError issue for saml mangemement command

---
 .../management/commands/saml.py               | 28 +++++++++--------
 .../management/commands/tests/__init__.py     |  3 ++
 .../management/commands/tests/test_saml.py    | 31 +++++++++++++++++++
 3 files changed, 49 insertions(+), 13 deletions(-)
 create mode 100644 common/djangoapps/third_party_auth/management/commands/tests/__init__.py
 create mode 100644 common/djangoapps/third_party_auth/management/commands/tests/test_saml.py

diff --git a/common/djangoapps/third_party_auth/management/commands/saml.py b/common/djangoapps/third_party_auth/management/commands/saml.py
index 8a1590ae66..238599e2cb 100644
--- a/common/djangoapps/third_party_auth/management/commands/saml.py
+++ b/common/djangoapps/third_party_auth/management/commands/saml.py
@@ -15,17 +15,19 @@ class Command(BaseCommand):
         parser.add_argument('--pull', action='store_true', help="Pull updated metadata from external IDPs")
 
     def handle(self, *args, **options):
-        if options['pull']:
-            log_handler = logging.StreamHandler(self.stdout)
-            log_handler.setLevel(logging.DEBUG)
-            log = logging.getLogger('third_party_auth.tasks')
-            log.propagate = False
-            log.addHandler(log_handler)
-            num_changed, num_failed, num_total = fetch_saml_metadata()
-            self.stdout.write(
-                "\nDone. Fetched {num_total} total. {num_changed} were updated and {num_failed} failed.\n".format(
-                    num_changed=num_changed, num_failed=num_failed, num_total=num_total
-                )
+        should_pull_saml_metadata = options.get('pull', False)
+
+        if not should_pull_saml_metadata:
+            raise CommandError("Command can only be used with '--pull' option.")
+
+        log_handler = logging.StreamHandler(self.stdout)
+        log_handler.setLevel(logging.DEBUG)
+        log = logging.getLogger('third_party_auth.tasks')
+        log.propagate = False
+        log.addHandler(log_handler)
+        num_changed, num_failed, num_total = fetch_saml_metadata()
+        self.stdout.write(
+            "\nDone. Fetched {num_total} total. {num_changed} were updated and {num_failed} failed.\n".format(
+                num_changed=num_changed, num_failed=num_failed, num_total=num_total
             )
-        else:
-            raise CommandError("Unknown argment: {}".format(subcommand))
+        )
diff --git a/common/djangoapps/third_party_auth/management/commands/tests/__init__.py b/common/djangoapps/third_party_auth/management/commands/tests/__init__.py
new file mode 100644
index 0000000000..c4f4248f89
--- /dev/null
+++ b/common/djangoapps/third_party_auth/management/commands/tests/__init__.py
@@ -0,0 +1,3 @@
+"""
+This directory contains tests for third_party_auth app.
+"""
diff --git a/common/djangoapps/third_party_auth/management/commands/tests/test_saml.py b/common/djangoapps/third_party_auth/management/commands/tests/test_saml.py
new file mode 100644
index 0000000000..12d790bd1e
--- /dev/null
+++ b/common/djangoapps/third_party_auth/management/commands/tests/test_saml.py
@@ -0,0 +1,31 @@
+"""
+Tests for `saml` management command, this command fetches saml metadata from providers and updates
+existing data accordingly.
+"""
+import unittest
+
+from django.test import TestCase
+from django.core.management import call_command
+from django.core.management.base import CommandError
+from django.conf import settings
+
+
+@unittest.skipUnless(settings.ROOT_URLCONF == 'lms.urls', 'Test only valid in lms')
+class TestSAMLCommand(TestCase):
+    """
+    Test django management command for fetching saml metadata.
+    """
+    def test_raises_command_error_for_invalid_arguments(self):
+        """
+        Test that management command raises `CommandError` with a proper message in case of
+        invalid command arguments.
+
+        This test would fail with an error if ValueError is raised.
+        """
+        # Call `saml` command without any argument so that it raises a CommandError
+        with self.assertRaisesMessage(CommandError, "Command can only be used with '--pull' option."):
+            call_command("saml")
+
+        # Call `saml` command without any argument so that it raises a CommandError
+        with self.assertRaisesMessage(CommandError, "Command can only be used with '--pull' option."):
+            call_command("saml", pull=False)