From 47e866cfc6be899ee690d1ea691eef9ac9d1bedf Mon Sep 17 00:00:00 2001
From: jianfengmao <jianfengmao@deephaven.io>
Date: Wed, 20 Mar 2024 21:56:46 -0600
Subject: [PATCH] Respond to latest review comments

---
 py/server/deephaven/table_factory.py  |  3 +++
 py/server/tests/test_table_factory.py | 16 +++++++++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/py/server/deephaven/table_factory.py b/py/server/deephaven/table_factory.py
index 445cf88874c..7efff4b534c 100644
--- a/py/server/deephaven/table_factory.py
+++ b/py/server/deephaven/table_factory.py
@@ -29,6 +29,7 @@
 _JTableDefinition = jpy.get_type("io.deephaven.engine.table.TableDefinition")
 _JTable = jpy.get_type("io.deephaven.engine.table.Table")
 _J_INPUT_TABLE_ATTRIBUTE = _JTable.INPUT_TABLE_ATTRIBUTE
+_J_InputTableUpdater = jpy.get_type("io.deephaven.engine.util.input.InputTableUpdater")
 _JRingTableTools = jpy.get_type("io.deephaven.engine.table.impl.sources.ring.RingTableTools")
 _JSupplier = jpy.get_type('java.util.function.Supplier')
 _JFunctionGeneratedTableFactory = jpy.get_type("io.deephaven.engine.table.impl.util.FunctionGeneratedTableFactory")
@@ -240,6 +241,8 @@ def __init__(self, j_table: jpy.JType):
         self.j_input_table = self.j_table.getAttribute(_J_INPUT_TABLE_ATTRIBUTE)
         if not self.j_input_table:
             raise DHError("the provided table input is not suitable for input tables.")
+        if not _J_InputTableUpdater.jclass.isInstance(self.j_input_table):
+            raise DHError("the provided table's InputTable attribute type is not of InputTableUpdater type.")
 
     def add(self, table: Table) -> None:
         """Synchronously writes rows from the provided table to this input table. If this is a keyed input table, added rows with keys
diff --git a/py/server/tests/test_table_factory.py b/py/server/tests/test_table_factory.py
index 8a526ceddc1..3b1cfe55062 100644
--- a/py/server/tests/test_table_factory.py
+++ b/py/server/tests/test_table_factory.py
@@ -375,10 +375,12 @@ def test_input_table(self):
 
         with self.subTest("custom input table creation"):
             place_holder_input_table = empty_table(1).update_view(["Key=`A`", "Value=10"]).with_attributes({_JTable.INPUT_TABLE_ATTRIBUTE: "Placeholder IT"}).j_table
-            # Confirming no error.
-            it = InputTable(place_holder_input_table)
 
-            self.assertTrue(isinstance(_wrapper.wrap_j_object(place_holder_input_table), InputTable))
+            with self.assertRaises(DHError) as cm:
+                InputTable(place_holder_input_table)
+            self.assertIn("not of InputTableUpdater type", str(cm.exception))
+
+            self.assertTrue(isinstance(_wrapper.wrap_j_object(place_holder_input_table), Table))
 
 
     def test_ring_table(self):
@@ -468,8 +470,12 @@ def test_j_input_wrapping(self):
         col_defs = {c.name: c.data_type for c in t.columns}
         append_only_input_table = input_table(col_defs=col_defs)
 
-        t = _wrapper.wrap_j_object(append_only_input_table.j_table)
-        self.assertTrue(isinstance(t, InputTable))
+        it = _wrapper.wrap_j_object(append_only_input_table.j_table)
+        self.assertTrue(isinstance(it, InputTable))
+
+        t = _wrapper.wrap_j_object(t.j_object)
+        self.assertFalse(isinstance(t, InputTable))
+        self.assertTrue(isinstance(t, Table))
 
 
 if __name__ == '__main__':