from tensorflow.python.framework import constant_op



def _ranking_train_input_fn():
  features = {
      "a.f1": constant_op.constant([[3.], [0.3], [1.]]),
      "a.f2": constant_op.constant([[0.1], [3.], [1.]]),
      "b.f1": constant_op.constant([[13.], [0.4], [5.]]),
      "b.f2": constant_op.constant([[1.], [3.], [0.01]]),
  }
  label = constant_op.constant([[0], [0], [1]], dtype=dtypes.int32)
  return features, label


def _eval_input_fn():
  features = {"x": constant_op.constant([[1.], [2.], [2.]])}
  label = constant_op.constant([[0], [1], [1]], dtype=dtypes.int32)
  return features, label


def _infer_ranking_train_input_fn():
  features = {
      "f1": constant_op.constant([[3.], [2], [1.]]),
      "f2": constant_op.constant([[0.1], [3.], [1.]])
  }
  return features, None


class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase):

  def setUp(self):
