class RpcTests(unittest.TestCase):
"""
Tests for the SubscriberDB rpc servicer and stub
"""
def setUp(self):
# Create an in-memory store
self._tmpfile = tempfile.TemporaryDirectory()
store = SqliteStore(self._tmpfile.name + '/')
# Bind the rpc server to a free port
self._rpc_server = grpc.server(
futures.ThreadPoolExecutor(max_workers=10), )
port = self._rpc_server.add_insecure_port('0.0.0.0:0')
# Add the servicer
self._servicer = SubscriberDBRpcServicer(
store=store, lte_processor=self._create_lte_processor_mock())
self._servicer.add_to_server(self._rpc_server)
self._rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
self._stub = SubscriberDBStub(channel)
def _create_lte_processor_mock(self):
lte_processor_mock = MagicMock()
sub_profile_mock = MagicMock()
lte_processor_mock.get_sub_profile.return_value = sub_profile_mock
sub_profile_mock.max_ul_bit_rate = 23
sub_profile_mock.max_dl_bit_rate = 42
return lte_processor_mock
def tearDown(self):
self._tmpfile.cleanup()
self._rpc_server.stop(0)
def test_get_invalid_subscriber(self):
"""
Test if the rpc call returns NOT_FOUND
"""
with self.assertRaises(grpc.RpcError) as err:
self._stub.GetSubscriberData(SIDUtils.to_pb('IMSI123'))
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
def test_add_delete_subscriber(self):
"""
Test if AddSubscriber and DeleteSubscriber rpc call works
"""
sid = SIDUtils.to_pb('IMSI1')
data = SubscriberData(sid=sid)
# Add subscriber
self._stub.AddSubscriber(data)
# Add subscriber again
with self.assertRaises(grpc.RpcError) as err:
self._stub.AddSubscriber(data)
self.assertEqual(err.exception.code(), grpc.StatusCode.ALREADY_EXISTS)
# See if we can get the data for the subscriber
self.assertEqual(self._stub.GetSubscriberData(sid).sid, data.sid)
self.assertEqual(len(self._stub.ListSubscribers(Void()).sids), 1)
self.assertEqual(self._stub.ListSubscribers(Void()).sids[0], sid)
# Delete the subscriber
self._stub.DeleteSubscriber(sid)
self.assertEqual(len(self._stub.ListSubscribers(Void()).sids), 0)
def test_update_subscriber(self):
"""
Test if UpdateSubscriber rpc call works
"""
sid = SIDUtils.to_pb('IMSI1')
data = SubscriberData(sid=sid)
# Add subscriber
self._stub.AddSubscriber(data)
sub = self._stub.GetSubscriberData(sid)
self.assertEqual(sub.lte.auth_key, b'')
self.assertEqual(sub.state.lte_auth_next_seq, 0)
# Update subscriber
update = SubscriberUpdate()
update.data.sid.CopyFrom(sid)
update.data.lte.auth_key = b'\xab\xcd'
update.data.state.lte_auth_next_seq = 1
update.mask.paths.append('lte.auth_key') # only auth_key
self._stub.UpdateSubscriber(update)
sub = self._stub.GetSubscriberData(sid)
self.assertEqual(sub.state.lte_auth_next_seq, 0) # no change
self.assertEqual(sub.lte.auth_key, b'\xab\xcd')
update.data.state.lte_auth_next_seq = 1
update.mask.paths.append('state.lte_auth_next_seq')
self._stub.UpdateSubscriber(update)
sub = self._stub.GetSubscriberData(sid)
self.assertEqual(sub.state.lte_auth_next_seq, 1)
# Delete the subscriber
self._stub.DeleteSubscriber(sid)
with self.assertRaises(grpc.RpcError) as err:
self._stub.UpdateSubscriber(update)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
def main():
"""Main routine for subscriberdb service.""" # noqa: D401
service = MagmaService('subscriberdb', mconfigs_pb2.SubscriberDB())
# Optionally pipe errors to Sentry
sentry_init(service_name=service.name)
# Initialize a store to keep all subscriber data.
store = SqliteStore(
service.config['db_path'],
loop=service.loop,
sid_digits=service.config['sid_last_n'],
)
# Initialize the processor
processor = Processor(
store,
get_default_sub_profile(service),
service.mconfig.sub_profiles,
service.mconfig.lte_auth_op,
service.mconfig.lte_auth_amf,
)
# Add all servicers to the server
subscriberdb_servicer = SubscriberDBRpcServicer(
store,
service.config.get('print_grpc_payload', False),
)
subscriberdb_servicer.add_to_server(service.rpc_server)
# Start a background thread to stream updates from the cloud
if service.config['enable_streaming']:
grpc_client_manager = GRPCClientManager(
service_name="subscriberdb",
service_stub=SubscriberDBCloudStub,
max_client_reuse=60,
)
sync_interval = _randomize_sync_interval(
service.config.get('subscriberdb_sync_interval'), )
subscriber_page_size = service.config.get('subscriber_page_size')
subscriberdb_cloud_client = SubscriberDBCloudClient(
service.loop,
store,
subscriber_page_size,
sync_interval,
grpc_client_manager,
)
subscriberdb_cloud_client.start()
else:
logging.info(
'enable_streaming set to False. Subscriber streaming '
'disabled!', )
# Wait until the datastore is populated by addition or resync before
# listening for clients.
async def serve(): # noqa: WPS430
if not store.list_subscribers():
# Waiting for subscribers to be added to store
await store.on_ready()
if service.config['s6a_over_grpc']:
logging.info('Running s6a over grpc')
s6a_proxy_servicer = S6aProxyRpcServicer(
processor,
service.config.get('print_grpc_payload', False),
)
s6a_proxy_servicer.add_to_server(service.rpc_server)
else:
logging.info('Running s6a over DIAMETER')
base_manager = base.BaseApplication(
service.config['mme_realm'],
service.config['mme_host_name'],
service.config['mme_host_address'],
)
s6a_manager = _get_s6a_manager(service, processor)
base_manager.register(s6a_manager)
# Setup the Diameter/s6a MME
s6a_server = service.loop.create_server(
lambda: S6aServer(
base_manager,
s6a_manager,
service.config['mme_realm'],
service.config['mme_host_name'],
loop=service.loop,
),
service.config['host_address'],
service.config['mme_port'],
)
asyncio.ensure_future(s6a_server, loop=service.loop)
asyncio.ensure_future(serve(), loop=service.loop)
# Run the service loop
service.run()
# Cleanup the service
service.close()