Switch the configuration system

I like this one much better.

metrik/batch.py

@@ -3,7 +3,9 @@ from luigi import build
 from datetime import datetime
 from argparse import ArgumentParser
 from dateutil.parser import parse
+from six import StringIO
 
+from metrik.conf import get_config
 from metrik.flows.rates_flow import LiborFlow
 from metrik.flows.equities_flow import EquitiesFlow
 from metrik import __version__
@@ -49,6 +51,7 @@ def handle_commandline():
     parser.add_argument('-f', '--flow', dest='flow', help='The flow to be run')
     parser.add_argument('-l', '--list-flows', dest='list', action='store_true',
                         help='List all available flows to be run.')
+    parser.add_argument('-o', '--config', action='store_true')
     parser.add_argument('-v', '--version', action='version',
                         version=__version__)
     args = parser.parse_args()
@@ -57,6 +60,11 @@ def handle_commandline():
         print(build_cron_file())
     elif args.list:
         print(list_flows())
+    elif args.config:
+        config = get_config()
+        s = StringIO
+        config.write(s)
+        print(s.getvalue())
     elif args.flow:
         if type(args.present) is datetime:
             run_flow(flows[args.flow], args.present, True)

metrik/conf.py

@@ -1,16 +1,35 @@
 import os
 import sys
+# from six.moves.configparser import RawConfigParser
+from ConfigParser import RawConfigParser
 
-IS_TRAVIS = 'TRAVIS_BUILD_NUMBER' in os.environ
-IS_PYTEST = hasattr(sys, '_called_from_test')
 
-TEST = IS_PYTEST or IS_TRAVIS
+def get_config_locations():
+    return ['/etc/metrik', os.path.expanduser('~/.metrik')]
 
-USER_AGENT = "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:47.0) Gecko/20100101 Firefox/47.0"
-MONGO_HOST = 'localhost'
-MONGO_PORT = 27017
 
-if TEST:
-    MONGO_DATABASE = 'metrik'
-else:
-    MONGO_DATABASE = 'metrik-test'
+def get_default_conf():
+    cur_file_dir = os.path.dirname(os.path.realpath(__file__))
+    return open(cur_file_dir + '/default.conf', 'r')
+
+
+# Normally it's terrible practice to put static calls into the signature,
+# but this is safe (for now) since the get_config_locations() won't change
+# during a run. I.e. it starts up and that's the only time it's ever needed.
+def get_config(extra_locations=get_config_locations()):
+    config = RawConfigParser()
+
+    config.readfp(get_default_conf())
+
+    conf_files = config.read(extra_locations)
+    for conf_file in conf_files:
+        config.readfp(open(conf_file, 'r'))
+
+    # Not a huge fan of when developers think they're smarter than the
+    # end-users, but I'm calling it a special case for testing
+    is_travis = 'TRAVIS_BUILD_NUMBER' in os.environ
+    is_pytest = hasattr(sys, '_called_from_test')
+    if is_pytest or is_travis:
+        config.set('metrik', 'mongo_database', 'metrik-test')
+
+    return config

metrik/default.conf

@@ -0,0 +1,11 @@
+[metrik]
+user_agent=Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:47.0) Gecko/20100101 Firefox/47.0
+mongo_host=localhost
+mongo_port=27017
+mongo_database=metrik
+
+[tradeking]
+consumer_key=
+consumer_secret=
+oauth_token=
+oauth_token_secret=

metrik/targets/mongo.py

@@ -1,13 +1,18 @@
 from luigi import Target
 from pymongo import MongoClient
-from metrik.conf import MONGO_HOST, MONGO_PORT, MONGO_DATABASE
+from metrik.conf import get_config
 from datetime import datetime
 
 
 class MongoTarget(Target):
 
     def __init__(self, collection, id):
-        self.connection = MongoClient(MONGO_HOST, MONGO_PORT)[MONGO_DATABASE]
+        config = get_config()
+        self.connection = MongoClient(
+            host=config.get('metrik', 'mongo_host'),
+            port=config.get('metrik', 'mongo_port'))[
+            config.get('metrik', 'mongo_database')
+        ]
         self.collection = self.connection[collection]
         self.id = id
 

metrik/tasks/base.py

@@ -10,7 +10,7 @@ from pymongo import MongoClient
 
 from metrik.targets.mongo import MongoTarget
 from metrik.targets.noop import NoOpTarget
-from metrik.conf import MONGO_HOST, MONGO_PORT, MONGO_DATABASE
+from metrik.conf import get_config
 
 
 class MongoCreateTask(Task):
@@ -70,52 +70,50 @@ class MongoNoBackCreateTask(MongoCreateTask):
 class MongoRateLimit(object):
     rate_limit_collection = 'rate_limit'
 
-    def __init__(self, service, limit, interval, max_tries=5, backoff=.5):
-        """
-
-        :param present:
-        :type present: datetime.datetime
-        :param service:
-        :param limit:
-        :param interval:
-        :type interval: datetime.timedelta
-        :param max_tries:
-        :param backoff:
-        """
-        self.service = service
-        self.limit = limit
-        self.interval = interval
-        self.max_tries = max_tries
-        self.backoff = backoff
-        self.db = MongoClient(host=MONGO_HOST, port=MONGO_PORT)[MONGO_DATABASE]
+    def __init__(self):
+        config = get_config()
+        self.db = MongoClient(
+            host=config.get('metrik', 'mongo_host'),
+            port=config.getint('metrik', 'mongo_port'))[
+            config.get('metrik', 'mongo_database')
+        ]
 
     def get_present(self):
         return datetime.datetime.now()
 
-    def query_locks(self, present):
+    def query_locks(self, present, interval, service):
         return self.db[self.rate_limit_collection].find(
-            {'_created_at': {'$gt': present - self.interval},
-             'service': self.service}).count()
+            {'_created_at': {'$gt': present - interval},
+             'service': service}).count()
 
-    def save_lock(self, present):
+    def save_lock(self, present, service):
         self.db[self.rate_limit_collection].save({
-            '_created_at': present, 'service': self.service
+            '_created_at': present, 'service': service
         })
 
-    def sleep_until(self, present):
-        future_time = present + self.interval * self.backoff
+    def sleep_until(self, present, interval, backoff):
+        future_time = present + interval * backoff
         return (future_time - present).total_seconds()
 
-    def acquire_lock(self):
+    def acquire_lock(self, service, limit, interval, max_tries=5, backoff=.5):
         num_tries = 0
-        while num_tries < self.max_tries:
+        while num_tries < max_tries:
             num_tries += 1
-            num_locks = self.query_locks(self.get_present())
-            if num_locks < self.limit:
-                self.save_lock(self.get_present())
+            num_locks = self.query_locks(
+                self.get_present(),
+                interval,
+                service
+            )
+
+            if num_locks < limit:
+                self.save_lock(self.get_present(), service)
                 return True
-            elif num_tries < self.max_tries:
-                sleep_amount = self.sleep_until(self.get_present())
+            elif num_tries < max_tries:
+                sleep_amount = self.sleep_until(
+                    self.get_present(),
+                    interval,
+                    backoff
+                )
                 sleep(sleep_amount)
 
         return False

requirements.txt

@@ -5,4 +5,5 @@ requests>=2.9.1
 pymongo>=3.2
 python-dateutil>=2.4.2
 pandas>=0.17.1
-argparse>=1.1.0
+argparse>=1.1.0
+requests-oauthlib>=0.4.0

setup.py

@@ -18,7 +18,8 @@ setup(
         'pytz >= 2016.6.1',
         'python-dateutil >= 2.4.2',
         'pandas >= 0.17.1',
-        'argparse >= 1.1.0'
+        'argparse >= 1.1.0',
+        'requests-oauthlib >= 0.4.0'
     ],
     setup_requires=[
         'pytest_runner'

test/mongo_test.py

@@ -1,18 +1,22 @@
 from unittest import TestCase
 from pymongo import MongoClient
 
-from metrik.conf import MONGO_DATABASE, MONGO_PORT, MONGO_HOST
+from metrik.conf import get_config
 from metrik.targets.mongo import MongoTarget
 
 
 class MongoTest(TestCase):
     def setUp(self):
-        self.client = MongoClient(MONGO_HOST, MONGO_PORT)
-        self.db = self.client[MONGO_DATABASE]
+        config = get_config()
+        self.client = MongoClient(
+            host=config.get('metrik', 'mongo_host'),
+            port=config.getint('metrik', 'mongo_port'))
+        self.db = self.client[config.get('metrik', 'mongo_database')]
 
     def tearDown(self):
         super(MongoTest, self).tearDown()
-        self.client.drop_database(MONGO_DATABASE)
+        config = get_config()
+        self.client.drop_database(config.get('metrik', 'mongo_database'))
 
 
 class MongoTestTest(MongoTest):

test/tasks/test_base.py

@@ -15,42 +15,41 @@ class BaseTaskTest(TestCase):
 class RateLimitTest(MongoTest):
     def test_save_creates_record(self):
         service = 'testing_ratelimit'
+        interval = timedelta(seconds=1)
         assert self.db[MongoRateLimit.rate_limit_collection].count() == 0
 
         present = datetime.now()
         onesec_back = present - timedelta(seconds=1)
-        ratelimit = MongoRateLimit(
-            service, 1, timedelta(seconds=1)
-        )
-        assert ratelimit.query_locks(onesec_back) == 0
+        ratelimit = MongoRateLimit()
+        assert ratelimit.query_locks(onesec_back, interval, service) == 0
 
-        ratelimit.save_lock(present)
+        ratelimit.save_lock(present, service)
         assert self.db[MongoRateLimit.rate_limit_collection].count() == 1
-        assert ratelimit.query_locks(onesec_back) == 1
+        assert ratelimit.query_locks(onesec_back, interval, service) == 1
 
     def test_save_creates_correct_service(self):
         service_1 = 'testing_ratelimit_1'
         service_2 = 'testing_ratelimit_2'
+        interval = timedelta(seconds=1)
 
-        ratelimit1 = MongoRateLimit(
-            service_1, 1, timedelta(seconds=1)
-        )
-        ratelimit2 = MongoRateLimit(
-            service_2, 1, timedelta(seconds=1)
-        )
+        ratelimit = MongoRateLimit()
 
         present = datetime.now()
         assert self.db[MongoRateLimit.rate_limit_collection].count() == 0
-        assert ratelimit1.query_locks(present) == 0
-        assert ratelimit2.query_locks(present) == 0
+        assert ratelimit.query_locks(present, interval, service_1) == 0
+        assert ratelimit.query_locks(present, interval, service_2) == 0
 
-        ratelimit1.save_lock(present)
+        ratelimit.save_lock(present, service_1)
         assert self.db[MongoRateLimit.rate_limit_collection].count() == 1
-        assert ratelimit1.query_locks(present) == 1
-        assert ratelimit2.query_locks(present) == 0
+        assert ratelimit.query_locks(present, interval, service_1) == 1
+        assert ratelimit.query_locks(present, interval, service_2) == 0
 
     def test_acquire_lock_fails(self):
         service = 'testing_ratelimit'
+        limit = 1
+        interval = timedelta(seconds=1)
+        max_tries = 1
+        backoff = 10
 
         # The first scenario is as follows:
         # We try to acquire a lock with 1 try, backoff is 10.
@@ -60,19 +59,22 @@ class RateLimitTest(MongoTest):
         # lock and are only allowed one try.
         # Ultimately, we are testing that the 'fail immediately'
         # switch gets triggered correctly
-        ratelimit = MongoRateLimit(
-            service, 1, timedelta(seconds=1), max_tries=1, backoff=10
-        )
+        ratelimit = MongoRateLimit()
 
         start = datetime.now()
-        ratelimit.save_lock(start)
-        did_acquire = ratelimit.acquire_lock()
+        ratelimit.save_lock(start, service)
+        did_acquire = ratelimit.acquire_lock(service, limit, interval,
+                                             max_tries, backoff)
         end = datetime.now()
         assert not did_acquire
         assert (end - start).total_seconds() < 1
 
     def test_acquire_lock_succeeds(self):
         service = 'testing_ratelimit'
+        limit = 1
+        interval = timedelta(seconds=1)
+        max_tries = 2
+        backoff = 1
 
         # The first scenario is as follows:
         # We try to acquire a lock with two tries, backoff is 1.
@@ -80,15 +82,14 @@ class RateLimitTest(MongoTest):
         # thus when we try to acquire on the first try, we should fail.
         # However, the backoff should kick in, and we acquire successfully
         # on the second try
-        ratelimit = MongoRateLimit(
-            service, 1, timedelta(seconds=1), max_tries=2, backoff=1
-        )
+        ratelimit = MongoRateLimit()
 
         start = datetime.now()
-        ratelimit.save_lock(start - timedelta(seconds=.5))
-        did_acquire = ratelimit.acquire_lock()
+        ratelimit.save_lock(start - timedelta(seconds=.5), service)
+        did_acquire = ratelimit.acquire_lock(service, limit, interval,
+                                             max_tries, backoff)
         end = datetime.now()
         # Check that we acquired the lock
         assert did_acquire
         # Check that we only used one backoff period
-        assert (end - start).total_seconds() < 2
+        assert (end - start).total_seconds() < 2

test/test_conf.py

@@ -0,0 +1,10 @@
+from unittest import TestCase
+
+from metrik.conf import get_config
+
+
+class ConfigurationTest(TestCase):
+
+    def test_config_returns_values(self):
+        config = get_config([])
+        assert config is not None