make test that "works"

Montana Low · Montana Low · commit c2de3d844566 · 2022-04-14T10:56:21.000-07:00
diff --git a/pgml/pgml/model.py b/pgml/pgml/model.py
@@ -41,7 +41,7 @@ def find(cls, id: int):
             FROM pgml.projects 
             WHERE id = {q(id)}
         """, 1)
-        if (result.nrows == 0):
+        if (len(result) == 0):
             return None
 
         project = Project()
@@ -71,7 +71,7 @@ def find_by_name(cls, name: str):
             FROM pgml.projects 
             WHERE name = {q(name)}
         """, 1)
-        if (result.nrows == 0):
+        if (len(result)== 0):
             return None
 
         project = Project()
@@ -159,7 +159,7 @@ def create(cls, relation_name: str, y_column_name: str, test_size: float or int,
             SET status = 'created' 
             WHERE id = {q(snapshot.id)} 
             RETURNING *
-        """)[0])
+        """, 1)[0])
         return snapshot
 
     def data(self):
@@ -172,8 +172,9 @@ def data(self):
             FROM pgml."snapshot_{self.id}"
         """)
 
+        print(data)
         # Sanity check the data
-        if data.nrows == 0:
+        if len(data) == 0:
             PgMLException(
                 f"Relation `{self.relation_name}` contains no rows. Did you pass the correct `relation_name`?"
             )
@@ -272,7 +273,7 @@ def find_deployed(cls, project_id: int):
             ORDER by deployments.created_at DESC
             LIMIT 1
         """)
-        if (result.nrows == 0):
+        if (len(result) == 0):
             return None
 
         model = Model()
diff --git a/pgml/tests/plpy.py b/pgml/tests/plpy.py
@@ -0,0 +1,16 @@
+from collections import deque
+
+execute_results = deque()
+
+def quote_literal(literal):
+    return "'" + literal + "'"
+
+def execute(sql, lines = 0):
+    if len(execute_results) > 0:
+        result = execute_results.popleft()
+        return result
+    else: 
+        return []
+
+def add_mock_result(result):
+    execute_results.append(result)
diff --git a/pgml/tests/test_model.py b/pgml/tests/test_model.py
@@ -0,0 +1,65 @@
+# stub out plpy
+from . import plpy
+import sys
+sys.modules['plpy'] = plpy
+
+import time
+import unittest
+from pgml import model
+
+class TestModel(unittest.TestCase):
+    def test_the_world(self):
+        plpy.add_mock_result(
+            [{"id": 1, "name": "Test", "objective": "regression", "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "relation_name": "test", "y_column_name": "test_y", "test_size": 0.1, "test_sampling": "random", "status": "new", "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            "OK"
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "relation_name": "test", "y_column_name": "test_y", "test_size": 0.1, "test_sampling": "random", "status": "created", "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            [
+                {"a": 1, "b": 2, "test_y": 3},
+                {"a": 2, "b": 3, "test_y": 4},
+                {"a": 3, "b": 4, "test_y": 5},
+            ]
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            [
+                {"a": 1, "b": 2, "test_y": 3},
+                {"a": 2, "b": 3, "test_y": 4},
+                {"a": 3, "b": 4, "test_y": 5},
+            ]
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+        
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+        plpy.add_mock_result(
+            [
+                {"a": 1, "b": 2, "test_y": 3},
+                {"a": 2, "b": 3, "test_y": 4},
+                {"a": 3, "b": 4, "test_y": 5},
+            ]
+        )
+        plpy.add_mock_result(
+            [{"id": 1, "project_id": 1, "snapshot_id": 1, "algorithm_name": "linear", "status": "new", "r2_score": None, "mean_squared_error": None, "pickle": None, "created_at": time.time(), "updated_at": time.time()}]
+        )
+        model.train("Test", "regression", "test", "test_y")
diff --git a/pgml/tests/test_train.py b/pgml/tests/test_train.py
