class ShardedProtoTableStorageTest(unittest.TestCase):
EXAMPLE_PROTO = ProtoUtil.text_to_message(message_type=NodeSnapshot,
text_str="""
node_name: "test"
children_names: ["child_1", "child_2"]
parents_names: ["parent_1", "parent_2"]
""")
TEST_DATA_DIR_1 = '/galaxy/ab-d/pslx/test_data/sharded_proto_table_1'
TEST_DATA_DIR_2 = '/galaxy/ab-d/pslx/test_data/sharded_proto_table_2'
TEST_DATA_DIR_3 = '/galaxy/ab-d/pslx/test_data/sharded_proto_table_3'
def test_initialize_from_dir(self):
shared_proto_table_storage = ShardedProtoTableStorage(size_per_shard=3)
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
self.assertFalse(shared_proto_table_storage.is_empty())
def test_write_1(self):
shared_proto_table_storage = ShardedProtoTableStorage(size_per_shard=3)
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
shared_proto_table_storage.write(data={
'test_7': self.EXAMPLE_PROTO,
})
self.assertFalse(shared_proto_table_storage.is_empty())
self.assertEqual(shared_proto_table_storage.get_num_shards(), 3)
gclient.rm_dir(self.TEST_DATA_DIR_1)
gclient_ext.cp_folder(self.TEST_DATA_DIR_2, self.TEST_DATA_DIR_1)
def test_write_2(self):
shared_proto_table_storage = ShardedProtoTableStorage(size_per_shard=3)
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
shared_proto_table_storage.write(
data={
'test_7': self.EXAMPLE_PROTO,
'test_8': self.EXAMPLE_PROTO,
'test_9': self.EXAMPLE_PROTO,
})
self.assertEqual(shared_proto_table_storage.get_num_entries(), 10)
self.assertEqual(shared_proto_table_storage.get_num_shards(), 4)
gclient.rm_dir(self.TEST_DATA_DIR_1)
gclient_ext.cp_folder(self.TEST_DATA_DIR_2, self.TEST_DATA_DIR_1)
def test_write_3(self):
shared_proto_table_storage = ShardedProtoTableStorage(size_per_shard=3)
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
shared_proto_table_storage.write(
data={
'test_7': self.EXAMPLE_PROTO,
'test_8': self.EXAMPLE_PROTO,
'test_0': self.EXAMPLE_PROTO,
})
self.assertEqual(shared_proto_table_storage.get_num_entries(), 9)
self.assertEqual(shared_proto_table_storage.get_num_shards(), 3)
gclient.rm_dir(self.TEST_DATA_DIR_1)
gclient_ext.cp_folder(self.TEST_DATA_DIR_2, self.TEST_DATA_DIR_1)
def test_read_1(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
data = shared_proto_table_storage.read(params={'key': 'test_0'})
self.assertTrue(data is not None)
def test_read_2(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
data = shared_proto_table_storage.read_multiple(
params={'keys': ['test_0', 'test_100']})
self.assertTrue('test_100' not in data)
def test_read_3(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
data = shared_proto_table_storage.read(params={'key': 'test_100'})
self.assertIsNone(data)
def test_read_4(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
data = shared_proto_table_storage.read_multiple(
params={'keys': ['test_0', 'test_1']})
self.assertTrue('test_0' in data and 'test_1' in data)
def test_read_5(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
shared_proto_table_storage.write(
data={
'test_7': self.EXAMPLE_PROTO,
'test_8': self.EXAMPLE_PROTO,
'test_9': self.EXAMPLE_PROTO,
})
data = shared_proto_table_storage.read_multiple(
params={'keys': ['test_9']})
self.assertTrue('test_9' in data)
gclient.rm_dir(self.TEST_DATA_DIR_1)
gclient_ext.cp_folder(self.TEST_DATA_DIR_2, self.TEST_DATA_DIR_1)
def test_resize_to_new_table(self):
shared_proto_table_storage = ShardedProtoTableStorage()
shared_proto_table_storage.initialize_from_dir(
dir_name=self.TEST_DATA_DIR_1)
new_table = shared_proto_table_storage.resize_to_new_table(
new_size_per_shard=2, new_dir_name=self.TEST_DATA_DIR_3)
self.assertEqual(new_table.get_num_entries(), 9)
self.assertEqual(new_table.get_num_shards(), 5)
gclient.rm_dir(self.TEST_DATA_DIR_3)
class ProtoTableStorageTest(unittest.TestCase):
EXAMPLE_PROTO_1 = ProtoUtil.text_to_message(message_type=NodeSnapshot,
text_str="""
node_name: "test"
children_names: ["child_1", "child_2"]
parents_names: ["parent_1", "parent_2"]
""")
EXAMPLE_PROTO_2 = ProtoUtil.text_to_message(message_type=OperatorSnapshot,
text_str="""
operator_name: "test"
data_model: BATCH
status: SUCCEEDED
node_snapshot: {
node_name: "test"
children_names: ["child_1", "child_2"]
parents_names: ["parent_1", "parent_2"]
}
""")
EXAMPLE_PROTO_3 = ProtoUtil.text_to_message(message_type=OperatorSnapshot,
text_str="""
operator_name: "test_1"
data_model: STREAMING
status: RUNNING
node_snapshot: {
node_name: "test"
children_names: ["child_1", "child_2"]
parents_names: ["parent_1", "parent_2"]
}
""")
TEST_DATA_1 = "/galaxy/ab-d/pslx/test_data/test_proto_table_data.pb"
TEST_DATA_2 = "/galaxy/ab-d/pslx/test_data/test_proto_table_data_2.pb"
TEST_DATA_3 = "/galaxy/ab-d/pslx/test_data/test_proto_table_data_3.pb"
TEST_DATA_4 = "/galaxy/ab-d/pslx/test_data/test_proto_table_data_4.pb"
def test_read_1(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={
'key': 'test',
'message_type': NodeSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_1)
def test_read_2(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={'key': 'test'})
self.assertEqual(result_proto,
ProtoUtil.message_to_any(self.EXAMPLE_PROTO_1))
def test_read_3(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={'key': 'test1'})
self.assertIsNone(result_proto)
def test_write_1(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1}, )
result_proto = proto_table_storage.read(params={
'key': 'test',
'message_type': NodeSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_1)
gclient_ext.cp_file(self.TEST_DATA_1, self.TEST_DATA_3)
def test_write_2(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_4)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1})
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_2})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_2)
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_3})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_3)
result = proto_table_storage.read_all()
self.assertDictEqual(
result, {
'test': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_1),
'test_1': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_3),
})
self.assertEqual(proto_table_storage.get_num_entries(), 2)
gclient_ext.cp_file(self.TEST_DATA_2, self.TEST_DATA_4)
def test_write_3(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_4)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1})
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_2})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_2)
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_3},
params={
'overwrite': False,
})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_2)
self.assertEqual(proto_table_storage.get_num_entries(), 2)
gclient_ext.cp_file(self.TEST_DATA_2, self.TEST_DATA_4)
def view_proto_table():
all_urls = sorted(server_urls)
if request.method == 'POST':
try:
server_url = request.form['server_url'].strip()
proto_table_path = request.form['proto_table_path'].strip()
value_type = request.form['value_type'].strip()
module = request.form['module'].strip()
pslx_frontend_logger.info("Proto table viewer selecting url [" + server_url + '] and input path [' +
proto_table_path + '] with value type [' + value_type +
'] in module name [' + module + '].')
result = client_map[server_url]['client'].view_proto(
proto_file_path=proto_table_path,
message_type="ProtoTable",
module="",
root_certificate=client_map[server_url]['root_certificate']
)
value_type = ProtoUtil.infer_message_type_from_str(
message_type_str=value_type,
modules=module
)
result_content = ProtoUtil.text_to_message(message_type=ProtoTable, text_str=result['proto_content'])
proto_contents = []
result_content = dict(result_content.data)
for key in sorted(result_content.keys()):
proto_val = ProtoUtil.any_to_message(
message_type=value_type,
any_message=result_content[key]
)
try:
proto_contents.append(
{
'key': key,
'val': ProtoUtil.message_to_text(proto_message=proto_val),
}
)
except Exception as err:
pslx_frontend_logger.error("Proto table viewer Parsing proto with error " + str(err) + '.')
proto_contents.append(
{
'key': key,
'val': str(proto_val),
}
)
all_urls.remove(server_url)
all_urls = [server_url] + all_urls
return render_template(
'proto_table_viewer.html',
proto_contents=proto_contents,
server_urls=all_urls
)
except Exception as err:
pslx_frontend_logger.error("Got error rendering proto_table_viewer.html: " + str(err) + '.')
return render_template(
'proto_table_viewer.html',
proto_contents=[],
server_urls=all_urls
)
else:
return render_template(
'proto_table_viewer.html',
proto_contents=[],
selected_server_url='',
server_urls=server_urls
)