from absl.testing import absltest
from absl.testing import parameterized

import numpy as np
import tensorflow as tf

from core.ensemble.dataset_lib import build_dataset_from_numerical_data


class MyTestCase(tf.test.TestCase, parameterized.TestCase):
    def test_something(self):
        x = np.random.rand(5,3)
        data = build_dataset_from_numerical_data(x, batch_size=1)
        for i, _x in enumerate(data):
            self.assertAllClose(_x[0, :], x[i, :])

        y = np.random.choice(2, 5)
        dataset = build_dataset_from_numerical_data((x, y), batch_size=1)
        for i, (_x, _y) in enumerate(dataset):
            self.assertAllClose(_x[0, :], x[i, :])
            self.assertEqual(_y, y[i])


if __name__ == '__main__':
    absltest.main()
