Fix tests

Feature flags

example.py

@@ -2,19 +2,28 @@
 
 # Import the library
 import posthog
+import time
 
 # You can find this key on the /setup page in PostHog
-posthog.api_key = '<your key>'
+posthog.api_key = ''
+posthog.personal_api_key = ''
 
 # Where you host PostHog, with no trailing /.
 # You can remove this line if you're using posthog.com
-# posthog.host = 'http://127.0.0.1:8000' 
+posthog.host = 'http://127.0.0.1:8000'
 
 # Capture an event
 posthog.capture('distinct_id', 'event', {'property1': 'value', 'property2': 'value'})
 
+print(posthog.feature_enabled('beta-feature', 'distinct_id'))
+
+print('sleeping')
+time.sleep(45)
+
+print(posthog.feature_enabled('beta-feature', 'distinct_id'))
+
 # # Alias a previous distinct id with a new one
-# posthog.alias('distinct_id', 'new_distinct_id')
+posthog.alias('distinct_id', 'new_distinct_id')
 
 # # Add properties to the person
-# posthog.identify('distinct_id', {'email': 'something@something.com'})
+posthog.identify('distinct_id', {'email': 'something@something.com'})

posthog/__init__.py

@@ -13,6 +13,7 @@ debug = False # type: bool
 send = True # type: bool
 sync_mode = False # type: bool
 disabled = False # type: bool
+personal_api_key = None # type: str
 
 default_client = None
 
@@ -100,6 +101,25 @@ def alias(
     """
     _proxy('alias', previous_id=previous_id, distinct_id=distinct_id, context=context, timestamp=timestamp, message_id=message_id)
 
+def feature_enabled(
+    key, # type: str,
+    distinct_id, # type: str,
+    default=None, # type: Optional[Any]
+    ):
+    # type: (...) -> bool
+    """
+    Use feature flags to enable or disable features for users.
+
+    For example:
+    ```python
+    if posthog.feature_enabled('beta feature', 'distinct id'):
+        # do something
+    ```
+
+    You can call `posthog.load_feature_flags()` before to make sure you're not doing unexpected requests.
+    """
+    return _proxy('feature_enabled', key=key, distinct_id=distinct_id, default=default)
+
 
 def page(*args, **kwargs):
     """Send a page call."""
@@ -135,7 +155,7 @@ def _proxy(method, *args, **kwargs):
     if not default_client:
         default_client = Client(api_key, host=host, debug=debug,
                                 on_error=on_error, send=send,
-                                sync_mode=sync_mode)
+                                sync_mode=sync_mode, personal_api_key=personal_api_key)
 
     fn = getattr(default_client, method)
-    fn(*args, **kwargs)
+    return fn(*args, **kwargs)

posthog/client.py

@@ -1,16 +1,18 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from uuid import uuid4
 import logging
 import numbers
 import atexit
+import hashlib
 
 from dateutil.tz import tzutc
 from six import string_types
 
 from posthog.utils import guess_timezone, clean
 from posthog.consumer import Consumer
-from posthog.request import post
+from posthog.request import post, get, APIError
 from posthog.version import VERSION
+from posthog.poller import Poller
 
 try:
     import queue
@@ -19,6 +21,7 @@ except ImportError:
 
 
 ID_TYPES = (numbers.Number, string_types)
+__LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
 
 
 class Client(object):
@@ -28,10 +31,12 @@ class Client(object):
     def __init__(self, api_key=None, host=None, debug=False,
                  max_queue_size=10000, send=True, on_error=None, flush_at=100,
                  flush_interval=0.5, gzip=False, max_retries=3,
-                 sync_mode=False, timeout=15, thread=1):
+                 sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None):
         require('api_key', api_key, string_types)
 
         self.queue = queue.Queue(max_queue_size)
+        
+        # api_key: This should be the Team API Key (token), public
         self.api_key = api_key
         self.on_error = on_error
         self.debug = debug
@@ -40,6 +45,11 @@ class Client(object):
         self.host = host
         self.gzip = gzip
         self.timeout = timeout
+        self.feature_flags = None
+        self.poll_interval = poll_interval
+        
+        # personal_api_key: This should be a generated Personal API Key, private
+        self.personal_api_key = personal_api_key
 
         if debug:
             self.log.setLevel(logging.DEBUG)
@@ -216,6 +226,76 @@ class Client(object):
         self.flush()
         self.join()
 
+    def _load_feature_flags(self):
+        if not self.personal_api_key:
+            self.log.warning('[FEATURE FLAGS] You have to specify a personal_api_key to use feature flags.')
+            self.feature_flags = []
+            return
+
+        try:
+            self.feature_flags = get(self.personal_api_key, '/api/feature_flag/', self.host)['results']
+        except APIError as e:
+            if e.status == 401:
+                raise APIError(
+                    status=401, 
+                    message='You are using a write-only key with feature flags. ' \
+                    'To use feature flags, please set a personal_api_key ' \
+                    'More information: https://posthog.com/docs/api/overview'
+                )
+        except Exception as e:
+            self.log.warning('[FEATURE FLAGS] Fetching feature flags failed with following error. We will retry in %s seconds.' % self.poll_interval)
+            self.log.warning(e)
+
+        self._last_feature_flag_poll = datetime.utcnow().replace(tzinfo=tzutc())
+
+    def load_feature_flags(self):
+        self._load_feature_flags()
+        poller = Poller(interval=timedelta(seconds=self.poll_interval), execute=self._load_feature_flags)
+        poller.start()
+
+    def feature_enabled(self, key, distinct_id, default=False):
+        require('key', key, string_types)
+        require('distinct_id', distinct_id, ID_TYPES)
+        error = False
+
+        if not self.personal_api_key:
+            self.log.warning('[FEATURE FLAGS] You have to specify a personal_api_key to use feature flags.')
+        if not self.feature_flags:
+            self.load_feature_flags()
+
+        # If loading in previous line failed
+        if not self.feature_flags:
+            response = default
+            error = True
+        else:
+            try:
+                feature_flag = [flag for flag in self.feature_flags if flag['key'] == key][0]
+            except IndexError:
+                return default
+
+            if feature_flag.get('is_simple_flag'):
+                response = _hash(key, distinct_id) <= (feature_flag['rollout_percentage'] / 100)
+            else:
+                try:
+                    request = get(self.api_key, '/decide/', self.host, timeout=1)
+                    response = key in request['featureFlags']
+                except Exception as e:
+                    response = default
+                    self.log.warning('[FEATURE FLAGS] Unable to get data for flag %s, because of the following error:' % key)
+                    self.log.warning(e)
+                    error = True
+
+        self.capture(distinct_id, '$feature_flag_called', {'$feature_flag': key, '$feature_flag_response': response})
+        return response
+
+# This function takes a distinct_id and a feature flag key and returns a float between 0 and 1.
+# Given the same distinct_id and key, it'll always return the same float. These floats are
+# uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
+# we can do _hash(key, distinct_id) < 0.2
+def _hash(key, distinct_id):
+    hash_key = "%s.%s" % (key, distinct_id)
+    hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
+    return hash_val / __LONG_SCALE__
 
 def require(name, field, data_type):
     """Require that the named `field` has the right `data_type`"""

posthog/poller.py

@@ -0,0 +1,19 @@
+import threading
+
+class Poller(threading.Thread):
+    def __init__(self, interval, execute, *args, **kwargs):
+        threading.Thread.__init__(self)
+        self.daemon = False
+        self.stopped = threading.Event()
+        self.interval = interval
+        self.execute = execute
+        self.args = args
+        self.kwargs = kwargs
+        
+    def stop(self):
+        self.stopped.set()
+        self.join()
+
+    def run(self):
+        while not self.stopped.wait(self.interval.total_seconds()):
+            self.execute(*self.args, **self.kwargs)

posthog/request.py

@@ -4,27 +4,29 @@ import logging
 import json
 from gzip import GzipFile
 from requests.auth import HTTPBasicAuth
-from requests import sessions
+import requests
 from io import BytesIO
 
 from posthog.version import VERSION
 from posthog.utils import remove_trailing_slash
 
-_session = sessions.Session()
+_session = requests.sessions.Session()
 
+DEFAULT_HOST = 'https://app.posthog.com'
+USER_AGENT = 'posthog-python/' + VERSION
 
 def post(api_key, host=None, gzip=False, timeout=15, **kwargs):
     """Post the `kwargs` to the API"""
     log = logging.getLogger('posthog')
     body = kwargs
     body["sentAt"] = datetime.utcnow().replace(tzinfo=tzutc()).isoformat()
-    url = remove_trailing_slash(host or 'https://t.posthog.com') + '/batch/'
+    url = remove_trailing_slash(host or DEFAULT_HOST) + '/batch/'
     body['api_key'] = api_key
     data = json.dumps(body, cls=DatetimeSerializer)
     log.debug('making request: %s', data)
     headers = {
         'Content-Type': 'application/json',
-        'User-Agent': 'analytics-python/' + VERSION
+        'User-Agent': USER_AGENT
     }
     if gzip:
         headers['Content-Encoding'] = 'gzip'
@@ -45,21 +47,40 @@ def post(api_key, host=None, gzip=False, timeout=15, **kwargs):
     try:
         payload = res.json()
         log.debug('received response: %s', payload)
-        raise APIError(res.status_code, payload['code'], payload['message'])
+        raise APIError(res.status_code, payload['detail'])
     except ValueError:
-        raise APIError(res.status_code, 'unknown', res.text)
+        raise APIError(res.status_code, res.text)
+
+def get(api_key, url, host=None, timeout=None):
+    log = logging.getLogger('posthog')
+    url = remove_trailing_slash(host or DEFAULT_HOST) + url
+    response = requests.get(
+        url,
+        headers={
+            'Authorization': 'Bearer %s' % api_key,
+            'User-Agent': USER_AGENT
+        },
+        timeout=timeout
+    )
+    if response.status_code == 200:
+        return response.json()
+    try:
+        payload = response.json()
+        log.debug('received response: %s', payload)
+        raise APIError(response.status_code, payload['detail'])
+    except ValueError:
+        raise APIError(response.status_code, response.text)
 
 
 class APIError(Exception):
 
-    def __init__(self, status, code, message):
+    def __init__(self, status, message):
         self.message = message
         self.status = status
-        self.code = code
 
     def __str__(self):
-        msg = "[PostHog] {0}: {1} ({2})"
-        return msg.format(self.code, self.message, self.status)
+        msg = "[PostHog] {0} ({1})"
+        return msg.format(self.message, self.status)
 
 
 class DatetimeSerializer(json.JSONEncoder):

posthog/test/client.py

@@ -3,21 +3,24 @@ import unittest
 import six
 import mock
 import time
+from freezegun import freeze_time
 
 from posthog.version import VERSION
 from posthog.client import Client
 from posthog.test.utils import TEST_API_KEY
 
+from posthog.request import APIError
+
 class TestClient(unittest.TestCase):
 
-    def fail(self, e, batch):
+    def set_fail(self, e, batch):
         """Mark the failure handler"""
         print('FAIL', e, batch)
         self.failed = True
 
     def setUp(self):
         self.failed = False
-        self.client = Client(TEST_API_KEY, on_error=self.fail)
+        self.client = Client(TEST_API_KEY, on_error=self.set_fail)
 
     def test_requires_api_key(self):
         self.assertRaises(AssertionError, Client)
@@ -221,3 +224,92 @@ class TestClient(unittest.TestCase):
         client = Client(TEST_API_KEY)
         for consumer in client.consumers:
             self.assertEquals(consumer.timeout, 15)
+
+    @mock.patch('posthog.client.Poller')
+    @mock.patch('posthog.client.get')
+    def test_load_feature_flags(self, patch_get, patch_poll):
+        patch_get.return_value = {
+                'results': [{
+                    'id': 1,
+                    'name': 'Beta Feature',
+                    'key': 'beta-feature'
+                }]
+            }
+        client = Client(TEST_API_KEY, personal_api_key='test')
+        with freeze_time('2020-01-01T12:01:00.0000Z'):
+            client.load_feature_flags()
+        self.assertEqual(client.feature_flags[0]['key'], 'beta-feature')
+        self.assertEqual(client._last_feature_flag_poll.isoformat(), "2020-01-01T12:01:00+00:00")
+        self.assertEqual(patch_poll.call_count, 1)
+        
+    def test_load_feature_flags_wrong_key(self):
+        client = Client(TEST_API_KEY, personal_api_key=TEST_API_KEY)
+        with freeze_time('2020-01-01T12:01:00.0000Z'):
+            self.assertRaises(APIError, client.load_feature_flags)
+
+    @mock.patch('posthog.client.get')
+    def test_feature_enabled_simple(self, patch_get):
+        client = Client(TEST_API_KEY)
+        client.feature_flags = [{
+            'id': 1,
+            'name': 'Beta Feature',
+            'key': 'beta-feature',
+            'is_simple_flag': True,
+            'rollout_percentage': 100
+        }]
+        self.assertTrue(client.feature_enabled('beta-feature', 'distinct_id'))
+
+    @mock.patch('posthog.client.get')
+    def test_feature_enabled_request(self, patch_get):
+        patch_get.return_value = {
+            'featureFlags': ['beta-feature']
+        }
+        client = Client(TEST_API_KEY)
+        client.feature_flags = [{
+            'id': 1,
+            'name': 'Beta Feature',
+            'key': 'beta-feature',
+            'is_simple_flag': False,
+            'rollout_percentage': 100
+        }]
+        self.assertTrue(client.feature_enabled('beta-feature', 'distinct_id'))
+
+    @mock.patch('posthog.client.Poller')
+    @mock.patch('posthog.client.get')
+    def test_feature_enabled_doesnt_exist(self, patch_get, patch_poll):
+        client = Client(TEST_API_KEY, personal_api_key='test')
+        client.feature_flags = []
+
+        self.assertFalse(client.feature_enabled('doesnt-exist', 'distinct_id'))
+        self.assertTrue(client.feature_enabled('doesnt-exist', 'distinct_id', True))
+
+    @mock.patch('posthog.client.Poller')
+    @mock.patch('posthog.client.get')
+    def test_personal_api_key_doesnt_exist(self, patch_get, patch_poll):
+        client = Client(TEST_API_KEY)
+        client.feature_flags = []
+
+        self.assertFalse(client.feature_enabled('doesnt-exist', 'distinct_id'))
+        self.assertTrue(client.feature_enabled('doesnt-exist', 'distinct_id', True))
+
+    @mock.patch('posthog.client.Poller')
+    @mock.patch('posthog.client.get')
+    def test_load_feature_flags_error(self, patch_get, patch_poll):
+        def raise_effect():
+            raise Exception('http exception')
+        patch_get.return_value.raiseError.side_effect = raise_effect
+        client = Client(TEST_API_KEY, personal_api_key='test')
+        client.feature_flags = []
+
+        self.assertFalse(client.feature_enabled('doesnt-exist', 'distinct_id'))
+
+    @mock.patch('posthog.client.Poller')
+    @mock.patch('posthog.client.get')
+    def test_call_identify_fails(self, patch_get, patch_poll):
+        def raise_effect():
+            raise Exception('http exception')
+        patch_get.return_value.raiseError.side_effect = raise_effect
+        client = Client(TEST_API_KEY, personal_api_key='test')
+        client.feature_flags = [{'key': 'example', 'is_simple_flag': False}]
+
+        self.assertFalse(client.feature_enabled('example', 'distinct_id'))

posthog/test/consumer.py

@@ -142,16 +142,16 @@ class TestConsumer(unittest.TestCase):
         # we should retry on server errors
         consumer = Consumer(None, TEST_API_KEY)
         self._test_request_retry(consumer, APIError(
-            500, 'code', 'Internal Server Error'), 2)
+            500, 'Internal Server Error'), 2)
 
         # we should retry on HTTP 429 errors
         consumer = Consumer(None, TEST_API_KEY)
         self._test_request_retry(consumer, APIError(
-            429, 'code', 'Too Many Requests'), 2)
+            429, 'Too Many Requests'), 2)
 
         # we should NOT retry on other client errors
         consumer = Consumer(None, TEST_API_KEY)
-        api_error = APIError(400, 'code', 'Client Errors')
+        api_error = APIError(400, 'Client Errors')
         try:
             self._test_request_retry(consumer, api_error, 1)
         except APIError:
@@ -162,7 +162,7 @@ class TestConsumer(unittest.TestCase):
         # test for number of exceptions raise > retries value
         consumer = Consumer(None, TEST_API_KEY, retries=3)
         self._test_request_retry(consumer, APIError(
-            500, 'code', 'Internal Server Error'), 3)
+            500, 'Internal Server Error'), 3)
 
     def test_pause(self):
         consumer = Consumer(None, TEST_API_KEY)

setup.py

@@ -24,7 +24,8 @@ install_requires = [
 ]
 
 tests_require = [
-    "mock>=2.0.0"
+    "mock>=2.0.0",
+    "freezegun==0.3.15"
 ]
 
 setup(