def setup_test_df():
"""Create a dataframe populated with some test columns."""
df = learn.DataFrame()
df["a"] = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockTwoOutputTransform("iue", "eui", "snt"), "out1")
df["b"] = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockTwoOutputTransform("iue", "eui", "snt"), "out2")
df["c"] = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockTwoOutputTransform("iue", "eui", "snt"), "out1")
return df
def test_build_multiple_output(self):
col = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockTwoOutputTransform("thb", "nth", "snt"), "out2")
result = col.build()
expected = "Fake Tensor 2"
self.assertEqual(expected, result)
def test_build_single_output(self):
col = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockOneOutputTransform("thb", "nth"), "out1")
result = col.build()
expected = "Fake Tensor 1"
self.assertEqual(expected, result)
def test_repr(self):
col = learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockTwoOutputTransform("thb", "nth", "snt"), "qux")
# note params are sorted by name
expected = ("MockTransform({'param_one': 'thb', 'param_three': 'snt', "
"'param_two': 'nth'})"
"(foobar)[qux]")
self.assertEqual(expected, repr(col))
def create_no_output_column():
return learn.TransformedColumn(
[mocks.MockColumn("foobar", [])],
mocks.MockZeroOutputTransform("thb", "nth"), None)