def setUp(self): # Create an in-memory store store = SqliteStore('file::memory:?cache=shared') # 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) 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 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) 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)
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()