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 main():
service = MagmaService('policydb', mconfigs_pb2.PolicyDB())
apn_rules_dict = ApnRuleAssignmentsDict()
assignments_dict = RuleAssignmentsDict()
basenames_dict = BaseNameDict()
rating_groups_dict = RatingGroupsDict()
sessiond_chan = ServiceRegistry.get_rpc_channel('sessiond',
ServiceRegistry.LOCAL)
session_mgr_stub = LocalSessionManagerStub(sessiond_chan)
sessiond_stub = SessionProxyResponderStub(sessiond_chan)
reauth_handler = ReAuthHandler(assignments_dict, sessiond_stub)
# Add all servicers to the server
session_servicer = SessionRpcServicer(service.mconfig,
rating_groups_dict,
basenames_dict,
apn_rules_dict)
session_servicer.add_to_server(service.rpc_server)
orc8r_chan = ServiceRegistry.get_rpc_channel('policydb',
ServiceRegistry.CLOUD)
policy_stub = PolicyAssignmentControllerStub(orc8r_chan)
policy_servicer = PolicyRpcServicer(reauth_handler, basenames_dict,
policy_stub)
policy_servicer.add_to_server(service.rpc_server)
# Start a background thread to stream updates from the cloud
if service.config['enable_streaming']:
stream = StreamerClient(
{
'policydb': PolicyDBStreamerCallback(),
'apn_rule_mappings': ApnRuleMappingsStreamerCallback(
session_mgr_stub,
basenames_dict,
apn_rules_dict,
),
'rule_mappings': RuleMappingsStreamerCallback(
reauth_handler,
basenames_dict,
assignments_dict,
apn_rules_dict,
),
'rating_groups': RatingGroupsStreamerCallback(
rating_groups_dict),
},
service.loop,
)
stream.start()
else:
logging.info('enable_streaming set to False. Streamer disabled!')
# Run the service loop
service.run()
# Cleanup the service
service.close()
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_EnableStaticRules(self):
"""
Check the happy path where everything succeeds.
"""
rules_by_sid = {}
rules_by_basename = {
"bn1": ChargingRuleNameSet(RuleNames=["p4", "p5"], ),
}
reauth_handler = ReAuthHandler(
rules_by_sid,
MockSessionProxyResponderStub(),
)
servicer = PolicyRpcServicer(
reauth_handler,
rules_by_basename,
MockPolicyAssignmentControllerStub(),
)
# Bind the rpc server to a free port
thread_pool = futures.ThreadPoolExecutor(max_workers=10)
rpc_server = grpc.server(thread_pool)
port = rpc_server.add_insecure_port('0.0.0.0:0')
# Create a mock "mconfig" for the servicer to use
mconfig = unittest.mock.Mock()
mconfig.ip_block = None
# Add the servicer
servicer.add_to_server(rpc_server)
rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
stub = PolicyDBStub(channel)
rules_by_basename["bn1"] = ChargingRuleNameSet(RuleNames=["p4",
"p5"], )
req = EnableStaticRuleRequest(
imsi="s1",
rule_ids=["p1", "p2", "p3"],
base_names=["bn1"],
)
stub.EnableStaticRules(req)
self.assertEqual(
len(rules_by_sid["s1"].installed_policies),
5,
'After a successful update, Redis should be tracking '
'5 active rules.',
)
def test_FailOrc8r(self):
""" Check that nothing is updated if orc8r is unreachable """
rules_by_sid = {}
rules_by_basename = {
"bn1": ChargingRuleNameSet(RuleNames=["p4", "p5"], ),
}
reauth_handler = ReAuthHandler(
rules_by_sid,
MockSessionProxyResponderStub(),
)
servicer = PolicyRpcServicer(
reauth_handler,
rules_by_basename,
MockPolicyAssignmentControllerStub2(),
)
# Bind the rpc server to a free port
thread_pool = futures.ThreadPoolExecutor(max_workers=10)
rpc_server = grpc.server(thread_pool)
port = rpc_server.add_insecure_port('0.0.0.0:0')
# Create a mock "mconfig" for the servicer to use
mconfig = unittest.mock.Mock()
mconfig.ip_block = None
# Add the servicer
servicer.add_to_server(rpc_server)
rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
stub = PolicyDBStub(channel)
req = EnableStaticRuleRequest(
imsi="s1",
rule_ids=["p1", "p2", "p3"],
base_names=["bn1"],
)
with self.assertRaises(grpc.RpcError):
stub.EnableStaticRules(req)
self.assertFalse(
"s1" in rules_by_sid,
"There should be no installed policies for s1",
)
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_SuccessfulUpdate(self):
"""
Test the happy path where updates come in for added rules, and sessiond
accepts the RAR without issue.
"""
install_dict = {}
handler = ReAuthHandler(
install_dict,
MockSessionProxyResponderStub(),
)
rar = PolicyReAuthRequest(
imsi='s1',
rules_to_install=[
StaticRuleInstall(rule_id=rule_id) for rule_id in ['p1', 'p2']
],
)
handler.handle_policy_re_auth(rar)
s1_policies = install_dict['s1'].installed_policies
expected = 2
self.assertEqual(
len(s1_policies),
expected,
'There should be 2 installed policies for s1',
)
self.assertTrue(
'p1' in s1_policies,
'Policy p1 should be marked installed for p1',
)
rar = PolicyReAuthRequest(
imsi='s1',
rules_to_install=[StaticRuleInstall(rule_id='p3')],
)
handler.handle_policy_re_auth(rar)
s1_policies = install_dict['s1'].installed_policies
expected = 3
self.assertEqual(
len(s1_policies),
expected,
'There should be 3 installed policies for s1',
)
self.assertTrue(
'p3' in s1_policies,
'Policy p1 should be marked installed for p1',
)