Merge pull request #5 from a5huynh/feature/vulcand-trust-headers

Vulcand Tests & Flag to trust headers

.gitignore

@@ -1,2 +1,3 @@
 .coverage
 htmlcov
+__pycache__

.travis.yml

@@ -1,6 +1,7 @@
 language: python
 python:
   - "3.4"
+  - "3.5"
 
 # install dependencies
 install: "pip install -r requirements.txt"

sidekick.py

@@ -3,7 +3,7 @@
 # @Author: ahuynh
 # @Date:   2015-06-10 16:51:36
 # @Last Modified by:   ahuynh
-# @Last Modified time: 2016-02-11 16:36:44
+# @Last Modified time: 2016-02-26 16:24:21
 '''
     The sidekick should essentially replace job of the following typical
     bash script that is used to announce a service to ETCD.
@@ -79,13 +79,16 @@ def parse_args( args ):
     parser.add_argument( '--vulcand', action='store', type=bool, default=False,
                          help='Selector for LB')
 
+    parser.add_argument( '--vulcand-trust-forwarded-headers', action='store',
+                         type=bool, default=False, help='Trust X-Forwarded-* headers' )
+
     parser.add_argument( '--type', action='store', default='http',
                          help='type for Vulcand')
 
     return parser.parse_args( args )
 
 
-def announce_services( services, etcd_folder, etcd_client, timeout, ttl, vulcand ):
+def announce_services( services, etcd_folder, etcd_client, timeout, ttl, vulcand, args ):
     for key, value in services:
         logger.info( 'Health check for {}'.format( key ) )
         healthy = check_health( value )
@@ -103,15 +106,18 @@ def announce_services( services, etcd_folder, etcd_client, timeout, ttl, vulcand
                     etcd_client.delete( server )
                     etcd_client.delete( frontend )
                 else:
+                    fend_data = {
+                        'Type': value['type'],
+                        'BackendId': value['domain'],
+                        'Route': 'Host(`{0}`)'.format( value['domain'] ),
+                    }
+                    # Should we trust the headers being forwared to vulcand?
+                    if args.vulcand_trust_forwarded_headers:
+                        fend_data[ 'Settings' ] = { 'TrustForwardHeader': True }
                     # Announce this server to ETCD
                     etcd_client.write( backend, json.dumps({ 'Type': value['type'] }), ttl=ttl)
                     etcd_client.write( server, json.dumps({ 'URL': 'http://{uri}'.format( **value ) }), ttl=ttl)
-                    etcd_client.write( frontend, json.dumps({
-                        'Type': value['type'],
-                        'BackendId': value['domain'],
-                        'Route': 'Host(`{0}`)'.format( value['domain'] )
-                    }), ttl=ttl)
-
+                    etcd_client.write( frontend, json.dumps( fend_data ), ttl=ttl)
             except etcd.EtcdException as e:
                 logging.error( e )
         else:
@@ -263,7 +269,8 @@ def main():
                            etcd_client,
                            args.timeout,
                            args.ttl,
-                           args.vulcand )
+                           args.vulcand,
+                           args )
 
 if __name__ == '__main__':
     main()

tests/__init__.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Author: ahuynh
+# @Date:   2016-01-28 10:04:58
+# @Last Modified by:   ahuynh
+# @Last Modified time: 2016-02-11 17:25:54
+import etcd
+
+
+class MockEtcd( object ):
+
+    def __init__( self, raise_exception=False ):
+        self._reset()
+        self.raise_exception = raise_exception
+
+    def _check_raise( self ):
+        if self.raise_exception:
+            raise etcd.EtcdException( 'Test Exception' )
+
+    def _reset( self ):
+        self.deleted = set([])
+        self.written = {}
+
+    def delete( self, key ):
+        self._check_raise()
+        self.deleted.add( key )
+
+    def write( self, key, value, ttl ):
+        self._check_raise()
+        self.written[ key ] = value

tests/backends/test_vulcand.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Author: ahuynh
+# @Date:   2016-02-11 14:28:02
+# @Last Modified by:   ahuynh
+# @Last Modified time: 2016-02-26 16:15:00
+import unittest
+
+from sidekick import announce_services, find_matching_container, parse_args
+from tests import MockEtcd
+from unittest.mock import patch
+
+
+class TestVulcandBackend( unittest.TestCase ):
+
+    def setUp( self ):
+
+        self.args = parse_args([
+            '--name', 'test',
+            '--ip', 'localhost',
+            '--check-ip', '0.0.0.0',
+            '--vulcand', 'True'
+        ])
+
+        self.etcd_client = MockEtcd()
+        self.container = {
+            'Image': 'image:latest',
+            'Ports': [{
+                'PrivatePort': 9200,
+                'IP': '0.0.0.0',
+                'Type': 'tcp',
+                'PublicPort': 9200
+            }, {
+                'PrivatePort': 9300,
+                'IP': '0.0.0.0',
+                'Type': 'tcp',
+                'PublicPort': 9300
+            }],
+            'Created': 1427906382,
+            'Names': ['/test'],
+            'Status': 'Up 2 days'
+        }
+
+    def test_vulcand_announce( self ):
+        """ Test `announce_services` functionality """
+        services = find_matching_container( [ self.container ], self.args )
+
+        # Successful health check
+        with patch( 'sidekick.check_health', return_value=True ):
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, True, self.args )
+            self.assertEqual( len( self.etcd_client.written.keys() ), 4 )
+
+        # Unsuccessful health check
+        self.etcd_client._reset()
+        with patch( 'sidekick.check_health', return_value=False ):
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, True, self.args )
+            self.assertEqual( len( self.etcd_client.deleted ), 4 )
+
+        # Correct etcd exception handling
+        self.etcd_client = MockEtcd( raise_exception=True )
+        with patch( 'logging.error' ) as mock_method:
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, True, self.args )
+            self.assertEquals( str(mock_method.call_args[0][0]), 'Test Exception' )

tests/test_sidekick.py

@@ -3,38 +3,20 @@
 # @Author: ahuynh
 # @Date:   2015-06-18 20:15:30
 # @Last Modified by:   ahuynh
-# @Last Modified time: 2016-02-11 16:23:22
+# @Last Modified time: 2016-02-26 16:14:09
 import unittest
 
-from collections import namedtuple
 from sidekick import announce_services, check_name, find_matching_container
-from sidekick import check_health, public_ports
-
-# Used to test command line arguments
-Args = namedtuple('Args', ['name', 'ip', 'check_ip', 'domain', 'vulcand', 'type'])
-
-
-class MockEtcd( object ):
-    def delete( self, value ):
-        pass
-
-    def write( self, value, ttl ):
-        pass
+from sidekick import check_health, public_ports, parse_args
+from tests import MockEtcd
+from unittest.mock import patch
 
 
 class TestSidekick( unittest.TestCase ):
 
     def setUp( self ):
 
-        self.args = Args(
-            name='test',
-            ip='localhost',
-            check_ip='0.0.0.0',
-            domain='example.com',
-            vulcand=False,
-            type='http'
-        )
-
+        self.args = parse_args( [ '--name', 'test', '--ip', 'localhost', '--check-ip', '0.0.0.0' ] )
         self.etcd_client  = MockEtcd()
 
         self.container = {
@@ -50,26 +32,51 @@ def setUp( self ):
                 'PublicPort': 9300}],
             'Created': 1427906382,
             'Names': ['/test'],
-            'Status': 'Up 2 days'}
+            'Status': 'Up 2 days'
+        }
 
     def test_announce_services( self ):
-        ''' Test `announce_services` functionality '''
+        """ Test `announce_services` functionality """
         services = find_matching_container( [self.container], self.args )
-        announce_services( services.items(), 'test', self.etcd_client, 0, 0, False )
+
+        # Test successful health check
+        with patch( 'sidekick.check_health', return_value=True ):
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, False, self.args )
+            self.assertEqual( len( self.etcd_client.written.keys() ), 2 )
+
+        # Test unsuccessful health check
+        with patch( 'sidekick.check_health', return_value=False ):
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, False, self.args )
+            self.assertEqual( len( self.etcd_client.deleted ), 2 )
+
+        # Test correct etcd exception handling
+        self.etcd_client = MockEtcd( raise_exception=True )
+        with patch( 'logging.error' ) as mock_method:
+            announce_services( services.items(), 'test', self.etcd_client, 0, 0, False, self.args )
+            self.assertEquals( str(mock_method.call_args[0][0]), 'Test Exception' )
 
     def test_check_health( self ):
-        ''' Test `check_health` functionality '''
+        """ Test `check_health` functionality """
         results = find_matching_container( [self.container], self.args )
+
         for value in results.values():
+            # Unsuccessful socket connection
             self.assertFalse( check_health( value ) )
 
+            # Successful socket connection
+            with patch( 'socket.socket.connect' ) as mock_method:
+                check_health( value )
+                args = mock_method.call_args[0][0]
+                self.assertEqual( args[0], value[ 'check_ip' ] )
+                self.assertEqual( args[1], value[ 'port' ] )
+
     def test_check_name( self ):
-        ''' Test `check_name` functionality '''
+        """ Test `check_name` functionality """
         self.assertTrue( check_name( self.container, 'test' ) )
         self.assertFalse( check_name( self.container, '/test' ) )
 
     def test_find_matching_container( self ):
-        ''' Test `find_matching_container` functionality '''
+        """ Test `find_matching_container` functionality """
         # Test a successful match
         results = find_matching_container( [self.container], self.args )
         self.assertEqual( len( results.items() ), 2 )
@@ -87,5 +94,5 @@ def test_find_matching_container( self ):
             find_matching_container( [no_open_ports], self.args )
 
     def test_public_ports( self ):
-        ''' Test `public_ports` functionality '''
+        """ Test `public_ports` functionality """
         self.assertEquals( len( public_ports( self.container ) ), 2 )