def testGetCommonFnArgs(self):
    source_data_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'testdata')
    examples = standard_artifacts.Examples()
    examples.uri = os.path.join(source_data_dir, 'transform/transformed_examples')
    examples.split_names = artifact_utils.encode_split_names(['train', 'eval'])
    transform_output = standard_artifacts.TransformGraph()
    transform_output.uri = os.path.join(source_data_dir, 'transform/transform_graph')
    schema = standard_artifacts.Schema()
    schema.uri = os.path.join(source_data_dir, 'schema_gen')
    base_model = standard_artifacts.Model()
    base_model.uri = os.path.join(source_data_dir, 'trainer/previous')
    input_dict = {standard_component_specs.EXAMPLES_KEY: [examples], standard_component_specs.TRANSFORM_GRAPH_KEY: [transform_output], standard_component_specs.SCHEMA_KEY: [schema], standard_component_specs.BASE_MODEL_KEY: [base_model]}
    exec_properties = {'train_args': proto_utils.proto_to_json(trainer_pb2.TrainArgs(num_steps=1000)), 'eval_args': proto_utils.proto_to_json(trainer_pb2.EvalArgs(num_steps=500))}
    fn_args = fn_args_utils.get_common_fn_args(input_dict, exec_properties, 'tempdir')
    self.assertEqual(fn_args.working_dir, 'tempdir')
    self.assertEqual(fn_args.train_steps, 1000)
    self.assertEqual(fn_args.eval_steps, 500)
    self.assertLen(fn_args.train_files, 1)
    self.assertEqual(fn_args.train_files[0], os.path.join(examples.uri, 'Split-train', '*'))
    self.assertLen(fn_args.eval_files, 1)
    self.assertEqual(fn_args.eval_files[0], os.path.join(examples.uri, 'Split-eval', '*'))
    self.assertEqual(fn_args.schema_path, os.path.join(schema.uri, 'schema.pbtxt'))
    self.assertRegexpMatches(fn_args.base_model, os.path.join(base_model.uri, 'Format-(Servo|Serving)/export/chicago-taxi/\\d+'))
    self.assertEqual(fn_args.transform_graph_path, transform_output.uri)
    self.assertIsInstance(fn_args.data_accessor, fn_args_utils.DataAccessor)
def testGetCommonFnArgs(self):
    source_data_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'testdata')
    examples = standard_artifacts.Examples()
    examples.uri = os.path.join(source_data_dir, 'transform/transformed_examples')
    examples.split_names = artifact_utils.encode_split_names(['train', 'eval'])
    transform_output = standard_artifacts.TransformGraph()
    transform_output.uri = os.path.join(source_data_dir, 'transform/transform_graph')
    schema = standard_artifacts.Schema()
    schema.uri = os.path.join(source_data_dir, 'schema_gen')
    base_model = standard_artifacts.Model()
    base_model.uri = os.path.join(source_data_dir, 'trainer/previous')
    input_dict = {standard_component_specs.EXAMPLES_KEY: [examples], standard_component_specs.TRANSFORM_GRAPH_KEY: [transform_output], standard_component_specs.SCHEMA_KEY: [schema], standard_component_specs.BASE_MODEL_KEY: [base_model]}
    exec_properties = {'train_args': proto_utils.proto_to_json(trainer_pb2.TrainArgs(num_steps=1000)), 'eval_args': proto_utils.proto_to_json(trainer_pb2.EvalArgs(num_steps=500))}
    fn_args = fn_args_utils.get_common_fn_args(input_dict, exec_properties, 'tempdir')
    self.assertEqual(fn_args.working_dir, 'tempdir')
    self.assertEqual(fn_args.train_steps, 1000)
    self.assertEqual(fn_args.eval_steps, 500)
    self.assertLen(fn_args.train_files, 1)
    self.assertEqual(fn_args.train_files[0], os.path.join(examples.uri, 'Split-train', '*'))
    self.assertLen(fn_args.eval_files, 1)
    self.assertEqual(fn_args.eval_files[0], os.path.join(examples.uri, 'Split-eval', '*'))
    self.assertEqual(fn_args.schema_path, os.path.join(schema.uri, 'schema.pbtxt'))
    self.assertRegexpMatches(fn_args.base_model, os.path.join(base_model.uri, 'Format-(Servo|Serving)/export/chicago-taxi/\\d+'))
    self.assertEqual(fn_args.transform_graph_path, transform_output.uri)
    self.assertIsInstance(fn_args.data_accessor, fn_args_utils.DataAccessor)