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']))