def test_thrifttype_sensitivity(self):
thrift_obj = lsdb_types.PrefixDatabase()
thrift_obj.thisNodeName = "some node"
raw_msg = serializer.serialize_thrift_object(thrift_obj)
recovered_obj = serializer.deserialize_thrift_object(
raw_msg, lsdb_types.PrefixEntry)
self.assertTrue(thrift_obj != recovered_obj)
def _run(self, client: OpenrCtrl.Client, node_name: str,
prefix_str: str) -> None:
key = Consts.STATIC_PREFIX_ALLOC_PARAM_KEY
# Retrieve previous allocation
resp = client.getKvStoreKeyVals([key])
allocs = None
if key in resp.keyVals:
allocs = serializer.deserialize_thrift_object(
resp.keyVals.get(key).value, alloc_types.StaticAllocation)
else:
allocs = alloc_types.StaticAllocation(nodePrefixes={})
# Return if there is no change
prefix = ipnetwork.ip_str_to_prefix(prefix_str)
if allocs.nodePrefixes.get(node_name) == prefix:
print("No changes needed. {}'s prefix is already set to {}".format(
node_name, prefix_str))
return
# Update value in KvStore
allocs.nodePrefixes[node_name] = prefix
value = serializer.serialize_thrift_object(allocs)
super(AllocationsSetCmd, self)._run(client, key, value, "breeze", None,
Consts.CONST_TTL_INF)
def run_set(self, node_name, prefix_str):
key = Consts.STATIC_PREFIX_ALLOC_PARAM_KEY
prefix = ipnetwork.ip_str_to_prefix(prefix_str)
# Retrieve previous allocation
resp = self.client.get_keys([key])
allocs = None
if key in resp.keyVals:
allocs = serializer.deserialize_thrift_object(
resp.keyVals.get(key).value, alloc_types.StaticAllocation
)
else:
allocs = alloc_types.StaticAllocation(nodePrefixes={})
# Return if there is no change
if allocs.nodePrefixes.get(node_name) == prefix:
print(
"No changes needed. {}'s prefix is already set to {}".format(
node_name, prefix_str
)
)
return
# Update value in KvStore
allocs.nodePrefixes[node_name] = prefix
value = serializer.serialize_thrift_object(allocs)
self.run(key, value, "breeze", None, Consts.CONST_TTL_INF)
def _run(self, client: OpenrCtrl.Client, node_name: str) -> None:
area = self.get_area_id()
key = Consts.STATIC_PREFIX_ALLOC_PARAM_KEY
# Retrieve previous allocation
resp = None
if area is None:
resp = client.getKvStoreKeyVals([key])
else:
resp = client.getKvStoreKeyValsArea([key], area)
allocs = None
if key in resp.keyVals:
allocs = serializer.deserialize_thrift_object(
resp.keyVals.get(key).value, openr_types.StaticAllocation)
else:
allocs = openr_types.StaticAllocation(nodePrefixes={node_name: ""})
# Return if there need no change
if node_name not in allocs.nodePrefixes:
print(
"No changes needed. {}'s prefix is not set".format(node_name))
return
# Update value in KvStore
del allocs.nodePrefixes[node_name]
value = serializer.serialize_thrift_object(allocs)
super(AllocationsUnsetCmd, self)._run(client, key, value, "breeze",
None, Consts.CONST_TTL_INF)
def send_and_recv_thrift_obj(self, thrift_obj_to_send,
thrift_type_to_recv):
req = serializer.serialize_thrift_object(thrift_obj_to_send,
self.cli_opts.proto_factory)
resp = None
if self.thrift_client:
try:
resp = self.thrift_client.command(self.module_type, req)
except Exception as e:
print(
"Tried to connect via thrift but could not. Exception: "
"{}".format(e),
file=sys.stderr,
)
self.cleanup_thrift()
raise e
else:
try:
self.zmq_client.send(req)
resp = self.zmq_client.recv()
# TODO: remove after Link monitor socket is changed to ROUTER everywhere
except Exception as e:
if OpenrModuleType.LINK_MONITOR == self.module_type:
dealer_client = self.get_zmq_client(zmq.DEALER)
dealer_client.send(req)
resp = dealer_client.recv()
else:
raise e
if thrift_type_to_recv is str:
return str(resp)
return serializer.deserialize_thrift_object(
resp, thrift_type_to_recv, self.cli_opts.proto_factory)
def test_exception_handling(self):
thrift_obj = lsdb_types.PrefixDatabase()
thrift_obj.thisNodeName = "some node"
raw_msg = serializer.serialize_thrift_object(thrift_obj)
# should raise exception due to inconsistency of protocol factor
with self.assertRaises(Exception):
serializer.deserialize_thrift_object(raw_msg,
lsdb_types.PrefixDatabase,
TJSONProtocolFactory)
def send_thrift_obj(self, thrift_obj):
''' Serialize thrift object and send to socket '''
raw_msg = serializer.serialize_thrift_object(thrift_obj,
self.proto_factory)
try:
self._socket.send(raw_msg)
except Exception as ex:
raise Exception("Failed sending message: {}".format(ex))
def test_reverse_equality(self):
for _ in range(100):
thrift_obj = lsdb_types.PrefixDatabase()
random_string = "".join(
random.choice(string.digits) for _ in range(10))
thrift_obj.thisNodeName = random_string
raw_msg = serializer.serialize_thrift_object(thrift_obj)
recovered_obj = serializer.deserialize_thrift_object(
raw_msg, lsdb_types.PrefixDatabase)
self.assertEqual(thrift_obj, recovered_obj)
for _ in range(100):
thrift_obj = lsdb_types.PrefixDatabase()
random_string = "".join(
random.choice(string.digits) for _ in range(10))
thrift_obj.thisNodeName = random_string
raw_msg = serializer.serialize_thrift_object(
thrift_obj, TJSONProtocolFactory)
recovered_obj = serializer.deserialize_thrift_object(
raw_msg, lsdb_types.PrefixDatabase, TJSONProtocolFactory)
self.assertEqual(thrift_obj, recovered_obj)
def send_and_recv_thrift_obj(self, thrift_obj_to_send,
thrift_type_to_recv):
req = serializer.serialize_thrift_object(thrift_obj_to_send,
self.cli_opts.proto_factory)
resp = None
if self.thrift_has_module:
# we should be able to connect to this module via the thriftCtrl
# server. if we think SSL is enabled on the server, lets try that
# first
if self.cli_opts.ssl:
try:
resp = self.get_thrift_client(True).command(
self.module_type, req)
except Exception as e:
print(
"Tried to connect via secure thrift but could not. "
"Exception: {}.".format(e),
file=sys.stderr,
)
if resp is None:
try:
resp = self.get_thrift_client(False).command(
self.module_type, req)
except Exception as e:
print(
"Tried to connect via plaintext thrift but could not. "
"Exception: {}.".format(e),
file=sys.stderr,
)
if resp is None:
try:
zmq_client = self.get_zmq_client()
zmq_client.send(req)
resp = zmq_client.recv()
# TODO: remove after Link monitor socket is changed to ROUTER everywhere
except Exception as e:
if OpenrModuleType.LINK_MONITOR == self.module_type:
zmq_client = self.get_zmq_client(zmq.DEALER)
zmq_client.send(req)
resp = zmq_client.recv()
else:
raise e
if thrift_type_to_recv is str:
return str(resp)
return serializer.deserialize_thrift_object(
resp, thrift_type_to_recv, self.cli_opts.proto_factory)
def _cs_client():
cs_client_inst = config_store_client.ConfigStoreClient(
zmq.Context(), "ipc:///tmp/config_store_cmd")
self.assertEqual(cs_client_inst.load('key1'), store_db['key1'])
with self.assertRaises(Exception):
cs_client_inst.load('key3')
self.assertTrue(cs_client_inst.erase('key1'))
with self.assertRaises(Exception):
cs_client_inst.load('key1')
value = serialize_thrift_object(
lm_types.DumpLinksReply(thisNodeName='node5'))
self.assertTrue(cs_client_inst.store('key5', value))
self.assertEqual(cs_client_inst.load('key5'), value)
def _cs_client():
cs_client_inst = config_store_client.ConfigStoreClient(
ctx, "inproc://openr_config_store_cmd")
self.assertEqual(cs_client_inst.load("key1"), store_db["key1"])
with self.assertRaises(Exception):
cs_client_inst.load("key3")
self.assertTrue(cs_client_inst.erase("key1"))
with self.assertRaises(Exception):
cs_client_inst.load("key1")
value = serialize_thrift_object(
lm_types.DumpLinksReply(thisNodeName="node5"))
self.assertTrue(cs_client_inst.store("key5", value))
self.assertEqual(cs_client_inst.load("key5"), value)
def run_unset(self, node_name):
key = Consts.STATIC_PREFIX_ALLOC_PARAM_KEY
# Retrieve previous allocation
resp = self.client.get_keys([key])
allocs = None
if key in resp.keyVals:
allocs = serializer.deserialize_thrift_object(
resp.keyVals.get(key).value, alloc_types.StaticAllocation
)
else:
allocs = alloc_types.StaticAllocation(nodePrefixes={node_name: ""})
# Return if there need no change
if node_name not in allocs.nodePrefixes:
print("No changes needed. {}'s prefix is not set".format(node_name))
# Update value in KvStore
del allocs.nodePrefixes[node_name]
value = serializer.serialize_thrift_object(allocs)
self.run(key, value, "breeze", None, Consts.CONST_TTL_INF)
from __future__ import unicode_literals
from __future__ import division
from openr.utils import socket
from openr.utils.serializer import serialize_thrift_object
from openr.clients import config_store_client
from openr.PersistentStore import ttypes as ps_types
from openr.LinkMonitor import ttypes as lm_types
import zmq
import unittest
from multiprocessing import Process
store_db = {
'key1':
serialize_thrift_object(lm_types.DumpLinksReply(thisNodeName='node1')),
'key2':
serialize_thrift_object(lm_types.DumpLinksReply(thisNodeName='node2'))
}
class ConfigStore():
def __init__(self, zmq_ctx, url):
self._cs_server_socket = socket.Socket(zmq_ctx, zmq.REP)
self._cs_server_socket.bind(url)
self._store_db = store_db
def process_request(self):
req = self._cs_server_socket.recv_thrift_obj(ps_types.StoreRequest)
if req.requestType == ps_types.StoreRequestType.LOAD:
