# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A simple smoke test that runs these examples for 1 training iteraton."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import pandas as pd

from six.moves import StringIO

import custom_estimator
import premade_estimator

FOUR_LINES = "\n".join([
    "1,52.40, 2823,152,2",
    "164, 99.80,176.60,66.20,1",
    "176,2824, 136,3.19,0",
    "2,177.30,66.30, 53.10,1",])

def four_lines_data():
  text = StringIO(FOUR_LINES)

  df = pd.read_csv(text, names=premade_estimator.COLUMNS)

  xy = (df, df.pop("Species"))
  return xy, xy


class RegressionTest(tf.test.TestCase):
  """Test the regression examples in this directory."""

  @tf.test.mock.patch.dict(premade_estimator.__dict__,
                           {"load_data": four_lines_data})
  def test_premade_estimator(self):
    premade_estimator.main([None, "--train_steps=1"])

  @tf.test.mock.patch.dict(custom_estimator.__dict__,
                           {"load_data": four_lines_data})
  def test_custom_estimator(self):
    custom_estimator.main([None, "--train_steps=1"])

if __name__ == "__main__":
  tf.test.main()
