22import copy
33import mock
44import unittest
5+ from collections import defaultdict
56
67# 3p
78from nose .plugins .attrib import attr
1112from utils .service_discovery .config_stores import get_config_store
1213from utils .service_discovery .consul_config_store import ConsulStore
1314from utils .service_discovery .etcd_config_store import EtcdStore
14- from utils .service_discovery .abstract_config_store import AbstractConfigStore , CONFIG_FROM_KUBE
15+ from utils .service_discovery .abstract_config_store import AbstractConfigStore , \
16+ _TemplateCache , CONFIG_FROM_KUBE , CONFIG_FROM_TEMPLATE , CONFIG_FROM_AUTOCONF
1517from utils .service_discovery .sd_backend import get_sd_backend
1618from utils .service_discovery .sd_docker_backend import SDDockerBackend , _SDDockerBackendConfigFetchState
1719
@@ -66,11 +68,11 @@ def client_read(path, **kwargs):
6668 if 'all' in kwargs :
6769 return {}
6870 else :
69- return TestServiceDiscovery .mock_tpls .get (image )[0 ][config_parts .index (config_part )]
71+ return TestServiceDiscovery .mock_raw_templates .get (image )[0 ][config_parts .index (config_part )]
7072
7173
7274def issue_read (identifier ):
73- return TestServiceDiscovery .mock_tpls .get (identifier )
75+ return TestServiceDiscovery .mock_raw_templates .get (identifier )
7476
7577@attr ('unix' )
7678class TestServiceDiscovery (unittest .TestCase ):
@@ -122,7 +124,7 @@ class TestServiceDiscovery(unittest.TestCase):
122124 }
123125
124126 # raw templates coming straight from the config store
125- mock_tpls = {
127+ mock_raw_templates = {
126128 # image_name: ('[check_name]', '[init_tpl]', '[instance_tpl]', expected_python_tpl_list)
127129 'image_0' : (
128130 ('["check_0"]' , '[{}]' , '[{"host": "%%host%%"}]' ),
@@ -509,12 +511,12 @@ def test_fill_tpl(self, *args):
509511 def test_get_auto_config (self ):
510512 """Test _get_auto_config"""
511513 expected_tpl = {
512- 'redis' : ('redisdb' , None , {"host" : "%%host%%" , "port" : "%%port%%" }),
513- 'consul' : ('consul' , None , {
514- "url" : "http://%%host%%:%%port%%" , "catalog_checks" : True , "new_leader_checks" : True
515- }) ,
516- 'redis:v1' : ('redisdb' , None , {"host" : "%%host%%" , "port" : "%%port%%" }),
517- 'foobar' : None
514+ 'redis' : [ ('redisdb' , None , {"host" : "%%host%%" , "port" : "%%port%%" })] ,
515+ 'consul' : [ ('consul' , None , {
516+ "url" : "http://%%host%%:%%port%%" , "catalog_checks" : True , "new_leader_checks" : True
517+ })] ,
518+ 'redis:v1' : [ ('redisdb' , None , {"host" : "%%host%%" , "port" : "%%port%%" })] ,
519+ 'foobar' : []
518520 }
519521
520522 config_store = get_config_store (self .auto_conf_agentConfig )
@@ -529,10 +531,10 @@ def test_get_check_tpls(self, mock_client_read):
529531 invalid_config = ['bad_image_0' , 'bad_image_1' ]
530532 config_store = get_config_store (self .auto_conf_agentConfig )
531533 for image in valid_config :
532- tpl = self .mock_tpls .get (image )[1 ]
534+ tpl = self .mock_raw_templates .get (image )[1 ]
533535 self .assertEquals (tpl , config_store .get_check_tpls (image ))
534536 for image in invalid_config :
535- tpl = self .mock_tpls .get (image )[1 ]
537+ tpl = self .mock_raw_templates .get (image )[1 ]
536538 self .assertEquals (tpl , config_store .get_check_tpls (image ))
537539
538540 @mock .patch .object (AbstractConfigStore , 'client_read' , side_effect = client_read )
@@ -542,7 +544,7 @@ def test_get_check_tpls_kube(self, mock_client_read):
542544 invalid_config = ['bad_image_0' ]
543545 config_store = get_config_store (self .auto_conf_agentConfig )
544546 for image in valid_config + invalid_config :
545- tpl = self .mock_tpls .get (image )[1 ]
547+ tpl = self .mock_raw_templates .get (image )[1 ]
546548 tpl = [(CONFIG_FROM_KUBE , t [1 ]) for t in tpl ]
547549 if tpl :
548550 self .assertNotEquals (
@@ -558,7 +560,7 @@ def test_get_check_tpls_kube(self, mock_client_read):
558560 ['service-discovery.datadoghq.com/foo.check_names' ,
559561 'service-discovery.datadoghq.com/foo.init_configs' ,
560562 'service-discovery.datadoghq.com/foo.instances' ],
561- self .mock_tpls [image ][0 ]))))
563+ self .mock_raw_templates [image ][0 ]))))
562564
563565 def test_get_config_id (self ):
564566 """Test get_config_id"""
@@ -570,8 +572,8 @@ def test_get_config_id(self):
570572 expected_ident )
571573 clear_singletons (self .auto_conf_agentConfig )
572574
573- @mock .patch .object (AbstractConfigStore , '_issue_read' , side_effect = issue_read )
574- def test_read_config_from_store (self , issue_read ):
575+ @mock .patch .object (_TemplateCache , '_issue_read' , side_effect = issue_read )
576+ def test_read_config_from_store (self , args ):
575577 """Test read_config_from_store"""
576578 valid_idents = [('nginx' , 'nginx' ), ('nginx:latest' , 'nginx:latest' ),
577579 ('custom-nginx' , 'custom-nginx' ), ('custom-nginx:latest' , 'custom-nginx' ),
@@ -582,7 +584,13 @@ def test_read_config_from_store(self, issue_read):
582584 for ident , expected_key in valid_idents :
583585 tpl = config_store .read_config_from_store (ident )
584586 # source is added after reading from the store
585- self .assertEquals (tpl , ('template' ,) + self .mock_tpls .get (expected_key ))
587+ self .assertEquals (
588+ tpl ,
589+ {
590+ CONFIG_FROM_AUTOCONF : None ,
591+ CONFIG_FROM_TEMPLATE : self .mock_raw_templates .get (expected_key )
592+ }
593+ )
586594 for ident in invalid_idents :
587595 self .assertEquals (config_store .read_config_from_store (ident ), [])
588596
@@ -600,3 +608,72 @@ def test_read_jmx_config_from_store(self, *args):
600608 for check in self .jmx_sd_configs :
601609 key = '{}_0' .format (check )
602610 self .assertEquals (jmx_configs [key ], valid_configs [key ])
611+
612+ # Template cache
613+ @mock .patch ('utils.service_discovery.abstract_config_store.get_auto_conf_images' )
614+ def test_populate_auto_conf (self , mock_get_auto_conf_images ):
615+ """test _populate_auto_conf"""
616+ auto_tpls = {
617+ 'foo' : [['check0' , 'check1' ], [{}, {}], [{}, {}]],
618+ 'bar' : [['check2' , 'check3' , 'check3' ], [{}, {}, {}], [{}, {'foo' : 'bar' }, {'bar' : 'foo' }]],
619+ }
620+ cache = _TemplateCache (issue_read , '' )
621+ cache .auto_conf_templates = defaultdict (lambda : [[]] * 3 )
622+ mock_get_auto_conf_images .return_value = auto_tpls
623+
624+ cache ._populate_auto_conf ()
625+ self .assertEquals (cache .auto_conf_templates ['foo' ], auto_tpls ['foo' ])
626+ self .assertEquals (cache .auto_conf_templates ['bar' ],
627+ [['check2' , 'check3' ], [{}, {}], [{}, {'foo' : 'bar' }]])
628+
629+ @mock .patch .object (_TemplateCache , '_issue_read' , return_value = None )
630+ def test_get_templates (self , args ):
631+ """test get_templates"""
632+ kv_tpls = {
633+ 'foo' : [['check0' , 'check1' ], [{}, {}], [{}, {}]],
634+ 'bar' : [['check2' , 'check3' ], [{}, {}], [{}, {}]],
635+ }
636+ auto_tpls = {
637+ 'foo' : [['check3' , 'check5' ], [{}, {}], [{}, {}]],
638+ 'bar' : [['check2' , 'check6' ], [{}, {}], [{}, {}]],
639+ 'foobar' : [['check4' ], [{}], [{}]],
640+ }
641+ cache = _TemplateCache (issue_read , '' )
642+ cache .kv_templates = kv_tpls
643+ cache .auto_conf_templates = auto_tpls
644+ self .assertEquals (cache .get_templates ('foo' ),
645+ {CONFIG_FROM_TEMPLATE : [['check0' , 'check1' ], [{}, {}], [{}, {}]],
646+ CONFIG_FROM_AUTOCONF : [['check3' , 'check5' ], [{}, {}], [{}, {}]]}
647+ )
648+
649+ self .assertEquals (cache .get_templates ('bar' ),
650+ # check3 must come from template not autoconf
651+ {CONFIG_FROM_TEMPLATE : [['check2' , 'check3' ], [{}, {}], [{}, {}]],
652+ CONFIG_FROM_AUTOCONF : [['check6' ], [{}], [{}]]}
653+ )
654+
655+ self .assertEquals (cache .get_templates ('foobar' ),
656+ {CONFIG_FROM_TEMPLATE : None ,
657+ CONFIG_FROM_AUTOCONF : [['check4' ], [{}], [{}]]}
658+ )
659+
660+ self .assertEquals (cache .get_templates ('baz' ), None )
661+
662+ def test_get_check_names (self ):
663+ """Test get_check_names"""
664+ kv_tpls = {
665+ 'foo' : [['check0' , 'check1' ], [{}, {}], [{}, {}]],
666+ 'bar' : [['check2' , 'check3' ], [{}, {}], [{}, {}]],
667+ }
668+ auto_tpls = {
669+ 'foo' : [['check4' , 'check5' ], [{}, {}], [{}, {}]],
670+ 'bar' : [['check2' , 'check6' ], [{}, {}], [{}, {}]],
671+ 'foobar' : None ,
672+ }
673+ cache = _TemplateCache (issue_read , '' )
674+ cache .kv_templates = kv_tpls
675+ cache .auto_conf_templates = auto_tpls
676+ self .assertEquals (cache .get_check_names ('foo' ), set (['check0' , 'check1' , 'check4' , 'check5' ]))
677+ self .assertEquals (cache .get_check_names ('bar' ), set (['check2' , 'check3' , 'check6' ]))
678+ self .assertEquals (cache .get_check_names ('foobar' ), set ())
679+ self .assertEquals (cache .get_check_names ('baz' ), set ())
0 commit comments