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 _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
class DirectorydRpcServiceTests(TestCase): @mock.patch("redis.Redis", MockRedis) def setUp(self): # Bind the rpc server to a free port thread_pool = futures.ThreadPoolExecutor(max_workers=10) self._rpc_server = grpc.server(thread_pool) port = self._rpc_server.add_insecure_port('0.0.0.0:0') # Add the servicer self._servicer = GatewayDirectoryServiceRpcServicer() 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 = GatewayDirectoryServiceStub(channel) @mock.patch("redis.Redis", MockRedis) def tearDown(self): self._rpc_server.stop(None) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) 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") @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) 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, {}) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) 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) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) 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) @mock.patch("redis.Redis", MockRedis) @mock.patch('snowflake.snowflake', get_mock_snowflake) 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()