class SubscriberDBCloudClientTests(unittest.TestCase):
"""Tests for the SubscriberDBCloudClient"""
def setUp(self):
"""Initialize client tests"""
# Create sqlite3 database for testing
self._tmpfile = tempfile.TemporaryDirectory()
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
store = SqliteStore(
'{filename}{slash}'.format(
filename=self._tmpfile.name,
slash='/',
), )
ServiceRegistry.add_service('test', '0.0.0.0', 0) # noqa: S104
ServiceRegistry._PROXY_CONFIG = {
'local_port': 1234,
'cloud_address': '',
'proxy_cloud_connections': False,
}
ServiceRegistry._REGISTRY = {
"services": {
"s6a_service": {
"ip_address": "0.0.0.0", # noqa: S104
"port": 2345,
},
},
}
self.service = MagicMock()
self.service.loop = self.loop
# 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 = MockSubscriberDBServer()
self._servicer.add_to_server(self._rpc_server)
self._rpc_server.start()
# Create a rpc stub
self.channel = grpc.insecure_channel(
'0.0.0.0:{port}'.format(port=port, ), )
grpc_client_manager = GRPCClientManager(
service_name="subscriberdb",
service_stub=SubscriberDBCloudStub,
max_client_reuse=60,
)
self.subscriberdb_cloud_client = SubscriberDBCloudClient(
loop=self.service.loop,
store=store,
subscriber_page_size=2,
sync_interval=10,
grpc_client_manager=grpc_client_manager,
)
self.subscriberdb_cloud_client.start()
def tearDown(self):
"""Clean up test setup"""
self._tmpfile.cleanup()
self._rpc_server.stop(None)
self.subscriberdb_cloud_client.stop()
def get_all_subscribers(self):
return [
SubscriberData(sid=SubscriberID(id="IMSI111", ), ),
SubscriberData(sid=SubscriberID(id="IMSI222", ), ),
SubscriberData(sid=SubscriberID(id="IMSI333", ), ),
SubscriberData(sid=SubscriberID(id="IMSI444", ), ),
SubscriberData(sid=SubscriberID(id="IMSI555", ), ),
SubscriberData(sid=SubscriberID(id="IMSI666", ), ),
]
@unittest.mock.patch(
'magma.common.service_registry.ServiceRegistry.get_rpc_channel', )
def test_get_all_subscribers_success(self, get_grpc_mock):
"""
Test ListSubscribers RPC happy path
Args:
get_grpc_mock: mock for service registry
"""
async def test(): # noqa: WPS430
get_grpc_mock.return_value = self.channel
ret = (await self.subscriberdb_cloud_client._get_all_subscribers())
self.assertTrue(ret is not None)
self.assertEqual(self.get_all_subscribers(), ret.subscribers)
self.assertEqual("root_digest_apple",
ret.root_digest.md5_base64_digest)
self.assertEqual(1, len(ret.leaf_digests))
self.assertEqual(
ret.leaf_digests[0].digest.md5_base64_digest,
"leaf_digests_apple",
)
self.assertEqual(ret.leaf_digests[0].id, "IMSI11111")
# Cancel the client's loop so there are no other activities
self.subscriberdb_cloud_client._periodic_task.cancel()
self.loop.run_until_complete(test())
@unittest.mock.patch(
'magma.common.service_registry.ServiceRegistry.get_rpc_channel', )
def test_get_all_subscribers_failure(self, get_grpc_mock):
"""
Test ListSubscribers RPC failures
Args:
get_grpc_mock: mock for service registry
"""
async def test(): # noqa: WPS430
get_grpc_mock.return_value = self.channel
# update page size to special value to trigger gRPC error
self.subscriberdb_cloud_client._subscriber_page_size = 1
ret = (await self.subscriberdb_cloud_client._get_all_subscribers())
self.assertTrue(ret is None)
# Cancel the client's loop so there are no other activities
self.subscriberdb_cloud_client._periodic_task.cancel()
self.loop.run_until_complete(test())
@unittest.mock.patch('magma.subscriberdb.client.S6aServiceStub')
def test_detach_deleted_subscribers(self, s6a_service_mock_stub):
"""
Test if the subscriberdb cloud client detaches deleted subscribers
Args:
s6a_service_mock_stub: mocked s6a stub
"""
# Mock out DeleteSubscriber.future
mock = unittest.mock.Mock()
mock.DeleteSubscriber.future.side_effect = [unittest.mock.Mock()]
s6a_service_mock_stub.side_effect = [mock]
# Call with no samples
old_sub_ids = ["IMSI202", "IMSI101"]
new_sub_ids = ["IMSI101", "IMSI202"]
self.subscriberdb_cloud_client._detach_deleted_subscribers(
old_sub_ids,
new_sub_ids,
)
s6a_service_mock_stub.DeleteSubscriber.future.assert_not_called()
self.subscriberdb_cloud_client._loop.stop()
# Call with one subscriber id deleted
old_sub_ids = ["IMSI202", "IMSI101", "IMSI303"]
new_sub_ids = ["IMSI202"]
self.subscriberdb_cloud_client._detach_deleted_subscribers(
old_sub_ids,
new_sub_ids,
)
mock.DeleteSubscriber.future.assert_called_once_with(
DeleteSubscriberRequest(imsi_list=["101", "303"]), )
@unittest.mock.patch(
'magma.common.service_registry.ServiceRegistry.get_rpc_channel', )
def test_check_subscribers_in_sync(self, get_grpc_mock):
"""
Test CheckInSync RPC success
Args:
get_grpc_mock: mock for service registry
"""
async def test(): # noqa: WPS430
get_grpc_mock.return_value = self.channel
in_sync = (
await
self.subscriberdb_cloud_client._check_subscribers_in_sync())
self.assertEqual(False, in_sync)
self.subscriberdb_cloud_client._store.update_root_digest(
"root_digest_apple")
in_sync = (
await
self.subscriberdb_cloud_client._check_subscribers_in_sync())
self.assertEqual(True, in_sync)
# Cancel the client's loop so there are no other activities
self.subscriberdb_cloud_client._periodic_task.cancel()
self.loop.run_until_complete(test())
@unittest.mock.patch(
'magma.common.service_registry.ServiceRegistry.get_rpc_channel', )
def test_sync_subscribers(self, get_grpc_mock):
"""
Test Sync RPC success
Args:
get_grpc_mock: mock for service registry
"""
async def test(): # noqa: WPS430
get_grpc_mock.return_value = self.channel
# resync is True if the changeset is too big
resync = (await self.subscriberdb_cloud_client._sync_subscribers())
self.assertEqual(True, resync)
self.subscriberdb_cloud_client._store.update_leaf_digests([
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest="digest_apple"),
),
LeafDigest(
id='IMSI00000',
digest=Digest(md5_base64_digest="digest_zebra"),
),
])
self.subscriberdb_cloud_client._store.add_subscriber(
subscriber_data_by_id('IMSI00000'), )
self.subscriberdb_cloud_client._store.add_subscriber(
subscriber_data_by_id('IMSI11111'), )
# the client subscriber db and leaf digests db are updated
# when resync is False
expected_leaf_digests = [
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest="digest_apple"),
),
LeafDigest(
id='IMSI22222',
digest=Digest(md5_base64_digest="digest_banana"),
),
LeafDigest(
id='IMSI33333',
digest=Digest(md5_base64_digest="digest_cherry"),
),
]
resync = (await self.subscriberdb_cloud_client._sync_subscribers())
self.assertEqual(False, resync)
self.assertEqual(
"root_digest_apple",
self.subscriberdb_cloud_client._store.get_current_root_digest(
),
)
self.assertEqual(
['IMSI11111', 'IMSI22222', 'IMSI33333'],
self.subscriberdb_cloud_client._store.list_subscribers(),
)
self.assertEqual(
expected_leaf_digests,
self.subscriberdb_cloud_client._store.get_current_leaf_digests(
),
)
# Cancel the client's loop so there are no other activities
self.subscriberdb_cloud_client._periodic_task.cancel()
self.loop.run_until_complete(test())
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()