def test_nested_value_provider_wrap_runtime(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg15', help='This keyword argument is a value provider') # set at runtime options = UserDefinedOptions([]) vp = NestedValueProvider(options.vpt_vp_arg15, lambda x: x + x) self.assertFalse(vp.is_accessible()) RuntimeValueProvider.set_runtime_options({'vpt_vp_arg15': 'abc'}) self.assertTrue(vp.is_accessible()) self.assertEqual(vp.get(), 'abcabc')
def GetFixedColumnsProvider(self): def DateTransform(yyyymmdd): if yyyymmdd == 'yesterday': d = (datetime.datetime.utcnow() - datetime.timedelta(days=1)).date() else: d = datetime.date(int(yyyymmdd[:4]), int(yyyymmdd[4:6]), int(yyyymmdd[6:8])) return beam.Row(date=d) return NestedValueProvider(self.end_date, DateTransform)
def test_nested_value_provider_caches_value(self): mock_fn = Mock() def translator(x): mock_fn() return x vp = NestedValueProvider(StaticValueProvider(int, 1), translator) vp.get() self.assertEqual(mock_fn.call_count, 1) vp.get() self.assertEqual(mock_fn.call_count, 1)
def test_nested_value_provider_wrap_static(self): vp = NestedValueProvider(StaticValueProvider(int, 1), lambda x: x + 1) self.assertTrue(vp.is_accessible()) self.assertEqual(vp.get(), 2)