def update_record(imsi: str, ip_addr: str) -> None:
"""
Make RPC call to 'UpdateRecord' method of local directoryD service
"""
try:
chan = ServiceRegistry.get_rpc_channel(DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL)
except ValueError:
logging.error('Cant get RPC channel to %s', DIRECTORYD_SERVICE_NAME)
return
client = GatewayDirectoryServiceStub(chan)
if not imsi.startswith("IMSI"):
imsi = "IMSI" + imsi
try:
# Location will be filled in by directory service
req = UpdateRecordRequest(id=imsi, location="hwid")
req.fields[IPV4_ADDR_KEY] = ip_addr
client.UpdateRecord(req, DEFAULT_GRPC_TIMEOUT)
except grpc.RpcError as err:
logging.error(
"UpdateRecordRequest error for id: %s, ipv4_addr: %s! [%s] %s",
imsi,
ip_addr,
err.code(),
err.details())
def test_redis_unavailable(self):
self._servicer._redis_dict = MockUnavailableRedis("localhost", 6380)
req = UpdateRecordRequest()
req.id = "IMSI557"
req.fields["mac_addr"] = "aa:bb:aa:bb:aa:bb"
with self.assertRaises(grpc.RpcError) as err:
self._stub.UpdateRecord(req)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNAVAILABLE)
get_req = GetDirectoryFieldRequest()
get_req.id = "IMSI557"
get_req.field_key = "mac_addr"
with self.assertRaises(grpc.RpcError) as err:
self._stub.GetDirectoryField(get_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNAVAILABLE)
del_req = DeleteRecordRequest()
del_req.id = "IMSI557"
with self.assertRaises(grpc.RpcError) as err:
self._stub.DeleteRecord(del_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNAVAILABLE)
void_req = Void()
with self.assertRaises(grpc.RpcError) as err:
self._stub.GetAllDirectoryRecords(void_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.UNAVAILABLE)
def test_update_record_bad_location(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI556"
req.location = "bbb-ccc"
self._stub.UpdateRecord(req)
actual_record = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record.location_history, ['aaa-bbb'])
self.assertEqual(actual_record.identifiers, {})
def update_record_handler(client, args):
update_record_request = UpdateRecordRequest()
update_record_request.id = args.id
if args.field_key is not None and args.field_value is not None:
update_record_request.fields[args.field_key] = args.field_value
try:
client.UpdateRecord(update_record_request)
print("Successfully updated record for ID: %s" % args.id)
except grpc.RpcError as e:
print("gRPC failed with %s: %s" % (e.code(), e.details()))
def test_get_all(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
req.fields["mac_addr"] = "aa:bb:aa:bb:aa:bb"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
req2 = UpdateRecordRequest()
req2.id = "IMSI556"
req2.fields["ipv4_addr"] = "192.168.127.11"
self._stub.UpdateRecord(req2)
self.assertTrue(req2.id in self._servicer._redis_dict)
void_req = Void()
ret = self._stub.GetAllDirectoryRecords(void_req)
self.assertEqual(2, len(ret.records))
for record in ret.records:
if record.id == "IMSI556":
self.assertEqual(record.fields["ipv4_addr"], "192.168.127.11")
elif record.id == "IMSI557":
self.assertEqual(record.fields["mac_addr"],
"aa:bb:aa:bb:aa:bb")
else:
raise AssertionError()
def test_delete_record(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
del_req = DeleteRecordRequest()
del_req.id = "IMSI557"
self._stub.DeleteRecord(del_req)
self.assertFalse(req.id in self._servicer._redis_dict)
with self.assertRaises(grpc.RpcError) as err:
self._stub.DeleteRecord(del_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
def test_get_field(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
req.fields["mac_addr"] = "aa:bb:aa:bb:aa:bb"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
get_req = GetDirectoryFieldRequest()
get_req.id = "IMSI557"
get_req.field_key = "mac_addr"
ret = self._stub.GetDirectoryField(get_req)
self.assertEqual("aa:bb:aa:bb:aa:bb", ret.value)
with self.assertRaises(grpc.RpcError) as err:
get_req.field_key = "ipv4_addr"
self._stub.GetDirectoryField(get_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
def test_update_record(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI555"
self._stub.UpdateRecord(req)
actual_record = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record.location_history, ['aaa-bbb'])
self.assertEqual(actual_record.identifiers, {})
req.fields["mac_addr"] = "aa:aa:bb:bb:cc:cc"
req.fields["ipv4_addr"] = "192.168.172.12"
self._stub.UpdateRecord(req)
actual_record2 = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record2.location_history, ["aaa-bbb"])
self.assertEqual(actual_record2.identifiers['mac_addr'],
"aa:aa:bb:bb:cc:cc")
self.assertEqual(actual_record2.identifiers['ipv4_addr'],
"192.168.172.12")
def _build_update_records_data(num_requests: int, input_file: str):
update_record_reqs = []
for i in range(num_requests):
id = str(i).zfill(15)
location = str(i).zfill(15)
request = UpdateRecordRequest(
id=id,
location=location,
)
request_dict = json_format.MessageToDict(request)
update_record_reqs.append(request_dict)
with open(input_file, 'w') as file:
json.dump(update_record_reqs, file, separators=(',', ':'))
def _load_subs(num_subs: int) -> List[DirectoryRecord]:
"""Load directory records"""
client = GatewayDirectoryServiceStub(
ServiceRegistry.get_rpc_channel(
DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL,
), )
sids = []
for i in range(num_subs):
mac_addr = (str(i) * 2 + ":") * 5 + (str(i) * 2)
ipv4_addr = str(i) * 3 + "." + str(i) * 3 + "." + str(
i) * 3 + "." + str(i) * 3
fields = {"mac-addr": mac_addr, "ipv4_addr": ipv4_addr}
sid = UpdateRecordRequest(
fields=fields,
id=str(i).zfill(15),
location=str(i).zfill(15),
)
client.UpdateRecord(sid)
sids.append(sid)
return sids