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 )