class TrackerTest(unittest.TestCase): def setUp(self): mock_config = Mock() mock_config.validate.return_value = True self.tracker = Tracker(mock_config) # pylint: disable=unused-argument @patch.object(Tracker, 'get_stmgr_topologies') @patch.object(Tracker, 'remove_topology') @patch.object(Tracker, 'add_new_topology') @patch.object(statemanagerfactory, 'get_all_state_managers') def test_first_synch_topologies(self, mock_get_all_state_managers, mock_add_new_topology, mock_remove_topology, mock_get_topologies_for_state_location): mock_state_manager_1 = Mock() mock_state_manager_1.name = 'mock_name1' mock_state_manager_2 = Mock() mock_state_manager_2.name = 'mock_name2' watches = {} mock_get_all_state_managers.return_value = [ mock_state_manager_1, mock_state_manager_2 ] mock_get_topologies_for_state_location.return_value = [] def side_effect1(on_topologies_watch): watches["1"] = on_topologies_watch on_topologies_watch(['top_name1', 'top_name2']) mock_state_manager_1.get_topologies = side_effect1 def side_effect2(on_topologies_watch): watches["2"] = on_topologies_watch on_topologies_watch(['top_name3', 'top_name4']) mock_state_manager_2.get_topologies = side_effect2 self.tracker.synch_topologies() mock_get_topologies_for_state_location.assert_has_calls( [call("mock_name2"), call("mock_name1")], any_order=True) mock_add_new_topology.assert_has_calls([ call(mock_state_manager_1, 'top_name1'), call(mock_state_manager_1, 'top_name2'), call(mock_state_manager_2, 'top_name3'), call(mock_state_manager_2, 'top_name4') ], any_order=True) @patch.object(Tracker, 'get_stmgr_topologies') @patch.object(Tracker, 'remove_topology') @patch.object(Tracker, 'add_new_topology') @patch.object(statemanagerfactory, 'get_all_state_managers') def test_synch_topologies_leading_with_add_and_remove_topologies( self, mock_get_all_state_managers, mock_add_new_topology, mock_remove_topology, mock_get_topologies_for_state_location): mock_state_manager_1 = Mock() mock_state_manager_1.name = 'mock_name1' mock_state_manager_2 = Mock() mock_state_manager_2.name = 'mock_name2' watches = {} mock_get_all_state_managers.return_value = [ mock_state_manager_1, mock_state_manager_2 ] mock_get_topologies_for_state_location.return_value = [] def side_effect1(on_topologies_watch): watches["1"] = on_topologies_watch on_topologies_watch(['top_name1', 'top_name2']) mock_state_manager_1.get_topologies = side_effect1 def side_effect2(on_topologies_watch): watches["2"] = on_topologies_watch on_topologies_watch(['top_name3', 'top_name4']) mock_state_manager_2.get_topologies = side_effect2 self.tracker.synch_topologies() mock_get_topologies_for_state_location.assert_has_calls( [call("mock_name2"), call("mock_name1")], any_order=True) mock_add_new_topology.assert_has_calls([ call(mock_state_manager_1, 'top_name1'), call(mock_state_manager_1, 'top_name2'), call(mock_state_manager_2, 'top_name3'), call(mock_state_manager_2, 'top_name4') ], any_order=True) self.assertEqual(4, mock_add_new_topology.call_count) self.assertEqual(0, mock_remove_topology.call_count) mock_get_topologies_for_state_location.reset_mock() mock_add_new_topology.reset_mock() mock_remove_topology.reset_mock() def get_topologies_for_state_location_side_effect(name): if name == 'mock_name1': return [ Topology('top_name1', 'mock_name1'), Topology('top_name2', 'mock_name1') ] if name == 'mock_name2': return [ Topology('top_name3', 'mock_name2'), Topology('top_name4', 'mock_name2') ] return [] # pylint: disable=line-too-long mock_get_topologies_for_state_location.side_effect = get_topologies_for_state_location_side_effect watches["1"](['top_name1', 'top_name3']) watches["2"](['top_name5', 'top_name6']) mock_add_new_topology.assert_has_calls([ call(mock_state_manager_1, 'top_name3'), call(mock_state_manager_2, 'top_name5'), call(mock_state_manager_2, 'top_name6') ], any_order=True) mock_remove_topology.assert_has_calls([ call('top_name2', 'mock_name1'), call('top_name3', 'mock_name2'), call('top_name4', 'mock_name2') ], any_order=False) self.assertEqual(3, mock_add_new_topology.call_count) self.assertEqual(3, mock_remove_topology.call_count) def fill_tracker_topologies(self): def create_mock_execution_state(cluster, role, environ): estate = protoEState.ExecutionState() estate.cluster = cluster estate.role = role estate.environ = environ return estate self.topology1 = Topology('top_name1', 'mock_name1') self.topology1.execution_state = create_mock_execution_state( 'cluster1', 'mark', 'env1') self.topology2 = Topology('top_name2', 'mock_name1') self.topology2.execution_state = create_mock_execution_state( 'cluster1', 'bob', 'env1') self.topology3 = Topology('top_name3', 'mock_name1') self.topology3.execution_state = create_mock_execution_state( 'cluster1', 'tom', 'env2') self.topology4 = Topology('top_name4', 'mock_name2') self.topology4.execution_state = create_mock_execution_state( 'cluster2', 'x', 'env1') self.topology5 = Topology('top_name5', 'mock_name2') self.topology5.execution_state = create_mock_execution_state( 'cluster2', 'x', 'env2') self.tracker.topologies = [ self.topology1, self.topology2, self.topology3, self.topology4, self.topology5 ] # pylint: disable=line-too-long def test_get_topology_by_cluster_environ_and_name(self): self.fill_tracker_topologies() self.assertEqual( self.topology1, self.tracker.get_topology('cluster1', 'mark', 'env1', 'top_name1')) self.assertEqual( self.topology1, self.tracker.get_topology('cluster1', None, 'env1', 'top_name1')) self.assertEqual( self.topology2, self.tracker.get_topology('cluster1', 'bob', 'env1', 'top_name2')) self.assertEqual( self.topology2, self.tracker.get_topology('cluster1', None, 'env1', 'top_name2')) self.assertEqual( self.topology3, self.tracker.get_topology('cluster1', 'tom', 'env2', 'top_name3')) self.assertEqual( self.topology3, self.tracker.get_topology('cluster1', None, 'env2', 'top_name3')) self.assertEqual( self.topology4, self.tracker.get_topology('cluster2', None, 'env1', 'top_name4')) self.assertEqual( self.topology5, self.tracker.get_topology('cluster2', None, 'env2', 'top_name5')) def test_get_topolies_for_state_location(self): self.fill_tracker_topologies() self.assertCountEqual([self.topology1, self.topology2, self.topology3], self.tracker.get_stmgr_topologies('mock_name1')) self.assertCountEqual([self.topology4, self.topology5], self.tracker.get_stmgr_topologies('mock_name2')) def test_add_new_topology(self): self.assertCountEqual([], self.tracker.topologies) mock_state_manager_1 = Mock() mock_state_manager_1.name = 'mock_name1' self.tracker.add_new_topology(mock_state_manager_1, 'top_name1') self.assertCountEqual(['top_name1'], [t.name for t in self.tracker.topologies]) self.tracker.add_new_topology(mock_state_manager_1, 'top_name2') self.assertCountEqual(['top_name1', 'top_name2'], [t.name for t in self.tracker.topologies]) self.assertEqual(2, mock_state_manager_1.get_pplan.call_count) self.assertEqual(2, mock_state_manager_1.get_execution_state.call_count) self.assertEqual(2, mock_state_manager_1.get_tmanager.call_count) def test_remove_topology(self): self.fill_tracker_topologies() self.tracker.remove_topology('top_name1', 'mock_name1') self.assertCountEqual( [self.topology2, self.topology3, self.topology4, self.topology5], self.tracker.topologies) self.tracker.remove_topology('top_name2', 'mock_name1') self.assertCountEqual([self.topology3, self.topology4, self.topology5], self.tracker.topologies) # Removing one that is not there should not have any affect self.tracker.remove_topology('top_name8', 'mock_name1') self.assertCountEqual([self.topology3, self.topology4, self.topology5], self.tracker.topologies) self.tracker.remove_topology('top_name4', 'mock_name2') self.assertCountEqual([self.topology3, self.topology5], self.tracker.topologies) def test_extract_physical_plan(self): # Create topology pb_pplan = MockProto().create_mock_simple_physical_plan() topology = Topology('topology_name', 'state_manager') topology.set_physical_plan(pb_pplan) # Extract physical plan pplan = self.tracker.extract_physical_plan(topology) # Mock topology doesn't have topology config and instances self.assertEqual(pplan['config'], {}) self.assertEqual(pplan['bolts'], {'mock_bolt': []}) self.assertEqual(pplan['spouts'], {'mock_spout': []}) self.assertEqual(pplan['components']['mock_bolt']['config'], {'topology.component.parallelism': '1'}) self.assertEqual(pplan['components']['mock_spout']['config'], {'topology.component.parallelism': '1'}) self.assertEqual(pplan['instances'], {}) self.assertEqual(pplan['stmgrs'], {}) def test_extract_packing_plan(self): # Create topology pb_pplan = MockProto().create_mock_simple_packing_plan() topology = Topology('topology_name', 'ExclamationTopology') topology.set_packing_plan(pb_pplan) # Extract packing plan packing_plan = self.tracker.extract_packing_plan(topology) self.assertEqual(packing_plan['id'], 'ExclamationTopology') self.assertEqual(packing_plan['container_plans'][0]['id'], 1) self.assertEqual( packing_plan['container_plans'][0]['required_resources'], { 'disk': 2048, 'ram': 1024, 'cpu': 1.0 }) self.assertEqual( packing_plan['container_plans'][0]['instances'][0], { 'component_index': 1, 'component_name': 'word', 'instance_resources': { 'cpu': 1.0, 'disk': 2048, 'ram': 1024 }, 'task_id': 1 })