mito-ai 0.1.50__py3-none-any.whl

mito_ai/tests/db/connections_test.py

@@ -0,0 +1,102 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+from mito_ai.tests.db.test_db_constants import SNOWFLAKE
+from mito_ai.tests.conftest import TOKEN
+
+
+# --- ADD CONNECTION ---
+
+
+def test_add_connection_with_auth(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=SNOWFLAKE,
+    )
+    if response.status_code != 200:
+        print(f"Error response: {response.text}")
+    assert response.status_code == 200
+
+
+def test_add_connection_with_no_auth(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        json=SNOWFLAKE,
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_add_connection_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token incorrect-token"},
+        json=SNOWFLAKE,
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- GET CONNECTIONS ---
+
+
+def test_get_connections_with_auth(jp_base_url: str, first_connection_id: str) -> None:
+    response = requests.get(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Ensure the response has the correct number of connections
+    assert len(response.json()) == 1
+    # Ensure the response has the correct connection details
+    assert response.json()[first_connection_id] == SNOWFLAKE
+
+
+def test_get_connections_with_no_auth(jp_base_url: str) -> None:
+    response = requests.get(jp_base_url + "/mito-ai/db/connections")
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_connections_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- DELETE CONNECTION ---
+
+
+def test_delete_connection_with_no_auth(jp_base_url: str, first_connection_id: str) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/connections/{first_connection_id}",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_delete_connection_with_incorrect_auth(
+    jp_base_url: str, first_connection_id: str
+) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/connections/{first_connection_id}",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_delete_connection_with_auth(
+    jp_base_url: str, first_connection_id: str
+) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/connections/{first_connection_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    # Ensure the connection was deleted
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert len(response.json()) == 0

mito_ai/tests/db/mssql_test.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+from mito_ai.tests.db.test_db_constants import MSSQL_CONNECTION_DETAILS
+from mito_ai.tests.conftest import TOKEN
+
+# To create a postgres database, run the following command:
+# docker-compose -f mito_ai/docker/mssql/compose.yml up
+# and then, to delete the database, run the following command:
+# docker-compose -f mito_ai/docker/mssql/compose.yml down -v
+
+
+def test_add_mssql_connection(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=MSSQL_CONNECTION_DETAILS,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1

mito_ai/tests/db/mysql_test.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+from mito_ai.tests.db.test_db_constants import MYSQL_CONNECTION_DETAILS
+from mito_ai.tests.conftest import TOKEN
+
+# To create a mysql database, run the following command:
+# docker-compose -f mito_ai/docker/mysql/compose.yml up
+# and then, to delete the database, run the following command:
+# docker-compose -f mito_ai/docker/mysql/compose.yml down -v
+
+
+def test_add_mysql_connection(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=MYSQL_CONNECTION_DETAILS,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1

mito_ai/tests/db/oracle_test.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+from mito_ai.tests.db.test_db_constants import ORACLE_CONNECTION_DETAILS
+from mito_ai.tests.conftest import TOKEN
+
+# To create a postgres database, run the following command:
+# docker-compose -f mito_ai/docker/postgres/compose.yml up
+# and then, to delete the database, run the following command:
+# docker-compose -f mito_ai/docker/postgres/compose.yml down -v
+
+
+def test_add_oracle_connection(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=ORACLE_CONNECTION_DETAILS,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1

mito_ai/tests/db/postgres_test.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import requests
+from mito_ai.tests.db.test_db_constants import POSTGRES_CONNECTION_DETAILS
+from mito_ai.tests.conftest import TOKEN
+
+# To create a postgres database, run the following command:
+# docker-compose -f mito_ai/docker/postgres/compose.yml up
+# and then, to delete the database, run the following command:
+# docker-compose -f mito_ai/docker/postgres/compose.yml down -v
+
+
+def test_add_postgres_connection(jp_base_url: str) -> None:
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=POSTGRES_CONNECTION_DETAILS,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1

mito_ai/tests/db/schema_test.py

@@ -0,0 +1,93 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import requests
+from mito_ai.tests.db.test_db_constants import SNOWFLAKE
+from mito_ai.tests.conftest import TOKEN
+
+# --- GET SCHEMAS ---
+
+
+def test_get_schemas_with_auth(jp_base_url: str) -> None:
+    # Add a connection to get a schema
+    response = requests.post(
+        jp_base_url + f"/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=SNOWFLAKE,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1
+
+
+def test_ensure_no_application_databases(jp_base_url: str, first_connection_id: str) -> None:
+    # We currently don't support Snowflake application databases.
+    # This test ensures that we don't crawl application databases.
+
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+
+    crawled_databases = response.json()[first_connection_id]["databases"].keys()
+    assert len(crawled_databases) > 0
+    assert "SNOWFLAKE" not in crawled_databases
+
+
+def test_get_schemas_with_no_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_schemas_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- DELETE SCHEMA ---
+
+
+def test_delete_schema_with_auth(jp_base_url: str, first_connection_id: str) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/schemas/{first_connection_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    # Ensure the schema was deleted
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    assert len(response.json()) == 0
+
+
+def test_delete_schema_with_no_auth(jp_base_url: str, first_connection_id: str) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/schemas/{first_connection_id}",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_delete_schema_with_incorrect_auth(
+    jp_base_url: str, first_connection_id: str
+) -> None:
+    response = requests.delete(
+        jp_base_url + f"/mito-ai/db/schemas/{first_connection_id}",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/db/sqlite_test.py

@@ -0,0 +1,31 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import requests
+from mito_ai.tests.db.test_db_constants import (
+    SQLITE_TEST_DB_PATH,
+    SQLITE_CONNECTION_DETAILS,
+)
+from mito_ai.tests.conftest import TOKEN
+
+
+def test_add_sqlite_connection(jp_base_url: str) -> None:
+    # Check for the sqlite database
+    assert os.path.exists(SQLITE_TEST_DB_PATH)
+
+    response = requests.post(
+        jp_base_url + "/mito-ai/db/connections",
+        headers={"Authorization": f"token {TOKEN}"},
+        json=SQLITE_CONNECTION_DETAILS,
+    )
+    assert response.status_code == 200
+
+    # Get the schemas
+    response = requests.get(
+        jp_base_url + f"/mito-ai/db/schemas",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+    # Esnure that there is one scema dict in the response
+    assert len(response.json()) == 1

mito_ai/tests/db/test_db_constants.py

@@ -0,0 +1,61 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from typing import Final
+from mito_ai.utils.schema import MITO_FOLDER
+
+DB_DIR_PATH: Final[str] = os.path.join(MITO_FOLDER, "db")
+CONNECTIONS_PATH: Final[str] = os.path.join(DB_DIR_PATH, "connections.json")
+SCHEMAS_PATH: Final[str] = os.path.join(DB_DIR_PATH, "schemas.json")
+BACKUP_DB_DIR_PATH: Final[str] = os.path.join(MITO_FOLDER, "db_backup")
+MSSQL_CONNECTION_DETAILS = {
+    "type": "mssql",
+    "username": os.environ.get("MSSQL_USERNAME"),
+    "password": os.environ.get("MSSQL_PASSWORD"),
+    "host": os.environ.get("MSSQL_HOST"),
+    "port": "1433",
+    "database": "Northwind",
+    "odbc_driver_version": "18",
+}
+MYSQL_CONNECTION_DETAILS = {
+    "type": "mysql",
+    "username": os.environ.get("MYSQL_USERNAME"),
+    "password": os.environ.get("MYSQL_PASSWORD"),
+    "host": os.environ.get("MYSQL_HOST"),
+    "port": "3306",
+    "database": "Northwind",
+}
+ORACLE_CONNECTION_DETAILS = {
+    "type": "oracle",
+    "username": os.environ.get("ORACLE_USERNAME"),
+    "password": os.environ.get("ORACLE_PASSWORD"),
+    "host": os.environ.get("ORACLE_HOST"),
+    "port": "1521",
+    "service_name": "ORCL",
+}
+POSTGRES_CONNECTION_DETAILS = {
+    "type": "postgres",
+    "username": os.environ.get("POSTGRES_USERNAME"),
+    "password": os.environ.get("POSTGRES_PASSWORD"),
+    "host": os.environ.get("POSTGRES_HOST"),
+    "port": "5432",
+    "database": "postgres",
+}
+SNOWFLAKE = {
+    "type": "snowflake",
+    "username": os.environ.get("SNOWFLAKE_USERNAME"),
+    "password": os.environ.get("SNOWFLAKE_PASSWORD"),
+    "account": os.environ.get("SNOWFLAKE_ACCOUNT"),
+    "warehouse": "COMPUTE_WH",
+}
+SQLITE_TEST_DB_PATH: Final[str] = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+    "tests",
+    "data",
+    "stock_data.sqlite3",
+)
+SQLITE_CONNECTION_DETAILS = {
+    "type": "sqlite",
+    "database": SQLITE_TEST_DB_PATH,
+}

mito_ai/tests/deploy_app/test_app_deploy_utils.py

@@ -0,0 +1,89 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import zipfile
+import logging
+from mito_ai.app_deploy.app_deploy_utils import add_files_to_zip
+from mito_ai.path_utils import AbsoluteNotebookDirPath
+
+class TestAddFilesToZip:
+    """Test cases for add_files_to_zip helper function"""
+
+    def test_files_added_correctly(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "file1.txt"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file1.txt", "file2.txt"], 'test-app-file-name.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file1.txt" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
+            
+    def test_renames_app_file(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "original-file-name.py"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["original-file-name.py", "file2.txt"], 'original-file-name.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "app.py" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
+
+    def test_directories_added_recursively(self, tmp_path):
+        """Ensure directories are added recursively with correct relative paths"""
+        nested = tmp_path / "folder"
+        nested.mkdir()
+        (nested / "nested1.txt").write_text("nested1 content")
+        subfolder = nested / "sub"
+        subfolder.mkdir()
+        (subfolder / "nested2.txt").write_text("nested2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["folder"], 'test-app.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "folder/nested1.txt" in names
+            assert "folder/sub/nested2.txt" in names
+
+    def test_missing_files_skipped(self, tmp_path, caplog):
+        """Ensure missing files do not break the function and warning is logged"""
+        caplog.set_level(logging.WARNING)
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["does_not_exist.txt"], 'test-app.py', logger=logging.getLogger())
+
+        # Zip should exist but be empty
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            assert zf.namelist() == []
+
+        # Check warning was logged
+        assert any("Skipping missing file" in record.message for record in caplog.records)
+
+    def test_arcname_paths_correct(self, tmp_path):
+        """Ensure arcname paths inside zip preserve relative paths to base_path"""
+        (tmp_path / "file.txt").write_text("content")
+        folder = tmp_path / "folder"
+        folder.mkdir()
+        (folder / "nested.txt").write_text("nested content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file.txt", "folder"], 'test-app.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file.txt" in names
+            assert "folder/nested.txt" in names

mito_ai/tests/file_uploads/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.