From 68cf236944eba157827238584dd6e9b5b6f1ec76 Mon Sep 17 00:00:00 2001
From: Joe Jevnik <joe@quantopian.com>
Date: Tue, 29 Dec 2015 10:12:39 -0500
Subject: [PATCH] TST: Add test case for adding columns in subclass

---
 tests/pipeline/test_term.py | 43 +++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/tests/pipeline/test_term.py b/tests/pipeline/test_term.py
index 94410c5d..f689bc33 100644
--- a/tests/pipeline/test_term.py
+++ b/tests/pipeline/test_term.py
@@ -32,6 +32,10 @@ class SubDataSet(SomeDataSet):
     pass
 
 
+class SubDataSetNewCol(SomeDataSet):
+    qux = Column(float64_dtype)
+
+
 class SomeFactor(Factor):
     dtype = float64_dtype
     window_length = 5
@@ -336,8 +340,8 @@ class SubDataSetTestCase(TestCase):
             column.name: column for column in SubDataSet.columns
         }
         self.assertEqual(
-            set(some_dataset_map),
-            set(sub_dataset_map),
+            {column.name for column in SomeDataSet.columns},
+            {column.name for column in SubDataSet.columns},
         )
         for k, some_dataset_column in some_dataset_map.items():
             sub_dataset_column = sub_dataset_map[k]
@@ -353,3 +357,38 @@ class SubDataSetTestCase(TestCase):
                 'subclass column %r should have the same dtype as the parent' %
                 k,
             )
+
+    def test_add_column(self):
+        some_dataset_map = {
+            column.name: column for column in SomeDataSet.columns
+        }
+        sub_dataset_new_col_map = {
+            column.name: column for column in SubDataSetNewCol.columns
+        }
+        sub_col_names = {column.name for column in SubDataSetNewCol.columns}
+
+        # check our extra col
+        self.assertIn('qux', sub_col_names)
+        self.assertEqual(
+            sub_dataset_new_col_map['qux'].dtype,
+            float64_dtype,
+        )
+
+        self.assertEqual(
+            {column.name for column in SomeDataSet.columns},
+            sub_col_names - {'qux'},
+        )
+        for k, some_dataset_column in some_dataset_map.items():
+            sub_dataset_column = sub_dataset_new_col_map[k]
+            self.assertIsNot(
+                some_dataset_column,
+                sub_dataset_column,
+                'subclass column %r should not have the same identity as'
+                ' the parent' % k,
+            )
+            self.assertEqual(
+                some_dataset_column.dtype,
+                sub_dataset_column.dtype,
+                'subclass column %r should have the same dtype as the parent' %
+                k,
+            )