def testRegisterContexts(self):
node_contexts = pipeline_pb2.NodeContexts()
self.load_proto_from_text(
os.path.join(self._testdata_dir, 'node_context_spec.pbtxt'),
node_contexts)
with metadata.Metadata(connection_config=self._connection_config) as m:
context_lib.prepare_contexts(metadata_handler=m,
node_contexts=node_contexts)
# Duplicated call should succeed.
contexts = context_lib.prepare_contexts(
metadata_handler=m, node_contexts=node_contexts)
got_context_type_one = m.store.get_context_type(
'my_context_type_one')
got_context_type_one.ClearField('id')
self.assertProtoEquals(
"""
name: 'my_context_type_one'
""", got_context_type_one)
got_context_type_two = m.store.get_context_type(
'my_context_type_two')
got_context_type_two.ClearField('id')
self.assertProtoEquals(
"""
name: 'my_context_type_two'
""", got_context_type_two)
self.assertEqual(
contexts[0],
m.store.get_context_by_type_and_name('my_context_type_one',
'my_context_one'))
self.assertEqual(
contexts[1],
m.store.get_context_by_type_and_name('my_context_type_one',
'my_context_two'))
self.assertEqual(
contexts[2],
m.store.get_context_by_type_and_name('my_context_type_two',
'my_context_three'))
self.assertEqual(
contexts[0].custom_properties['property_a'].int_value, 1)
self.assertEqual(
contexts[1].custom_properties['property_a'].int_value, 2)
self.assertEqual(
contexts[2].custom_properties['property_a'].int_value, 3)
self.assertEqual(
contexts[2].custom_properties['property_b'].string_value, '4')
def _generate_contexts(self, metadata_handler):
context_spec = pipeline_pb2.NodeContexts()
text_format.Parse(
"""
contexts {
type {name: 'pipeline_context'}
name {
field_value {string_value: 'my_pipeline'}
}
}
contexts {
type {name: 'component_context'}
name {
field_value {string_value: 'my_component'}
}
}""", context_spec)
return context_lib.prepare_contexts(metadata_handler, context_spec)
def testRegisterContexts(self):
node_contexts = pipeline_pb2.NodeContexts()
self.load_proto_from_text('node_context_spec.pbtxt', node_contexts)
with metadata.Metadata(connection_config=self._connection_config) as m:
context_lib.register_contexts_if_not_exists(
metadata_handler=m, node_contexts=node_contexts)
# Duplicated call should succeed.
contexts = context_lib.register_contexts_if_not_exists(
metadata_handler=m, node_contexts=node_contexts)
self.assertProtoEquals(
"""
id: 1
name: 'my_context_type_one'
""", m.store.get_context_type('my_context_type_one'))
self.assertProtoEquals(
"""
id: 2
name: 'my_context_type_two'
""", m.store.get_context_type('my_context_type_two'))
self.assertEqual(
contexts[0],
m.store.get_context_by_type_and_name('my_context_type_one',
'my_context_one'))
self.assertEqual(
contexts[1],
m.store.get_context_by_type_and_name('my_context_type_one',
'my_context_two'))
self.assertEqual(
contexts[2],
m.store.get_context_by_type_and_name('my_context_type_two',
'my_context_three'))
self.assertEqual(
contexts[0].custom_properties['property_a'].int_value, 1)
self.assertEqual(
contexts[1].custom_properties['property_a'].int_value, 2)
self.assertEqual(
contexts[2].custom_properties['property_a'].int_value, 3)
self.assertEqual(
contexts[2].custom_properties['property_b'].string_value, '4')
def testRegisterContexts(self):
node_contexts = pipeline_pb2.NodeContexts()
self.load_proto_from_text(
os.path.join(self._testdata_dir, 'node_context_spec.pbtxt'),
node_contexts)
with metadata.Metadata(connection_config=self._connection_config) as m:
context_lib.prepare_contexts(metadata_handler=m,
node_contexts=node_contexts)
# Duplicated call should succeed.
contexts = context_lib.prepare_contexts(
metadata_handler=m, node_contexts=node_contexts)
self.assertProtoEquals(
"""
id: 1
name: 'my_context_type_one'
""", m.store.get_context_type('my_context_type_one'))
self.assertProtoEquals(
"""
id: 2
name: 'my_context_type_two'
""", m.store.get_context_type('my_context_type_two'))
self.assertProtoEquals(
"""
type_id: 1
name: "my_context_one"
custom_properties {
key: "property_a"
value {
int_value: 1
}
}
""", contexts[0])
self.assertProtoEquals(
"""
type_id: 1
name: "my_context_two"
custom_properties {
key: "property_a"
value {
int_value: 2
}
}
""", contexts[1])
self.assertProtoEquals(
"""
type_id: 2
name: "my_context_three"
custom_properties {
key: "property_a"
value {
int_value: 3
}
}
custom_properties {
key: "property_b"
value {
string_value: '4'
}
}
""", contexts[2])
self.assertEqual(
contexts[0].custom_properties['property_a'].int_value, 1)
self.assertEqual(
contexts[1].custom_properties['property_a'].int_value, 2)
self.assertEqual(
contexts[2].custom_properties['property_a'].int_value, 3)
self.assertEqual(
contexts[2].custom_properties['property_b'].string_value, '4')