Create smoke test for Snowflake notebooks. (#1619)

tests/e2e/data/environment.yml

@@ -0,0 +1,8 @@
+name: app_environment
+channels:
+  - snowflake
+dependencies:
+  - trulens-connectors-snowflake=*
+  - trulens-core=*
+  - trulens-feedback=*
+  - trulens-providers-cortex=*

tests/e2e/data/simple.ipynb

@@ -0,0 +1,79 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import trulens.connectors.snowflake\n",
+    "\n",
+    "print(trulens.connectors.snowflake.__version__)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import trulens.core\n",
+    "\n",
+    "print(trulens.core.__version__)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import trulens.feedback\n",
+    "\n",
+    "print(trulens.feedback.__version__)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import trulens.providers.cortex\n",
+    "\n",
+    "print(trulens.providers.cortex.__version__)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# State the obvious.\n",
+    "\n",
+    "print(\"Kojikun is the world's cutest baby!\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "trulens_3_11",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/e2e/test_snowflake_connection.py

@@ -40,7 +40,7 @@ def test_connecting_to_premade_schema(self):
         schema_name += str(uuid.uuid4()).replace("-", "_")
         schema_name = schema_name.upper()
         self.assertNotIn(schema_name, self.list_schemas())
-        self._snowflake_schemas_to_delete.append(schema_name)
+        self._snowflake_schemas_to_delete.add(schema_name)
         self.run_query(f"CREATE SCHEMA {schema_name}")
         self.run_query(
             f"CREATE TABLE {self._database}.{schema_name}.MY_TABLE (TEST_COLUMN NUMBER)"

tests/e2e/test_snowflake_notebooks.py

@@ -0,0 +1,48 @@
+from unittest import main
+import uuid
+
+from tests.util.snowflake_test_case import SnowflakeTestCase
+
+_STAGE_NAME = "SNOWFLAKE_NOTEBOOKS"
+
+
+class TestSnowflakeNotebooks(SnowflakeTestCase):
+    def test_simple_notebook(self) -> None:
+        self._upload_and_run_notebook(
+            "test_simple_notebook",
+            "simple",
+            "tests/e2e/data/",
+        )
+
+    def _upload_and_run_notebook(
+        self, schema_base_name: str, name: str, path: str
+    ) -> None:
+        try:
+            schema_name = (
+                f"{schema_base_name}_{str(uuid.uuid4()).replace('-', '_')}"
+            )
+            self.create_and_use_schema(schema_name)
+            self.run_query(f"CREATE STAGE {_STAGE_NAME}")
+            self.run_query(
+                f"PUT file://{path}/{name}.ipynb @{_STAGE_NAME} AUTO_COMPRESS = FALSE"
+            )
+            self.run_query(
+                f"PUT file://{path}/environment.yml @{_STAGE_NAME} AUTO_COMPRESS = FALSE"
+            )
+            self.run_query(
+                f"""
+                CREATE NOTEBOOK {name}
+                FROM '@{_STAGE_NAME}'
+                MAIN_FILE = '{name}.ipynb'
+                QUERY_WAREHOUSE = {self._snowflake_connection_parameters["warehouse"]}
+                """
+            )
+            self.run_query(f"ALTER NOTEBOOK {name} ADD LIVE VERSION FROM LAST")
+            self.run_query(f"EXECUTE NOTEBOOK {name}()")
+        finally:
+            self.run_query(f"DROP STAGE IF EXISTS {_STAGE_NAME}")
+            self.run_query(f"DROP NOTEBOOK IF EXISTS {name}")
+
+
+if __name__ == "__main__":
+    main()

tests/util/snowflake_test_case.py

@@ -4,7 +4,7 @@
 
 import logging
 import os
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 from unittest import TestCase
 from unittest import main
 import uuid
@@ -30,7 +30,7 @@ def setUp(self):
         self._snowpark_session = Session.builder.configs(
             self._snowflake_connection_parameters
         ).create()
-        self._snowflake_schemas_to_delete = []
+        self._snowflake_schemas_to_delete = set()
 
     def tearDown(self):
         # [HACK!] Clean up any instances of `TruSession` so tests don't interfere with each other.
@@ -92,7 +92,7 @@ def get_session(
         self._schema = self._schema.upper()
         if not schema_already_exists:
             self.assertNotIn(self._schema, self.list_schemas())
-            self._snowflake_schemas_to_delete.append(self._schema)
+            self._snowflake_schemas_to_delete.add(self._schema)
         if not connect_via_snowpark_session:
             connector = snowflake_connector.SnowflakeConnector(
                 schema=self._schema,
@@ -102,9 +102,7 @@ def get_session(
             )
         else:
             if not schema_already_exists:
-                snowflake_connector.SnowflakeConnector._create_snowflake_schema_if_not_exists(
-                    self._snowpark_session, self._schema
-                )
+                self.create_and_use_schema(self._schema)
             connector = snowflake_connector.SnowflakeConnector(
                 snowpark_session=self._snowpark_session,
                 init_server_side=True,
@@ -114,8 +112,15 @@ def get_session(
         self.assertIn(self._schema, self.list_schemas())
         return session
 
-    def run_query(self, q: str) -> List[Row]:
-        return self._snowpark_session.sql(q).collect()
+    def run_query(
+        self, q: str, bindings: Optional[List[Any]] = None
+    ) -> List[Row]:
+        return self._snowpark_session.sql(q, bindings).collect()
+
+    def create_and_use_schema(self, schema_name: str) -> None:
+        self.run_query("CREATE SCHEMA IDENTIFIER(?)", [schema_name])
+        self._snowflake_schemas_to_delete.add(schema_name)
+        self._snowpark_session.use_schema(schema_name)
 
 
 if __name__ == "__main__":