Skip to content
This repository has been archived by the owner on Oct 28, 2020. It is now read-only.

Commit

Permalink
Add vulcand trust forwarded header functionality and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@a5huynh
a5huynh committed Feb 27, 2016
1 parent e0e20f3 commit 9203803
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 17 deletions.
25 changes: 16 additions & 9 deletions sidekick.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 )
Expand All @@ -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:
Expand Down Expand Up @@ -263,7 +269,8 @@ def main():
etcd_client,
args.timeout,
args.ttl,
args.vulcand )
args.vulcand,
args )

if __name__ == '__main__':
main()
8 changes: 4 additions & 4 deletions tests/backends/test_vulcand.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# @Author: ahuynh
# @Date: 2016-02-11 14:28:02
# @Last Modified by: ahuynh
# @Last Modified time: 2016-02-11 17:27:07
# @Last Modified time: 2016-02-26 16:15:00
import unittest

from sidekick import announce_services, find_matching_container, parse_args
Expand Down Expand Up @@ -47,17 +47,17 @@ def test_vulcand_announce( self ):

# Successful health check
with patch( 'sidekick.check_health', return_value=True ):
announce_services( services.items(), 'test', self.etcd_client, 0, 0, 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 )
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 )
announce_services( services.items(), 'test', self.etcd_client, 0, 0, True, self.args )
self.assertEquals( str(mock_method.call_args[0][0]), 'Test Exception' )
8 changes: 4 additions & 4 deletions tests/test_sidekick.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# @Author: ahuynh
# @Date: 2015-06-18 20:15:30
# @Last Modified by: ahuynh
# @Last Modified time: 2016-02-11 17:27:50
# @Last Modified time: 2016-02-26 16:14:09
import unittest

from sidekick import announce_services, check_name, find_matching_container
Expand Down Expand Up @@ -41,18 +41,18 @@ def test_announce_services( self ):

# Test successful health check
with patch( 'sidekick.check_health', return_value=True ):
announce_services( services.items(), 'test', self.etcd_client, 0, 0, False )
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 )
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 )
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 ):
Expand Down

0 comments on commit 9203803

Please sign in to comment.