Add more tests

superset/commands/importers/exceptions.py

@@ -26,3 +26,8 @@ class IncorrectVersionError(CommandException):
 class NoValidFilesFoundError(CommandException):
     status = 400
     message = "No valid import files were found"
+
+
+class IncorrectFormatError(CommandException):
+    status = 422
+    message = "File has the incorrect format"

superset/importexport/api.py

@@ -17,13 +17,16 @@
 import json
 from datetime import datetime
 from io import BytesIO
-from zipfile import ZipFile
+from zipfile import is_zipfile, ZipFile
 
 from flask import request, Response, send_file
 from flask_appbuilder.api import BaseApi, expose, protect
 
 from superset.commands.export.assets import ExportAssetsCommand
-from superset.commands.importers.exceptions import NoValidFilesFoundError
+from superset.commands.importers.exceptions import (
+    IncorrectFormatError,
+    NoValidFilesFoundError,
+)
 from superset.commands.importers.v1.assets import ImportAssetsCommand
 from superset.commands.importers.v1.utils import get_contents_from_bundle
 from superset.extensions import event_logger
@@ -140,6 +143,8 @@ def import_(self) -> Response:
         upload = request.files.get("bundle")
         if not upload:
             return self.response_400()
+        if not is_zipfile(upload):
+            raise IncorrectFormatError("Not a ZIP file")
 
         with ZipFile(upload) as bundle:
             contents = get_contents_from_bundle(bundle)

tests/unit_tests/conftest.py

@@ -14,8 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=redefined-outer-name
+# pylint: disable=redefined-outer-name, import-outside-toplevel
 
+import importlib
 from typing import Any, Iterator
 
 import pytest
@@ -70,6 +71,12 @@ def app() -> Iterator[SupersetApp]:
     app_initializer = SupersetAppInitializer(app)
     app_initializer.init_app()
 
+    # reload base views to ensure error handlers are applied to the app
+    with app.app_context():
+        import superset.views.base
+
+        importlib.reload(superset.views.base)
+
     yield app
 
 

tests/unit_tests/importexport/api_test.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name, import-outside-toplevel
 
 import json
 from io import BytesIO
@@ -30,7 +31,6 @@ def test_export_assets(mocker: MockFixture, client: Any) -> None:
     """
     Test exporting assets.
     """
-    # pylint: disable=import-outside-toplevel
     from superset.commands.importers.v1.utils import get_contents_from_bundle
 
     # grant access
@@ -47,7 +47,6 @@ def test_export_assets(mocker: MockFixture, client: Any) -> None:
         ("databases/example.yaml", "<DATABASE CONTENTS>"),
     ]
 
-    # pylint: disable=invalid-name
     ExportAssetsCommand = mocker.patch("superset.importexport.api.ExportAssetsCommand")
     ExportAssetsCommand().run.return_value = mocked_contents[:]
 
@@ -80,7 +79,6 @@ def test_import_assets(mocker: MockFixture, client: Any) -> None:
         "databases/example.yaml": "<DATABASE CONTENTS>",
     }
 
-    # pylint: disable=invalid-name
     ImportAssetsCommand = mocker.patch("superset.importexport.api.ImportAssetsCommand")
 
     root = Path("assets_export")
@@ -105,3 +103,152 @@ def test_import_assets(mocker: MockFixture, client: Any) -> None:
 
     passwords = {"assets_export/databases/imported_database.yaml": "SECRET"}
     ImportAssetsCommand.assert_called_with(mocked_contents, passwords=passwords)
+
+
+def test_import_assets_not_zip(mocker: MockFixture, client: Any) -> None:
+    """
+    Test error message when the upload is not a ZIP file.
+    """
+    # grant access
+    mocker.patch(
+        "flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
+    )
+    mocker.patch.object(security_manager, "has_access", return_value=True)
+
+    buf = BytesIO(b"definitely_not_a_zip_file")
+    form_data = {
+        "bundle": (buf, "broken.txt"),
+    }
+    response = client.post(
+        "/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
+    )
+    assert response.status_code == 422
+    assert response.json == {
+        "errors": [
+            {
+                "message": "Not a ZIP file",
+                "error_type": "GENERIC_COMMAND_ERROR",
+                "level": "warning",
+                "extra": {
+                    "issue_codes": [
+                        {
+                            "code": 1010,
+                            "message": (
+                                "Issue 1010 - Superset encountered an error while "
+                                "running a command."
+                            ),
+                        }
+                    ]
+                },
+            }
+        ]
+    }
+
+
+def test_import_assets_no_form_data(mocker: MockFixture, client: Any) -> None:
+    """
+    Test error message when the upload has no form data.
+    """
+    # grant access
+    mocker.patch(
+        "flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
+    )
+    mocker.patch.object(security_manager, "has_access", return_value=True)
+
+    response = client.post("/api/v1/assets/import/", data="some_content")
+    assert response.status_code == 400
+    assert response.json == {
+        "errors": [
+            {
+                "message": "Request MIME type is not 'multipart/form-data'",
+                "error_type": "INVALID_PAYLOAD_FORMAT_ERROR",
+                "level": "error",
+                "extra": {
+                    "issue_codes": [
+                        {
+                            "code": 1019,
+                            "message": (
+                                "Issue 1019 - The submitted payload has the incorrect "
+                                "format."
+                            ),
+                        }
+                    ]
+                },
+            }
+        ]
+    }
+
+
+def test_import_assets_incorrect_form_data(mocker: MockFixture, client: Any) -> None:
+    """
+    Test error message when the upload form data has the wrong key.
+    """
+    # grant access
+    mocker.patch(
+        "flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
+    )
+    mocker.patch.object(security_manager, "has_access", return_value=True)
+
+    buf = BytesIO(b"definitely_not_a_zip_file")
+    form_data = {
+        "wrong": (buf, "broken.txt"),
+    }
+    response = client.post(
+        "/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
+    )
+    assert response.status_code == 400
+    assert response.json == {"message": "Arguments are not correct"}
+
+
+def test_import_assets_no_contents(mocker: MockFixture, client: Any) -> None:
+    """
+    Test error message when the ZIP bundle has no contents.
+    """
+    # grant access
+    mocker.patch(
+        "flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
+    )
+    mocker.patch.object(security_manager, "has_access", return_value=True)
+
+    mocked_contents = {
+        "README.txt": "Something is wrong",
+    }
+
+    root = Path("assets_export")
+    buf = BytesIO()
+    with ZipFile(buf, "w") as bundle:
+        for path, contents in mocked_contents.items():
+            with bundle.open(str(root / path), "w") as fp:
+                fp.write(contents.encode())
+    buf.seek(0)
+
+    form_data = {
+        "bundle": (buf, "assets_export.zip"),
+        "passwords": json.dumps(
+            {"assets_export/databases/imported_database.yaml": "SECRET"}
+        ),
+    }
+    response = client.post(
+        "/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
+    )
+    assert response.status_code == 400
+    assert response.json == {
+        "errors": [
+            {
+                "message": "No valid import files were found",
+                "error_type": "GENERIC_COMMAND_ERROR",
+                "level": "warning",
+                "extra": {
+                    "issue_codes": [
+                        {
+                            "code": 1010,
+                            "message": (
+                                "Issue 1010 - Superset encountered an error while "
+                                "running a command."
+                            ),
+                        }
+                    ]
+                },
+            }
+        ]
+    }