def testGetNodeMapBasic(self):
meta_graph_def = meta_graph_pb2.MetaGraphDef()
meta_graph_def.collection_def[
'my_collection/%s' %
encoding.KEY_SUFFIX].bytes_list.value[:] = map(
encoding.encode_key, ['alpha', 'bravo', 'charlie'])
meta_graph_def.collection_def[
'my_collection/fruits'].bytes_list.value[:] = [
b'apple', b'banana', b'cherry'
]
expected = {
'alpha': {
'fruits': b'apple'
},
'bravo': {
'fruits': b'banana'
},
'charlie': {
'fruits': b'cherry'
}
}
self.assertDictEqual(
expected,
graph_ref.get_node_map(meta_graph_def, 'my_collection',
['fruits']))
def testGetNodeMapMultiple(self):
meta_graph_def = meta_graph_pb2.MetaGraphDef()
meta_graph_def.collection_def[
'my_collection/%s' %
encoding.KEY_SUFFIX].bytes_list.value[:] = map(
encoding.encode_key, ['alpha', 'bravo', 'charlie'])
meta_graph_def.collection_def[
'my_collection/fruits'].bytes_list.value[:] = [
'apple', 'banana', 'cherry'
]
meta_graph_def.collection_def[
'my_collection/animals'].bytes_list.value[:] = [
'aardvark', 'badger', 'camel'
]
expected = {
'alpha': {
'fruits': 'apple',
'animals': 'aardvark'
},
'bravo': {
'fruits': 'banana',
'animals': 'badger'
},
'charlie': {
'fruits': 'cherry',
'animals': 'camel'
}
}
self.assertDictEqual(
expected,
graph_ref.get_node_map(meta_graph_def, 'my_collection',
['fruits', 'animals']))
def testGetNodeMapEmpty(self):
meta_graph_def = meta_graph_pb2.MetaGraphDef()
self.assertDictEqual({},
graph_ref.get_node_map(meta_graph_def,
'my_collection',
['fruits']))
