Merge pull request #45 from vanna-ai/create-dataset-bug

fix create dataset bug
vanna-ai · Jul 25, 2023 · 9c4c50f · 9c4c50f
2 parents cc25d4e + 0dc0daf
commit 9c4c50f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/src/vanna/__init__.py b/src/vanna/__init__.py
@@ -240,6 +240,9 @@ def create_dataset(dataset: str, db_type: str) -> bool:
 
     status = Status(**d['result'])
 
+    if status.success:
+        __org = dataset
+
     return status.success
 
 def add_user_to_dataset(dataset: str, email: str, is_admin: bool) -> bool:

diff --git a/tests/test_vanna.py b/tests/test_vanna.py
@@ -199,4 +199,11 @@ def test_remove_training_data():
         rv = vn.remove_training_data(row['id'])
         assert rv == True
 
-        assert vn.get_training_data().shape[0] == 2-index
+        assert vn.get_training_data().shape[0] == 2-index
+
+def test_create_dataset_and_add_user():
+    created = vn.create_dataset('test_org2', 'Snowflake')
+    assert created == True
+
+    added = vn.add_user_to_dataset(dataset='test_org2', email="[email protected]", is_admin=False)
+    assert added == True