From c8c7f39b562989d6ea865eafb59cc88d4957997d Mon Sep 17 00:00:00 2001
From: Tyrel Souza <tyrelsouza@gmail.com>
Date: Tue, 24 May 2016 15:07:36 -0400
Subject: [PATCH] Test api key in env var, or passed in

---
 celigo/core.py | 15 ++++++++++++---
 tests.py       | 18 ++++++++++++++----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/celigo/core.py b/celigo/core.py
index e1060d3..2eae9db 100644
--- a/celigo/core.py
+++ b/celigo/core.py
@@ -19,8 +19,10 @@ class BackupCeligo(object):
     """
     def __init__(self,
                  data_dir,
+                 api_key=None,
                  headers=None,
-                 base_url=None,):
+                 base_url=None,
+                 ):
         """
         :param data_dir: The directory to read and store JSON files
             for imports and configurations.
@@ -32,9 +34,16 @@ class BackupCeligo(object):
         if not base_url:
             base_url = DEFAULT_BASE_URL
 
-        self.api_key = os.environ.get("CELIGO_API_KEY", None)
+        # Set api key if it was passed in
+        self.api_key = api_key
         if not self.api_key:
-            raise Exception("Please set CELIGO_API_KEY environment variable")
+            # if this is still falsy, set from environment variable.
+            self.api_key = os.environ.get("CELIGO_API_KEY")
+            if not self.api_key:
+                # Needs an API key, raise exceptions
+                raise Exception(
+                    "Please pass in api_key, or set the CELIGO_API_KEY "
+                    "environment variable.")
         self.data_dir = data_dir
         self.base_url = base_url
         self.imports_cache = {}
diff --git a/tests.py b/tests.py
index 17b9efb..1bacf6a 100644
--- a/tests.py
+++ b/tests.py
@@ -24,16 +24,26 @@ class CeligoTest(unittest.TestCase):
         data_dir = "fakedir"
         celigo.BackupCeligo(data_dir)
 
-    def test_needs_api_key_env(self):
+    def test_needs_api_key(self):
         """
-        Test that we need to pass in a data directory.
+        Test that we need to either pass in the api_key,
+        or set the CELIGO_API_KEY environment variable.
         """
         data_dir = "fakedir"
         os.environ['CELIGO_API_KEY'] = ''
-        with self.assertRaises(Exception):
+        with self.assertRaisesRegexp(Exception,
+            'Please pass in api_key.*'):
             celigo.BackupCeligo(data_dir)
         os.environ['CELIGO_API_KEY'] = self.FAKE_API_KEY
-        celigo.BackupCeligo(data_dir)
+        bc = celigo.BackupCeligo(data_dir)
+        # Make sure it uses the env var one
+        self.assertEqual(bc.api_key, self.FAKE_API_KEY)
+
+        # Unset env var and make sure it uses the passed in one
+        os.environ['CELIGO_API_KEY'] = ''
+        bc = celigo.BackupCeligo(data_dir, api_key='fake_passed_in')
+        self.assertEqual(bc.api_key, 'fake_passed_in')
+
 
     def test_ensure_directories(self):
         """