def test_leaf_digests(self):
"""
Test if leaf digests gets & updates work as expected
"""
self.assertEqual(self._store.get_current_leaf_digests(), [])
digests1 = [
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest='digest_apple'),
),
LeafDigest(
id='IMSI22222',
digest=Digest(md5_base64_digest='digest_banana'),
),
]
self._store.update_leaf_digests(digests1)
self.assertEqual(self._store.get_current_leaf_digests(), digests1)
digests2 = [
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest='digest_apple'),
),
LeafDigest(
id='IMSI33333',
digest=Digest(md5_base64_digest='digest_cherry'),
),
LeafDigest(
id='IMSI44444',
digest=Digest(md5_base64_digest='digest_dragonfruit'),
),
]
self._store.update_leaf_digests(digests2)
self.assertEqual(self._store.get_current_leaf_digests(), digests2)
def Sync(self, request: SyncRequest, context) -> SyncResponse:
"""
Mock to trigger Sync-related test cases
Args:
request: SyncRequest
context: request context
Returns:
SyncResponse
"""
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"),
),
]
client_leaf_digest_ids = [
digest.id for digest in request.leaf_digests
]
to_renew = []
deleted = []
if 'IMSI11111' not in client_leaf_digest_ids:
to_renew.append(subscriber_data_by_id('IMSI11111'))
if 'IMSI22222' not in client_leaf_digest_ids:
to_renew.append(subscriber_data_by_id('IMSI22222'))
if 'IMSI33333' not in client_leaf_digest_ids:
to_renew.append(subscriber_data_by_id('IMSI33333'))
if 'IMSI00000' in client_leaf_digest_ids:
deleted.append('IMSI00000')
resync = len(to_renew) >= 3
marshaled_to_renew = []
for data in to_renew:
anyVal = Any()
anyVal.Pack(data)
marshaled_to_renew.append(anyVal)
return SyncResponse(
resync=resync,
digests=DigestTree(
root_digest=Digest(md5_base64_digest="root_digest_apple"),
leaf_digests=leaf_digests,
),
changeset=Changeset(
to_renew=marshaled_to_renew,
deleted=deleted,
),
)
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(),
)
def ListSubscribers(self, request: ListSubscribersRequest,
context) -> ListSubscribersResponse: # noqa: N802
"""
List subscribers is a mock to trigger various test cases
Args:
request: ListSubscribersRequest
context: request context
Raises:
RpcError: If page size is 1
Returns:
ListSubscribersResponse
"""
# Add in logic to allow error handling testing
root_digest = Digest(md5_base64_digest="")
leaf_digests = []
if request.page_size == 1:
raise grpc.RpcError("Test Exception")
if request.page_token == "":
next_page_token = "aaa" # noqa: S105
subscribers = [
SubscriberData(sid=SubscriberID(id="IMSI111")),
SubscriberData(sid=SubscriberID(id="IMSI222")),
]
root_digest = Digest(md5_base64_digest="root_digest_apple")
leaf_digests = [
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest="leaf_digests_apple"),
),
]
elif request.page_token == "aaa":
next_page_token = "bbb" # noqa: S105
subscribers = [
SubscriberData(sid=SubscriberID(id="IMSI333")),
SubscriberData(sid=SubscriberID(id="IMSI444")),
]
else:
next_page_token = "" # noqa: S105
subscribers = [
SubscriberData(sid=SubscriberID(id="IMSI555")),
SubscriberData(sid=SubscriberID(id="IMSI666")),
]
return ListSubscribersResponse(
subscribers=subscribers,
next_page_token=next_page_token,
digests=DigestTree(
root_digest=root_digest,
leaf_digests=leaf_digests,
),
)
async def _check_subscribers_in_sync(self) -> bool:
"""
Check if the local subscriber data is up-to-date with the cloud by
comparing root digests
Returns:
boolean value for whether the local data is in sync
"""
subscriberdb_cloud_client = self._grpc_client_manager.get_client()
req = CheckInSyncRequest(root_digest=Digest(
md5_base64_digest=self._store.get_current_root_digest(), ), )
try:
res = await grpc_async_wrapper(
subscriberdb_cloud_client.CheckInSync.future(
req,
self.SUBSCRIBERDB_REQUEST_TIMEOUT,
),
self._loop,
)
except grpc.RpcError as err:
logging.error(
"Check subscribers in sync request error! [%s] %s",
err.code(),
err.details(),
)
return False
return res.in_sync
def test_subscriber_deletion_digests(self):
"""
Test if subscriber deletion also unconditionally removes digest info.
Regression test for #9029.
"""
(sid1, _) = self._add_subscriber('IMSI11111')
(sid2, _) = self._add_subscriber('IMSI22222')
self.assertEqual(self._store.list_subscribers(), [sid1, sid2])
root_digest = "apple"
leaf_digest = LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest="digest_apple"),
)
self._store.update_root_digest(root_digest)
self._store.update_leaf_digests([leaf_digest])
self.assertNotEqual(self._store.get_current_root_digest(), "")
self.assertNotEqual(self._store.get_current_leaf_digests(), [])
self._store.delete_subscriber(sid2)
self.assertEqual(self._store.list_subscribers(), [sid1])
# Deleting a subscriber also deletes all digest info
self.assertEqual(self._store.get_current_root_digest(), "")
self.assertEqual(self._store.get_current_leaf_digests(), [])
def test_leaf_digests_update(self):
"""
Test if leaf digests update triggers ready
"""
self.assertEqual(self._store._on_digests_ready.event.is_set(), False)
self._store.update_leaf_digests([
LeafDigest(
id='IMSI11111',
digest=Digest(md5_base64_digest='digest_cherry'),
),
])
async def defer():
await self._store.on_digests_ready()
self.loop.run_until_complete(defer())
self.assertEqual(self._store._on_digests_ready.event.is_set(), True)
def get_current_leaf_digests(self) -> List[LeafDigest]:
digests = []
conn = sqlite3.connect(self._leaf_digests_db_location, uri=True)
try:
with conn:
res = conn.execute(
"SELECT sid, digest FROM subscriber_leaf_digests ", )
for row in res:
digest = LeafDigest(
id=row[0],
digest=Digest(md5_base64_digest=row[1]),
)
digests.append(digest)
finally:
conn.close()
return digests