def test_MultiUpdate(self): """ Test consecutive updates """ assignments_dict = {} apn_rules_dict = {} # type: Dict[str, SubscriberPolicySet] basenames_dict = {} callback = RuleMappingsStreamerCallback( ReAuthHandler(assignments_dict, MockSessionProxyResponderStub3()), basenames_dict, assignments_dict, apn_rules_dict, ) # Construct a set of updates, keyed by subscriber ID updates = [ DataUpdate( key="s1", value=AssignedPolicies( assigned_policies=["p1", "p2"], ).SerializeToString(), ), DataUpdate( key="s2", value=AssignedPolicies( assigned_policies=["p2", "p3"], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) updates = [ DataUpdate( key="s1", value=AssignedPolicies( assigned_policies=["p4", "p5"], ).SerializeToString(), ), DataUpdate( key="s2", value=AssignedPolicies( assigned_policies=["p4", "p5"], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) s1_policies = assignments_dict["s1"].installed_policies expected = 2 self.assertEqual(len(s1_policies), expected, 'There should be 2 active policies for s1') self.assertTrue("p5" in s1_policies, 'Policy p5 should be active for subscriber s1') self.assertTrue("p4" in s1_policies, 'Policy p4 should be active for subscriber s1') s2_policies = assignments_dict["s2"].installed_policies expected = 2 self.assertEqual(len(s2_policies), expected, 'There should be 2 active policies for s2') self.assertTrue("p5" in s2_policies, 'Policy p5 should be active for subscriber s2') self.assertTrue("p4" in s2_policies, 'Policy p4 should be active for subscriber s2')
def test_SuccessfulUpdate(self): """ Test the happy path where updates come in for added rules, and sessiond accepts the RAR without issue. """ assignments_dict = {} basenames_dict = { 'bn1': ChargingRuleNameSet(RuleNames=['p5']), 'bn2': ChargingRuleNameSet(RuleNames=['p6']), } callback = RuleMappingsStreamerCallback( ReAuthHandler(assignments_dict, MockSessionProxyResponderStub1()), basenames_dict, assignments_dict, ) # Construct a set of updates, keyed by subscriber ID updates = [ DataUpdate( key="s1", value=AssignedPolicies( assigned_policies=["p1", "p2"], assigned_base_names=["bn1"], ).SerializeToString(), ), DataUpdate( key="s2", value=AssignedPolicies( assigned_policies=["p2", "p3"], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) # Since we used a stub which always succeeds when a RAR is made, # We should expect the assignments_dict to be updated s1_policies = assignments_dict["s1"].installed_policies expected = 3 self.assertEqual(len(s1_policies), expected, 'There should be 3 active ' 'policies for s1') self.assertTrue("p1" in s1_policies, 'Policy p1 should be active for ' 'subscriber s1') self.assertTrue("p5" in s1_policies, 'Policy p5 should be active for ' 'subscriber s1') s2_policies = assignments_dict["s2"].installed_policies expected = 2 self.assertEqual(len(s2_policies), expected, 'There should be 2 active ' 'policies for s2') self.assertTrue("p3" in s2_policies, 'Policy p3 should be active for ' 'subscriber s2')
def test_FailedOnePolicy(self): """ Test when sessiond answers to the RAR with a failure for installing p2. """ assignments_dict = {} apn_rules_dict = {} # type: Dict[str, SubscriberPolicySet] basenames_dict = {} callback = RuleMappingsStreamerCallback( ReAuthHandler(assignments_dict, MockSessionProxyResponderStub3()), basenames_dict, assignments_dict, apn_rules_dict, ) # Construct a set of updates, keyed by subscriber ID updates = [ DataUpdate( key="s1", value=AssignedPolicies( assigned_policies=["p1", "p2"], ).SerializeToString(), ), DataUpdate( key="s2", value=AssignedPolicies( assigned_policies=["p2", "p3"], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) s1_policies = assignments_dict["s1"].installed_policies expected = 1 self.assertEqual(len(s1_policies), expected, 'There should be 1 active ' 'policies for s1') self.assertTrue("p1" in s1_policies, 'Policy p1 should be active for ' 'subscriber s1') s2_policies = assignments_dict["s2"].installed_policies expected = 1 self.assertEqual(len(s2_policies), expected, 'There should be 1 active ' 'policies for s2') self.assertTrue("p3" in s2_policies, 'Policy p3 should be active for ' 'subscriber s2')
def test_FailedUpdate(self): """ Test when sessiond answers to the RAR with a failure for any re-auth. """ assignments_dict = {} apn_rules_dict = {} # type: Dict[str, SubscriberPolicySet] basenames_dict = {} callback = RuleMappingsStreamerCallback( ReAuthHandler(assignments_dict, MockSessionProxyResponderStub2()), basenames_dict, assignments_dict, apn_rules_dict, ) # Construct a set of updates, keyed by subscriber ID updates = [ DataUpdate( key="s1", value=AssignedPolicies( assigned_policies=["p1", "p2"], ).SerializeToString(), ), DataUpdate( key="s2", value=AssignedPolicies( assigned_policies=["p2", "p3"], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) # Since we used a stub which always succeeds when a RAR is made, # We should expect the assignments_dict to be updated self.assertFalse("s1" in assignments_dict, 'There should be no entry ' 'for subscriber s1 since update failed') self.assertFalse("s2" in assignments_dict, 'There should be no entry ' 'for subscriber s2 since update failed')
def test_Update(self): """ Test the happy path where updates come in for rules, and sessiond accepts the SessionRules without issue. """ # Expected call arguments to SetSessionRules allow_all_flow_list = [ FlowDescription( match=FlowMatch( direction=FlowMatch.Direction.Value("UPLINK"), ), action=FlowDescription.Action.Value("PERMIT"), ), FlowDescription( match=FlowMatch( direction=FlowMatch.Direction.Value("DOWNLINK"), ), action=FlowDescription.Action.Value("PERMIT"), ), ] # type: List[FlowDescription] no_tracking_type = PolicyRule.TrackingType.Value("NO_TRACKING") expected = SessionRules(rules_per_subscriber=[ RulesPerSubscriber( imsi='imsi_1', rule_set=[ RuleSet( apply_subscriber_wide=False, apn="apn1", static_rules=[ StaticRuleInstall(rule_id="p1"), ], dynamic_rules=[ DynamicRuleInstall(policy_rule=PolicyRule( id="allowlist_sid-imsi_1-apn1", priority=2, flow_list=allow_all_flow_list, tracking_type=no_tracking_type, )) ], ), ]), RulesPerSubscriber( imsi='imsi_2', rule_set=[ RuleSet( apply_subscriber_wide=False, apn="apn1", static_rules=[ StaticRuleInstall(rule_id="p5"), ], dynamic_rules=[ DynamicRuleInstall(policy_rule=PolicyRule( id="allowlist_sid-imsi_2-apn1", priority=2, flow_list=allow_all_flow_list, tracking_type=no_tracking_type, )) ], ), ]) ]) # Setup the test apn_rules_dict = {} basenames_dict = { 'bn1': ChargingRuleNameSet(RuleNames=['p5']), 'bn2': ChargingRuleNameSet(RuleNames=['p6']), } stub = MockLocalSessionManagerStub() stub_call_args = [] # type: List[SessionRules] side_effect = get_SetSessionRules_side_effect(stub_call_args) stub.SetSessionRules = Mock(side_effect=side_effect) callback = ApnRuleMappingsStreamerCallback( stub, basenames_dict, apn_rules_dict, ) # Construct a set of updates, keyed by subscriber ID updates = [ DataUpdate( key="imsi_1", value=SubscriberPolicySet(rules_per_apn=[ ApnPolicySet( apn="apn1", assigned_base_names=[], assigned_policies=["p1"], ), ], ).SerializeToString(), ), DataUpdate( key="imsi_2", value=SubscriberPolicySet(rules_per_apn=[ ApnPolicySet( apn="apn1", assigned_base_names=["bn1"], assigned_policies=[], ), ], ).SerializeToString(), ), ] callback.process_update("stream", updates, False) # Since we used a stub which always succeeds when a RAR is made, # We should expect the assignments_dict to be updated imsi_1_policies = apn_rules_dict["imsi_1"] self.assertEqual(len(imsi_1_policies.rules_per_apn), 1, 'There should be 1 active APNs for imsi_1') self.assertEqual(len(stub_call_args), 1, 'Stub should have been called once') called_with = stub_call_args[0].SerializeToString() self.assertEqual(called_with, expected.SerializeToString(), 'SetSessionRules call has incorrect arguments') # Stream down a second update, and now IMSI_1 gets access to a new APN updates_2 = [ DataUpdate( key="imsi_1", value=SubscriberPolicySet(rules_per_apn=[ ApnPolicySet( apn="apn2", assigned_base_names=["bn1"], assigned_policies=[], ), ], ).SerializeToString(), ), DataUpdate( key="imsi_2", value=SubscriberPolicySet( global_base_names=["bn2"], global_policies=[], rules_per_apn=[ ApnPolicySet( apn="apn1", assigned_base_names=[], assigned_policies=[], ), ], ).SerializeToString(), ), ] expected_2 = SessionRules(rules_per_subscriber=[ RulesPerSubscriber( imsi='imsi_1', rule_set=[ RuleSet( apply_subscriber_wide=False, apn="apn2", static_rules=[ StaticRuleInstall(rule_id="p5"), ], dynamic_rules=[ DynamicRuleInstall(policy_rule=PolicyRule( id="allowlist_sid-imsi_1-apn2", priority=2, flow_list=allow_all_flow_list, tracking_type=no_tracking_type, )) ], ), ]), RulesPerSubscriber( imsi='imsi_2', rule_set=[ RuleSet( apply_subscriber_wide=False, apn="apn1", static_rules=[ StaticRuleInstall(rule_id="p6"), ], dynamic_rules=[ DynamicRuleInstall(policy_rule=PolicyRule( id="allowlist_sid-imsi_2-apn1", priority=2, flow_list=allow_all_flow_list, tracking_type=no_tracking_type, )) ], ), ]), ]) callback.process_update("stream", updates_2, False) imsi_1_policies = apn_rules_dict["imsi_1"] self.assertEqual(len(imsi_1_policies.rules_per_apn), 1, 'There should be 1 active APNs for imsi_1') self.assertEqual(len(stub_call_args), 2, 'Stub should have been called twice') called_with = stub_call_args[1].SerializeToString() self.assertEqual(called_with, expected_2.SerializeToString(), 'SetSessionRules call has incorrect arguments')
def test_process_update(self, config_mock): """ Test an update with the following characteristics: - There is an unrecognized Any type - There is an unrecognized service - The magmad mconfig is updated and flips the new feature flag on - Another service has its config updated - Another service does not have its config updated """ # Set up fixtures old_mconfig, new_mconfig = self._get_old_and_new_mconfig_fixtures() new_mconfig_ser = MessageToJson(new_mconfig) # Add a config with a bad type to the new mconfig (MessageToJson errors # out if we include it in the proto message) new_mconfig_json_deser = json.loads(new_mconfig_ser) new_mconfig_json_deser['configs']['configsByKey']['oops'] = { '@type': 'type.googleapis.com/not.a.real.type', 'value': 'value', } new_mconfig_ser = json.dumps(new_mconfig_json_deser) updates = [ DataUpdate( value=new_mconfig_ser.encode(), key='mock', ), ] # Set up mock dependencies config_mock.return_value = ['metricsd'] @asyncio.coroutine def _mock_restart_services(): return 'mock' service_manager_mock = mock.Mock() service_manager_mock.restart_services = mock.MagicMock( wraps=_mock_restart_services, ) mconfig_manager_mock = mock.Mock() mconfig_manager_mock.load_mconfig.return_value = old_mconfig mconfig_manager_mock.update_stored_mconfig = mock.MagicMock() # use original impl of deserialize mconfig_manager_mock.deserialize_mconfig = mock.MagicMock( wraps=functools.partial( StreamedMconfigManager.deserialize_mconfig, mconfig_manager_mock, ), ) magma_service_mock = mock.Mock() magma_service_mock.reload_mconfig = mock.MagicMock() loop_mock = mock.MagicMock() # Run function, assert calls callback = StreamingMconfigCallback( ['metricsd'], service_manager_mock, magma_service_mock, mconfig_manager_mock, allow_unknown_fields=False, loop=loop_mock, ) callback.process_update(MCONFIG_VIEW_STREAM_NAME, updates, True) # magmad and metricsd configs changed but only metricsd should be # restarted (magmad won't restart for config change) service_manager_mock.restart_services.assert_called_once_with( ['metricsd'], ) # We should have written the whole serialized mconfig, bad keys and all mconfig_manager_mock.update_stored_mconfig.assert_called_once_with( new_mconfig_ser, ) # Should have reloaded magmad mconfigs to pick up change magma_service_mock.reload_mconfig.assert_called_once_with()
def test_update(self, config_mock): """ Test that mconfig updates are handled correctly """ # Set up fixture data # Update will simulate gateway moving from # test_mconfig -> updated_mconfig test_mconfig = GatewayConfigs() updated_mconfig = GatewayConfigs() some_any = Any() magmad = MagmaD(log_level=1) some_any.Pack(magmad) test_mconfig.configs_by_key['magmad'].CopyFrom(some_any) updated_mconfig.configs_by_key['magmad'].CopyFrom(some_any) metricsd = MetricsD(log_level=2) some_any.Pack(metricsd) test_mconfig.configs_by_key['metricsd'].CopyFrom(some_any) metricsd = MetricsD(log_level=3) some_any.Pack(metricsd) updated_mconfig.configs_by_key['metricsd'].CopyFrom(some_any) # Set up mock dependencies config_mock.return_value = { 'magma_services': ['magmad', 'metricsd'], } @asyncio.coroutine def _mock_restart_services(): return "blah" @asyncio.coroutine def _mock_update_dynamic_services(): return "mockResponse" service_manager_mock = MagicMock() magmad_service_mock = MagicMock() mconfig_manager_mock = MconfigManagerImpl() load_mock = patch.object( mconfig_manager_mock, 'load_mconfig', MagicMock(return_value=test_mconfig), ) update_mock = patch.object( mconfig_manager_mock, 'update_stored_mconfig', Mock(), ) restart_service_mock = patch.object( service_manager_mock, 'restart_services', MagicMock(wraps=_mock_restart_services), ) update_dynamic_services_mock = patch.object( service_manager_mock, 'update_dynamic_services', MagicMock(wraps=_mock_update_dynamic_services), ) processed_updates_mock = patch('magma.magmad.events.processed_updates') class ServiceMconfigMock: def __init__(self, service, mconfig_struct): pass dynamic_services = [] mock_mcfg = patch('magma.magmad.config_manager.load_service_mconfig', MagicMock(wraps=ServiceMconfigMock)) with load_mock as loader,\ update_mock as updater, \ restart_service_mock as restarter,\ update_dynamic_services_mock as dynamic_services,\ mock_mcfg as mock_c,\ processed_updates_mock as processed_updates: loop = asyncio.new_event_loop() config_manager = ConfigManager( ['magmad', 'metricsd'], service_manager_mock, magmad_service_mock, mconfig_manager_mock, allow_unknown_fields=False, loop=loop, ) # Process an empty set of updates updates = [] config_manager.process_update(CONFIG_STREAM_NAME, updates, False) # No services should be updated or restarted due to empty updates restarter.assert_not_called() updater.assert_not_called() # Verify that config update restarts all services update_str = MessageToJson(updated_mconfig) updates = [ DataUpdate(value=''.encode('utf-8'), key='some key'), DataUpdate( value=update_str.encode('utf-8'), key='last key', ), ] config_manager.process_update(CONFIG_STREAM_NAME, updates, False) # Only metricsd config was updated, hence should be restarted self.assertEqual(loader.call_count, 1) restarter.assert_called_once_with(['metricsd']) updater.assert_called_once_with(update_str) configs_by_service = { 'magmad': updated_mconfig.configs_by_key['magmad'], 'metricsd': updated_mconfig.configs_by_key['metricsd'], } processed_updates.assert_called_once_with(configs_by_service) restarter.reset_mock() updater.reset_mock() processed_updates.reset_mock() updated_mconfig.configs_by_key['shared_mconfig'].CopyFrom(some_any) update_str = MessageToJson(updated_mconfig) updates = [ DataUpdate( value=update_str.encode('utf-8'), key='last key', ), ] config_manager.process_update(CONFIG_STREAM_NAME, updates, False) # shared config update should restart all services restarter.assert_called_once_with(['magmad', 'metricsd']) updater.assert_called_once_with(update_str) dynamic_services.assert_called_once_with([]) processed_updates.assert_called_once_with(configs_by_service) self.assertEqual(mock_c.call_count, 1) restarter.reset_mock() updater.reset_mock() processed_updates.reset_mock() dynamic_services.reset_mock()