Repository: opus-42/superset-api-client (language: python)

Issue:
Add integration test suite
Since the sole purpose of this client is to interface with the Superset API, testing should be focused on running against an actual Superset API instead of mocking Superset API endpoints.

An integration test suite can build on top of the docker-compose config already present.


Discussion / hints:

Fix patch:
diff --git a/.env b/.env
index 8d73393..23ff34b 100644
--- a/.env
+++ b/.env
@@ -1,2 +1,2 @@
-SUPERSET_VERSION=1.3.1
+SUPERSET_VERSION=2.0.0
 LOAD_EXAMPLES=0
diff --git a/.github/workflows/python-package.yaml b/.github/workflows/python-package.yaml
index 52dd2d3..ad973ac 100644
--- a/.github/workflows/python-package.yaml
+++ b/.github/workflows/python-package.yaml
@@ -18,12 +18,14 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install flake8 pytest
           pip install .[dev]
       - name: Lint with flake8
         run: |
           # fully comply to pep8
           flake8 . --count --max-complexity=10 --max-line-length=127 --statistics
-      - name: Test with pytest
+      - name: Test with pytest and coverage
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          pytest tests
+          pytest -cov=supersetapiclient tests
+          coveralls --service=github
diff --git a/docker-compose.yml b/docker-compose.yml
index 15cf35e..77b6cea 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -10,7 +10,14 @@ services:
       - "127.0.0.1:8080:8080"
     environment:
       - LOAD_EXAMPLES="${LOAD_EXAMPLES}"
-      # - SUPERSET_CONFIG_PATH=/app/superset_home/test_config.py
+      - SUPERSET_CONFIG_PATH=/app/superset_home/test_config.py
     volumes:
       - ./scripts/entrypoint.sh:/app/superset_home/entrypoint.sh
-      # - ./tests/superset_test_config.py:/app/superset_home/test_config.py
+      - ./tests/superset_config.py:/app/superset_home/test_config.py
+  postgres:
+    image: postgres:14.1-alpine
+    hostname: pg
+    restart: unless-stopped
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh
index afd256b..de00340 100644
--- a/scripts/entrypoint.sh
+++ b/scripts/entrypoint.sh
@@ -3,7 +3,7 @@
 superset db upgrade
 superset init
 superset fab create-admin --username admin --firstname admin --lastname admin --email admin --password admin
-if [ "${LOAD_EXAMPLES}" -eq 1 ]; then
+if [ "${LOAD_EXAMPLES}" -eq "1" ]; then
   superset load_examples
 fi
-superset run -h 0.0.0.0 -p 8080
+FLASK_ENV=development superset run -h 0.0.0.0 -p 8080
diff --git a/setup.cfg b/setup.cfg
index da77217..650b8ad 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -20,5 +20,5 @@ install_requires =
 python_requires = >=3.6
 
 [options.extras_require]
-dev = flake8; pytest; requests-mock; pytest-docker
+dev = flake8; pytest; requests-mock; pytest-docker; pytest-cov; coveralls
 build = build
diff --git a/supersetapiclient/base.py b/supersetapiclient/base.py
index 13ba3a7..315b5e3 100644
--- a/supersetapiclient/base.py
+++ b/supersetapiclient/base.py
@@ -3,7 +3,7 @@
 import dataclasses
 try:
     from functools import cached_property
-except ImportError:
+except ImportError:  # pragma: no cover
     # Python<3.8
     from cached_property import cached_property
 import json
@@ -12,9 +12,9 @@
 
 import os.path
 import yaml
-from requests import Response, HTTPError
+from requests import HTTPError
 
-from supersetapiclient.exceptions import ComplexBadRequestError, NotFound, BadRequestError
+from supersetapiclient.exceptions import ComplexBadRequestError, NotFound, BadRequestError, MultipleFound
 
 logger = logging.getLogger(__name__)
 
@@ -45,7 +45,6 @@ def raise_for_status(response):
 
 class Object:
     _parent = None
-    EXPORTABLE = False
     JSON_FIELDS = []
 
     @classmethod
@@ -56,10 +55,7 @@ def fields(cls):
     @classmethod
     def field_names(cls):
         """Get field names."""
-        return set(
-            f.name
-            for f in dataclasses.fields(cls)
-        )
+        return {f.name for f in cls.fields()}
 
     @classmethod
     def from_json(cls, json: dict):
@@ -74,6 +70,18 @@ def from_json(cls, json: dict):
         field_names = cls.field_names()
         return cls(**{k: v for k, v in json.items() if k in field_names})
 
+    def to_json(self, columns):
+        o = {}
+        for c in columns:
+            if not hasattr(self, c):
+                # Column that is not implemented yet
+                continue
+            value = getattr(self, c)
+            if c in self.JSON_FIELDS:
+                value = json.dumps(value)
+            o[c] = value
+        return o
+
     def __post_init__(self):
         for f in self.JSON_FIELDS:
             setattr(self, f, json.loads(getattr(self, f) or "{}"))
@@ -85,46 +93,9 @@ def base_url(self) -> str:
             str(self.id)
         )
 
-    @property
-    def import_url(self) -> str:
-        return self._parent.client.join_urls(
-            self._parent.base_url,
-            str(self.id)
-        )
-
-    @property
-    def export_url(self) -> str:
-        """Export url for a single object."""
-        # Note that params should be passed on run
-        # to bind to a specific object
-        return self._parent.client.join_urls(
-            self.base_url,
-            "export"
-        )
-
-    @property
-    def test_connection_url(self) -> str:
-        return self._parent.client.join_urls(
-            self._parent.base_url,
-            str(self.id)
-        )
-
     def export(self, path: Union[Path, str]) -> None:
         """Export object to path"""
-        if not self.EXPORTABLE:
-            raise NotImplementedError(
-                "Export is not defined for this object."
-            )
-
-        # Get export response
-        client = self._parent.client
-        response = client.get(self.export_url, params={
-            "q": [self.id]  # Object must have an id field to be exported
-        })
-        raise_for_status(response)
-
-        with open(path, "w", encoding="utf-8") as f:
-            f.write(response.text)
+        self._parent.export(ids=[self.id], path=path)
 
     def fetch(self) -> None:
         """Fetch additional object information."""
@@ -136,24 +107,20 @@ def fetch(self) -> None:
         o = o.get("result")
         for k, v in o.items():
             if k in field_names:
-                setattr(self, k, v)
+                if k in self.JSON_FIELDS:
+                    setattr(self, k, json.loads(v or "{}"))
+                else:
+                    setattr(self, k, v)
 
     def save(self) -> None:
         """Save object information."""
-        o = {}
-        for c in self._parent.edit_columns:
-            if hasattr(self, c):
-                value = getattr(self, c)
-
-                if c in self.JSON_FIELDS:
-                    value = json.dumps(value)
-                o[c] = value
-
+        o = self.to_json(columns=self._parent.edit_columns)
         response = self._parent.client.put(self.base_url, json=o)
-        if response.status_code in [400, 422]:
-            logger.error(response.text)
         raise_for_status(response)
 
+    def delete(self) -> bool:
+        return self._parent.delete(id=self.id)
+
 
 class ObjectFactories:
     endpoint = ""
@@ -186,10 +153,7 @@ def _infos(self):
                 "q": json.dumps(self._INFO_QUERY)
             })
 
-        if response.status_code != 200:
-            logger.error(f"Unable to build object factory for {self.endpoint}")
-            raise_for_status(response)
-
+        raise_for_status(response)
         return response.json()
 
     @property
@@ -232,27 +196,6 @@ def export_url(self):
             "export/"
         )
 
-    @property
-    def test_connection_url(self):
-        """Base url for these objects."""
-        return self.client.join_urls(
-            self.client.base_url,
-            self.endpoint,
-            "test_connection"
-        )
-
-    @staticmethod
-    def _handle_reponse_status(response: Response) -> None:
-        """Handle response status."""
-        if response.status_code not in (200, 201):
-            logger.error(
-                f"Unable to proceed, API return {response.status_code}"
-            )
-            logger.error(f"Full API response is {response.text}")
-
-        # Finally raising for status
-        raise_for_status(response)
-
     def get(self, id: int):
         """Get an object by id."""
         url = self.base_url + str(id)
@@ -307,39 +250,28 @@ def find(self, page_size: int = 100, page: int = 0, **kwargs):
 
     def count(self):
         """Count objects."""
-
-        # To do : add kwargs parameters for more flexibility
         response = self.client.get(self.base_url)
-
-        if response.status_code not in (200, 201):
-            logger.error(response.text)
         raise_for_status(response)
-
-        json = response.json()
-        return (json['count'])
+        return response.json()['count']
 
     def find_one(self, **kwargs):
         """Find only object or raise an Exception."""
         objects = self.find(**kwargs)
         if len(objects) == 0:
-            raise NotFound(f"No {self.base_object.__name__} has been found.")
+            raise NotFound(f"No {self.base_object.__name__} found")
+        if len(objects) > 1:
+            raise MultipleFound(f"Multiple {self.base_object.__name__} found")
         return objects[0]
 
     def add(self, obj) -> int:
-        """Create a object on remote."""
-
-        o = {}
-        for c in self.add_columns:
-            if hasattr(obj, c):
-                value = getattr(obj, c)
-
-                if c in obj.JSON_FIELDS:
-                    value = json.dumps(value)
-                o[c] = value
+        """Create an object on remote."""
 
+        o = obj.to_json(columns=self.add_columns)
         response = self.client.post(self.base_url, json=o)
         raise_for_status(response)
-        return response.json().get("id")
+        obj.id = response.json().get("id")
+        obj._parent = self
+        return obj.id
 
     def export(self, ids: List[int], path: Union[Path, str]) -> None:
         """Export object into an importable file"""
@@ -352,39 +284,33 @@ def export(self, ids: List[int], path: Union[Path, str]) -> None:
                 "q": f"[{ids_array}]"
             })
 
-        if response.status_code not in (200, 201):
-            logger.error(response.text)
         raise_for_status(response)
 
         content_type = response.headers["content-type"].strip()
-        if content_type.startswith("application/text"):
-            data = response.text
-            data = yaml.load(data, Loader=yaml.FullLoader)
+        if content_type.startswith("application/text"):  # pragma: no cover
+            # Superset 1.x
+            data = yaml.load(response.text, Loader=yaml.FullLoader)
             with open(path, "w", encoding="utf-8") as f:
                 yaml.dump(data, f, default_flow_style=False)
+            return
+        if content_type.startswith("application/json"):  # pragma: no cover
+            # Superset 1.x
+            data = response.json()
+            with open(path, 'w', encoding='utf-8') as f:
+                json.dump(data, f, ensure_ascii=False, indent=4)
+            return
         if content_type.startswith("application/zip"):
             data = response.content
             with open(path, 'wb') as f:
                 f.write(data)
-        else:
-            data = response.json()
-            with open(path, 'w', encoding='utf-8') as f:
-                json.dump(data, f, ensure_ascii=False, indent=4)
-
-        return data
+            return
+        raise ValueError(f"Unknown content type {content_type}")
 
-    def delete(self, id) -> int:
-        """Delete a object on remote."""
+    def delete(self, id: int) -> bool:
+        """Delete an object on remote."""
         response = self.client.delete(self.base_url + str(id))
-
-        if response.status_code not in (200, 201):
-            logger.error(response.text)
         raise_for_status(response)
-
-        if response.json().get('message') == 'OK':
-            return True
-        else:
-            return False
+        return response.json().get('message') == 'OK'
 
     def import_file(self, file_path, overwrite=False, passwords=None) -> dict:
         """Import a file on remote.
@@ -416,20 +342,4 @@ def import_file(self, file_path, overwrite=False, passwords=None) -> dict:
         raise_for_status(response)
 
         # If import is successful, the following is returned: {'message': 'OK'}
-        return response.json()
-
-    def test_connection(self, obj):
-        """Import a file on remote."""
-        url = self.test_connection_url
-        connection_columns = ['database_name', 'sqlalchemy_uri']
-        o = {}
-        for c in connection_columns:
-            if hasattr(obj, c):
-                value = getattr(obj, c)
-                o[c] = value
-
-        response = self.client.post(url, json=o)
-        if response.json().get('message') == 'OK':
-            return True
-        else:
-            return False
+        return response.json().get('message') == 'OK'
diff --git a/supersetapiclient/charts.py b/supersetapiclient/charts.py
index bead657..e3044a3 100644
--- a/supersetapiclient/charts.py
+++ b/supersetapiclient/charts.py
@@ -1,5 +1,6 @@
 """Charts."""
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Optional, List
 
 from supersetapiclient.base import (
     Object, ObjectFactories, json_field, default_string
@@ -12,14 +13,19 @@ class Chart(Object):
         "params"
     ]
 
-    id: int
-    description: default_string()
-    slice_name: default_string()
-    params: json_field()
-    datasource_id: int = None
-    datasource_type: str = default_string
-    viz_type: str = ""
+    id: Optional[int] = None
+    description: str = default_string()
+    slice_name: str = default_string()
     params: dict = json_field()
+    datasource_id: Optional[int] = None
+    datasource_type: str = default_string()
+    viz_type: str = ""
+    dashboards: List[int] = field(default_factory=list)
+
+    def to_json(self, columns):
+        o = super().to_json(columns)
+        o["dashboards"] = self.dashboards
+        return o
 
 
 class Charts(ObjectFactories):
diff --git a/supersetapiclient/client.py b/supersetapiclient/client.py
index 44d1418..b5fef7f 100644
--- a/supersetapiclient/client.py
+++ b/supersetapiclient/client.py
@@ -3,7 +3,7 @@
 import logging
 try:
     from functools import cached_property
-except ImportError:
+except ImportError:  # pragma: no cover
     # Python<3.8
     from cached_property import cached_property
 
@@ -11,6 +11,7 @@
 import requests.exceptions
 import requests_oauthlib
 
+from supersetapiclient.base import raise_for_status
 from supersetapiclient.dashboards import Dashboards
 from supersetapiclient.charts import Charts
 from supersetapiclient.datasets import Datasets
@@ -120,13 +121,12 @@ def authenticate(self) -> dict:
             "provider": self.provider,
             "refresh": "true"
         })
-        response.raise_for_status()
+        raise_for_status(response)
         return response.json()
 
     def token_refresher(self, r, *args, **kwargs):
         """A requests response hook for token refresh."""
         if r.status_code == 401:
-
             # Check if token has expired
             try:
                 msg = r.json().get("msg")
@@ -141,7 +141,7 @@ def token_refresher(self, r, *args, **kwargs):
             refresh_r = requests_oauthlib.OAuth2Session(
                 token=tmp_token
             ).post(self.refresh_endpoint)
-            refresh_r.raise_for_status()
+            raise_for_status(refresh_r)
 
             new_token = refresh_r.json()
             if "refresh_token" not in new_token:
@@ -175,7 +175,7 @@ def run(self, database_id, query, query_limit=None):
         if query_limit:
             payload["queryLimit"] = query_limit
         response = self.post(self._sql_endpoint, json=payload)
-        response.raise_for_status()
+        raise_for_status(response)
         result = response.json()
         display_limit = result.get("displayLimit", None)
         display_limit_reached = result.get("displayLimitReached", False)
@@ -211,7 +211,7 @@ def csrf_token(self, session) -> str:
             self.join_urls(self.base_url, "/security/csrf_token/"),
             headers={"Referer": f"{self.base_url}"},
         )
-        csrf_response.raise_for_status()  # Check CSRF Token went well
+        raise_for_status(csrf_response)  # Check CSRF Token went well
         return csrf_response.json().get("result")
 
 
diff --git a/supersetapiclient/dashboards.py b/supersetapiclient/dashboards.py
index c12af27..48480ec 100644
--- a/supersetapiclient/dashboards.py
+++ b/supersetapiclient/dashboards.py
@@ -1,7 +1,6 @@
 """Dashboards."""
-from typing import List
-
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import List, Optional
 
 from supersetapiclient.base import (
     Object, ObjectFactories, json_field, default_string
@@ -14,21 +13,19 @@ class Dashboard(Object):
         "json_metadata",
         "position_json"
     ]
-    EXPORTABLE = True
 
-    id: int
     dashboard_title: str
     published: bool
+    id: Optional[int] = None
     json_metadata: dict = json_field()
     position_json: dict = json_field()
     changed_by: str = default_string()
-    changed_by: str = default_string()
     slug: str = default_string()
     changed_by_name: str = default_string()
     changed_by_url: str = default_string()
     css: str = default_string()
     changed_on: str = default_string()
-    charts: list = default_string()
+    charts: List[str] = field(default_factory=list)
 
     @property
     def colors(self) -> dict:
@@ -47,14 +44,11 @@ def update_colors(self, value: dict) -> None:
         self.colors = colors
 
     def get_charts(self) -> List[int]:
-        """Get dashboard Slice IDs."""
+        """Get chart objects"""
         charts = []
-        for _, v in self.position_json.items():
-            if isinstance(v, dict):
-                if v.get("type") == "CHART":
-                    meta = v.get("meta", {})
-                    if meta.get("chartId") is not None:
-                        charts.append(meta.get("chartId"))
+        for slice_name in self.charts:
+            c = self._parent.client.charts.find_one(slice_name=slice_name)
+            charts.append(c)
         return charts
 
 
diff --git a/supersetapiclient/databases.py b/supersetapiclient/databases.py
index 198d609..d2cdecd 100644
--- a/supersetapiclient/databases.py
+++ b/supersetapiclient/databases.py
@@ -1,36 +1,70 @@
 """Databases."""
 from dataclasses import dataclass
+from typing import Optional
 
 from supersetapiclient.base import (
-    Object, ObjectFactories, json_field, default_string
+    Object, ObjectFactories, default_string, json_field
 )
 
 
 @dataclass
 class Database(Object):
-    JSON_FIELDS = ["metadata_cache_timeout"]
+    JSON_FIELDS = [
+        "extra",
+    ]
 
-    id: int
     database_name: str
-    allow_ctas: bool
-    allow_cvas: bool
-    allow_dml: bool
-    allow_multi_schema_metadata_fetch: bool
-    allow_run_async: bool
-    metadata_cache_timeout: dict = json_field()
+    id: Optional[int] = None
+    allow_ctas: bool = True
+    allow_cvas: bool = True
+    allow_dml: bool = True
+    allow_multi_schema_metadata_fetch: bool = True
+    allow_run_async: bool = True
+    expose_in_sqllab: bool = True
+    cache_timeout: Optional[int] = None
     encrypted_extra: str = default_string()
     engine: str = default_string()
-    extra: str = json_field()
+    extra: dict = json_field()
     force_ctas_schema: str = default_string()
     server_cert: str = default_string()
     sqlalchemy_uri: str = default_string()
 
+    def to_json(self, *args, **kwargs):
+        if not self.extra:
+            self.extra = {
+                "metadata_params": {},
+                "engine_params": {},
+                "metadata_cache_timeout": {},
+                "schemas_allowed_for_csv_upload": [],
+            }
+        return super().to_json(*args, **kwargs)
+
     def run(self, query, query_limit=None):
         return self._parent.client.run(
             database_id=self.id, query=query, query_limit=query_limit
         )
 
+    def test_connection(self):
+        return self._parent.test_connection(self)
+
 
 class Databases(ObjectFactories):
     endpoint = "/database/"
     base_object = Database
+
+    @property
+    def test_connection_url(self):
+        """Base url for these objects."""
+        return self.client.join_urls(
+            self.client.base_url,
+            self.endpoint,
+            "test_connection"
+        )
+
+    def test_connection(self, obj):
+        """Test connection to a database"""
+        url = self.test_connection_url
+        connection_columns = ['database_name', 'sqlalchemy_uri']
+        o = {c: getattr(obj, c) for c in connection_columns}
+        response = self.client.post(url, json=o)
+        return response.json().get('message') == 'OK'
diff --git a/supersetapiclient/datasets.py b/supersetapiclient/datasets.py
index 50cb8fd..6daf468 100644
--- a/supersetapiclient/datasets.py
+++ b/supersetapiclient/datasets.py
@@ -1,5 +1,6 @@
 """Dashboards."""
 from dataclasses import dataclass, field
+from typing import Optional
 
 from supersetapiclient.base import (
     Object, ObjectFactories
@@ -9,15 +10,14 @@
 @dataclass
 class Dataset(Object):
     JSON_FIELDS = []
-    EXPORTABLE = True
 
-    id: int
-    table_name: str
+    id: Optional[int] = None
+    table_name: str = ""
     schema: str = ""
     columns: list = field(default_factory=list)
     description: str = ""
     kind: str = ""
-    database_id: int = None
+    database_id: Optional[int] = None
     sql: str = ""
 
     @classmethod
@@ -28,6 +28,15 @@ def from_json(cls, json: dict):
             res.database_id = database.get("id")
         return res
 
+    def to_json(self, *args, **kwargs):
+        o = super().to_json(*args, **kwargs)
+        o.pop("columns", None)
+        if self.id:
+            o["database_id"] = self.database_id
+        else:
+            o["database"] = self.database_id
+        return o
+
     def run(self, query_limit=None):
         if not self.sql:
             raise ValueError("Cannot run a dataset with no SQL")
diff --git a/supersetapiclient/exceptions.py b/supersetapiclient/exceptions.py
index 7b46659..fe5b0e2 100644
--- a/supersetapiclient/exceptions.py
+++ b/supersetapiclient/exceptions.py
@@ -8,6 +8,10 @@ class NotFound(Exception):
     pass
 
 
+class MultipleFound(Exception):
+    pass
+
+
 class QueryLimitReached(Exception):
     pass
 
@@ -18,7 +22,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def __str__(self):
-        return self.message
+        return json.dumps(self.message, indent=4)
 
 
 class ComplexBadRequestError(HTTPError):
diff --git a/supersetapiclient/saved_queries.py b/supersetapiclient/saved_queries.py
index 0c378d1..f3e82cb 100644
--- a/supersetapiclient/saved_queries.py
+++ b/supersetapiclient/saved_queries.py
@@ -1,5 +1,6 @@
 """Saved queries."""
 from dataclasses import dataclass
+from typing import Optional
 
 from supersetapiclient.base import (
     Object, ObjectFactories, default_string
@@ -10,8 +11,8 @@
 class SavedQuery(Object):
     JSON_FIELDS = []
 
-    id: int
     label: str
+    id: Optional[int] = None
     description: str = default_string()
     sql: str = default_string()
     db_id: int = None


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index f7a11bf..c784fe5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -17,6 +17,12 @@
 API_MOCKS = Path(__file__).parent / "mocks" / "endpoints"
 
 
+class CustomClient(SupersetClient):
+    @property
+    def _sql_endpoint(self) -> str:
+        return self.join_urls(self.base_url, "execute_sql_json/")
+
+
 @pytest.fixture
 def permanent_requests(requests_mock): # noqa
 
@@ -63,7 +69,7 @@ def is_responsive(url):
 
 
 @pytest.fixture(scope="session")
-def superset_api(docker_ip, docker_services):
+def superset_url(docker_ip, docker_services):
     """Ensure that Superset API is up and responsive."""
 
     # `port_for` takes a container port and returns the corresponding host port
@@ -75,9 +81,40 @@ def superset_api(docker_ip, docker_services):
     docker_services.wait_until_responsive(
         timeout=600, pause=1.0, check=lambda: is_responsive(url)
     )
-    yield SupersetClient(url, "admin", "admin")
+    yield url
+
+
+@pytest.fixture
+def superset_api(superset_url):
+    yield CustomClient(superset_url, "admin", "admin")
 
 
 @pytest.fixture(scope="session")
 def docker_compose_file(pytestconfig):
     return str(Path(__file__).parent.parent / "docker-compose.yml")
+
+
+# Fixtures to enable testing without starting and stopping containers on each run.
+# See https://github.com/avast/pytest-docker/issues/46#issuecomment-887408396
+def pytest_addoption(parser):
+    """Add the --keepalive option for pytest."""
+    parser.addoption("--keepalive", "-K", action="store_true", default=False, help="Keep Docker containers alive")
+
+
+@pytest.fixture(scope="session")
+def keepalive(request):
+    return request.config.option.keepalive
+
+
+@pytest.fixture(scope="session")
+def docker_compose_project_name(keepalive, docker_compose_project_name):
+    if keepalive:
+        return "test_superset_app"
+    return docker_compose_project_name
+
+
+@pytest.fixture(scope="session")
+def docker_cleanup(keepalive, docker_cleanup):
+    if keepalive:
+        return "version"
+    return docker_cleanup
diff --git a/tests/integration/test_basic.py b/tests/integration/test_basic.py
index 997548a..df8e0de 100644
--- a/tests/integration/test_basic.py
+++ b/tests/integration/test_basic.py
@@ -1,2 +1,529 @@
-def test_connection(superset_api):
-    assert superset_api.databases.find() == []
+import json
+import random
+import string
+import tempfile
+
+import pytest
+import requests
+import requests.exceptions
+
+from supersetapiclient.charts import Chart
+from supersetapiclient.client import SupersetClient, raise_for_status
+from supersetapiclient.dashboards import Dashboard
+from supersetapiclient.databases import Database
+from supersetapiclient.datasets import Dataset
+from supersetapiclient.exceptions import BadRequestError, ComplexBadRequestError, QueryLimitReached, NotFound, \
+    MultipleFound
+from supersetapiclient.saved_queries import SavedQuery
+
+
+def random_str(length, lowercase=False):
+    # Create a random string
+    chars = string.ascii_lowercase if lowercase else string.ascii_letters
+    return "".join(random.choice(chars) for _ in range(length))
+
+
+@pytest.fixture
+def database(superset_api):
+    db = Database(
+        database_name=random_str(8),
+        sqlalchemy_uri="postgresql+psycopg2://postgres:postgres@pg:5432/postgres",
+    )
+    superset_api.databases.add(db)
+    yield db
+    try:
+        db.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+@pytest.fixture
+def database_with_table(database):
+    schema, table_name = random_str(8, lowercase=True), random_str(8, lowercase=True)
+    database.run(f"CREATE SCHEMA IF NOT EXISTS {schema}")
+    database.run(f"CREATE TABLE IF NOT EXISTS {schema}.{table_name} (i integer)")
+    yield database, schema, table_name
+    database.run(f"DROP TABLE IF EXISTS {schema}.{table_name}")
+    database.run(f"DROP SCHEMA IF EXISTS {schema}")
+
+
+@pytest.fixture
+def dataset(superset_api, database_with_table):
+    database, schema, table_name = database_with_table
+    ds = Dataset(
+        database_id=database.id,
+        schema=schema,
+        table_name=table_name,
+        description="My Test Table",
+    )
+    superset_api.datasets.add(ds)
+    yield ds
+    try:
+        ds.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+@pytest.fixture
+def virtual_dataset(superset_api, database):
+    # We're using 'dataset' and 'virtual_dataset' in the same test case so need to create a new table here because
+    # two databases can't cover the same table.
+    schema, table_name = random_str(8, lowercase=True), random_str(8, lowercase=True)
+    database.run(f"CREATE SCHEMA IF NOT EXISTS {schema}")
+    database.run(f"CREATE TABLE IF NOT EXISTS {schema}.{table_name} (i integer)")
+    ds = Dataset(
+        database_id=database.id,
+        schema=schema,
+        table_name=table_name,
+        description="My Virtual Table",
+        sql=f"SELECT i FROM {schema}.{table_name} ORDER BY i LIMIT 3",
+    )
+    superset_api.datasets.add(ds)
+    yield ds
+    database.run(f"DROP TABLE IF EXISTS {schema}.{table_name}")
+    database.run(f"DROP SCHEMA IF EXISTS {schema}")
+    try:
+        ds.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+@pytest.fixture
+def saved_query(superset_api, database_with_table):
+    database, schema, table_name = database_with_table
+    sq = SavedQuery(
+        db_id=database.id,
+        label=random_str(8),
+        schema=schema,
+        sql=f"SELECT i FROM {schema}.{table_name}",
+    )
+    superset_api.saved_queries.add(sq)
+    yield sq
+    try:
+        sq.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+@pytest.fixture
+def chart(superset_api, dataset):
+    c = Chart(
+        datasource_id=dataset.id,
+        datasource_type="table",
+        slice_name=random_str(8),
+        viz_type="table",
+    )
+    superset_api.charts.add(c)
+    yield c
+    try:
+        c.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+@pytest.fixture
+def dashboard(superset_api, dataset):
+    d = Dashboard(
+        dashboard_title=random_str(8),
+        published=True,
+        slug=random_str(8),
+    )
+    superset_api.dashboards.add(d)
+    yield d
+    try:
+        d.delete()
+    except BadRequestError as e:
+        if not (e.response.status_code == 404 and e.message == "Not found"):
+            raise
+
+
+class TestEntities:
+    def test_databases(self, superset_api, database):
+        # Test getting an invalid item ID
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.databases.get(id=0)
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        assert isinstance(database.id, int)
+        assert superset_api.databases.get(id=database.id).database_name == database.database_name
+
+        # Test fetch
+        database.fetch()
+
+        # Test count
+        assert superset_api.databases.count() >= 1
+
+        # Test creating a conflicting item
+        db_id, database.id = database.id, None
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.databases.add(database)
+        database.id = db_id
+        assert exc_info.value.response.status_code == 422
+        assert exc_info.value.message == {"database_name": "A database with the same name already exists."}
+        assert str(exc_info.value) == """\
+{
+    "database_name": "A database with the same name already exists."
+}"""
+
+        # Test connection to DB
+        assert database.test_connection()
+
+        # Test running SQL
+        schema = random_str(8)
+        try:
+            assert database.run(f"CREATE SCHEMA IF NOT EXISTS {schema}") == ([], [])
+        finally:
+            database.run(f"DROP SCHEMA IF EXISTS {schema}")
+
+        # Test running invalid SQL
+        with pytest.raises(ComplexBadRequestError) as exc_info:
+            database.run("DROP SCHEMA xxx")
+        assert exc_info.value.response.status_code == 400
+        assert str(exc_info.value) == """\
+[
+    {
+        "message": "postgresql error: schema \\"xxx\\" does not exist\\n",
+        "error_type": "GENERIC_DB_ENGINE_ERROR",
+        "extra": {
+            "engine_name": "PostgreSQL",
+            "issue_codes": [
+                {
+                    "code": 1002,
+                    "message": "Issue 1002 - The database returned an unexpected error."
+                }
+            ]
+        }
+    }
+]"""
+
+        # Test updating the item
+        database.database_name = "XXX"
+        database.save()
+        assert superset_api.databases.get(id=database.id).database_name == "XXX"
+
+        # Test exporting and importing the item
+        with tempfile.NamedTemporaryFile(suffix=".zip") as f:
+            database.export(f.name)
+            superset_api.databases.import_file(f.name, overwrite=True)
+
+        # Test deleting the item
+        database.delete()
+
+        # Test deleting a non-existent item
+        with pytest.raises(BadRequestError) as exc_info:
+            database.delete()
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+    def test_datasets(self, superset_api, dataset, virtual_dataset):
+        # Test getting an invalid item ID
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.datasets.get(id=0)
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        assert isinstance(dataset.id, int)
+        assert superset_api.datasets.get(id=dataset.id).schema == dataset.schema
+        assert superset_api.datasets.get(id=dataset.id).table_name == dataset.table_name
+
+        # Test fetch
+        dataset.fetch()
+
+        # Test count
+        assert superset_api.datasets.count() >= 1
+
+        # Test creating a conflicting item
+        ds_id, dataset.id = dataset.id, None
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.datasets.add(dataset)
+        dataset.id = ds_id
+        assert exc_info.value.response.status_code == 422
+        assert exc_info.value.message == {"table_name": [f"Dataset {dataset.table_name} already exists"]}
+
+        # Test running SQL
+        with pytest.raises(ValueError) as exc_info:
+            dataset.run()
+        assert exc_info.value.args[0] == "Cannot run a dataset with no SQL"
+
+        # Test updating the item
+        dataset.table_name = "XXX"
+        dataset.save()
+        assert superset_api.datasets.get(id=dataset.id).table_name == "XXX"
+
+        # Test exporting and importing the item
+        with tempfile.NamedTemporaryFile(suffix=".zip") as f:
+            dataset.export(f.name)
+            superset_api.datasets.import_file(f.name, overwrite=True)
+
+        # Test deleting the item
+        dataset.delete()
+
+        # Test deleting a non-existent item
+        with pytest.raises(BadRequestError) as exc_info:
+            dataset.delete()
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        # Test virtual dataset
+        assert virtual_dataset.run() == ([], [])
+        # Insert more values than the query limit defined in superset_config.py::SqlExecuteRestApi.query_limit
+        values = ", ".join(f"({i})" for i in range(1, 21))
+        superset_api.run(
+            database_id=virtual_dataset.database_id,
+            query=f"INSERT INTO {virtual_dataset.schema}.{virtual_dataset.table_name} (i) VALUES {values}"
+        )
+        # We're getting results limited by the LIMIT value in virtual_dataset.sql
+        assert virtual_dataset.run() == (
+            [{'is_dttm': False, 'name': 'i', 'type': 'INTEGER'}],
+            [{'i': 1}, {'i': 2}, {'i': 3}],
+        )
+        # Limit the query result size. TODO: This doesn't actually work in Superset 2.x for some reason
+        assert virtual_dataset.run(query_limit=2) == (
+            [{'is_dttm': False, 'name': 'i', 'type': 'INTEGER'}],
+            [{'i': 1}, {'i': 2}, {'i': 3}],
+        )
+        virtual_dataset.sql = f"SELECT i FROM {virtual_dataset.schema}.{virtual_dataset.table_name} ORDER BY i LIMIT 15"
+        # We're hitting the query limit defined in superset_config.py::SqlExecuteRestApi.query_limit
+        with pytest.raises(QueryLimitReached) as exc_info:
+            virtual_dataset.run()
+        assert exc_info.value.args[0] == \
+               "You have exceeded the maximum number of rows that can be returned (10). Either set the `query_limit` " \
+               "attribute to a lower number than this, or add LIMIT keywords to your SQL statement to limit the " \
+               "number of rows returned."
+
+    def test_saved_queries(self, superset_api, saved_query):
+        # Test getting an invalid item ID
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.saved_queries.get(id=0)
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        assert isinstance(saved_query.id, int)
+        assert superset_api.saved_queries.get(id=saved_query.id).schema == saved_query.schema
+        assert superset_api.saved_queries.get(id=saved_query.id).label == saved_query.label
+
+        # Test fetch
+        saved_query.fetch()
+
+        # Test count
+        assert superset_api.saved_queries.count() >= 1
+
+        # Conflicting saved queries can apparently not be created
+
+        # Test running SQL
+        assert saved_query.run() == ([], [])
+        table_path = saved_query.sql.split("FROM")[-1].strip()  # Hacky, but it's the only way we can get the table name
+        superset_api.run(
+            database_id=saved_query.db_id,
+            query=f"INSERT INTO {table_path} (i) VALUES (1), (2), (3)"
+        )
+        assert saved_query.run() == (
+            [{'is_dttm': False, 'name': 'i', 'type': 'INTEGER'}],
+            [{'i': 1}, {'i': 2}, {'i': 3}],
+        )
+
+        # Test updating the item
+        saved_query.label = "XXX"
+        saved_query.save()
+        assert superset_api.saved_queries.get(id=saved_query.id).label == "XXX"
+
+        # Test exporting and importing the item
+        with tempfile.NamedTemporaryFile(suffix=".zip") as f:
+            saved_query.export(f.name)
+            superset_api.saved_queries.import_file(f.name, overwrite=True)
+
+        # Test deleting the item
+        saved_query.delete()
+
+        # Test deleting a non-existent item
+        with pytest.raises(BadRequestError) as exc_info:
+            saved_query.delete()
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+    def test_charts(self, superset_api, chart):
+        # Test getting an invalid item ID
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.charts.get(id=0)
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        assert isinstance(chart.id, int)
+        assert superset_api.charts.get(id=chart.id).slice_name == chart.slice_name
+        assert superset_api.charts.get(id=chart.id).viz_type == chart.viz_type
+
+        # Test fetch
+        chart.fetch()
+
+        # Test count
+        assert superset_api.charts.count() >= 1
+
+        # Conflicting charts can apparently not be created
+
+        # Test updating the item
+        chart.slice_name = "XXX"
+        chart.save()
+        assert superset_api.charts.get(id=chart.id).slice_name == "XXX"
+
+        # Test exporting and importing the item
+        with tempfile.NamedTemporaryFile(suffix=".zip") as f:
+            chart.export(f.name)
+            superset_api.charts.import_file(f.name, overwrite=True)
+
+        # Test deleting the item
+        chart.delete()
+
+        # Test deleting a non-existent item
+        with pytest.raises(BadRequestError) as exc_info:
+            chart.delete()
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+    def test_dashboards(self, superset_api, dashboard, chart):
+        # Test getting an invalid item ID
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.dashboards.get(id=0)
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+        assert isinstance(dashboard.id, int)
+        assert superset_api.dashboards.get(id=dashboard.id).dashboard_title == dashboard.dashboard_title
+
+        # Test fetch
+        dashboard.fetch()
+
+        # Test count
+        assert superset_api.dashboards.count() >= 1
+
+        # Test creating a conflicting item
+        d_id, dashboard.id = dashboard.id, None
+        with pytest.raises(BadRequestError) as exc_info:
+            superset_api.dashboards.add(dashboard)
+        dashboard.id = d_id
+        assert exc_info.value.response.status_code == 422
+        assert exc_info.value.message == {'slug': ['Must be unique']}
+
+        # Test updating the item
+        dashboard.dashboard_title = "XXX"
+        dashboard.save()
+        assert superset_api.dashboards.get(id=dashboard.id).dashboard_title == "XXX"
+
+        # Test changing colors
+        dashboard.update_colors({
+            "label": "#fcba03"
+        })
+        dashboard.save()
+        assert superset_api.dashboards.get(id=dashboard.id).colors == {"label": "#fcba03"}
+
+        # Test connected charts
+        chart.dashboards.append(dashboard.id)
+        chart.save()
+        assert [d["id"] for d in superset_api.charts.get(id=chart.id).dashboards] == [dashboard.id]
+        assert [c.id for c in superset_api.dashboards.get(id=dashboard.id).get_charts()] == [chart.id]
+
+        # Test exporting and importing the item
+        with tempfile.NamedTemporaryFile(suffix=".zip") as f:
+            dashboard.export(f.name)
+            superset_api.dashboards.import_file(f.name, overwrite=True)
+
+        # Test deleting the item
+        dashboard.delete()
+
+        # Test deleting a non-existent item
+        with pytest.raises(BadRequestError) as exc_info:
+            dashboard.delete()
+        assert exc_info.value.response.status_code == 404
+        assert exc_info.value.message == "Not found"
+
+    def test_find_one(self, superset_api):
+        title = random_str(8)
+        with pytest.raises(NotFound):
+            superset_api.dashboards.find_one(dashboard_title=title)
+
+        superset_api.dashboards.add(Dashboard(dashboard_title=title, published=True, slug=random_str(8)))
+        assert superset_api.dashboards.find_one(dashboard_title=title).dashboard_title == title
+
+        superset_api.dashboards.add(Dashboard(dashboard_title=title, published=True, slug=random_str(8)))
+        with pytest.raises(MultipleFound):
+            superset_api.dashboards.find_one(dashboard_title=title)
+
+
+class TestClient:
+    def test_no_verify(self, superset_url):
+        superset_api = SupersetClient(superset_url, "admin", "admin", verify=False)
+        assert superset_api.databases.count() >= 0  # Call something that makes an actual HTTP call
+
+    def test_refresh_token(self, requests_mock, superset_api):
+        # Mock test a token expiry message from the API
+        requests_mock.real_http = True
+        superset_api.databases.find()  # Call something that sets up the initial token
+        url = superset_api.databases.base_url
+
+        # HTTP 401 but no token expiry message in response
+        requests_mock.get(url, status_code=401)
+        r = superset_api.get(url)
+        assert superset_api.token_refresher(r=r).status_code == 401
+
+        # HTTP 401 but wrong message in response
+        requests_mock.get(url, status_code=401, content=json.dumps({"msg": "XXX"}).encode())
+        r = superset_api.get(url)
+        assert superset_api.token_refresher(r=r).status_code == 401
+
+        # HTTP 401 and proper token expiry message in response
+        requests_mock.get(url, [
+            dict(status_code=401, content=json.dumps({"msg": "Token has expired"}).encode()),
+            dict(status_code=200, content=json.dumps([]).encode()),
+        ])
+        r = superset_api.get(url)
+        assert superset_api.token_refresher(r=r).status_code == 200
+
+    def test_authenticate(self, monkeypatch, superset_url):
+        superset_api = SupersetClient(superset_url)  # No username or password yet
+
+        monkeypatch.setattr("getpass.getuser", lambda: "admin")
+        monkeypatch.setattr("getpass.getpass", lambda: "admin")
+        superset_api.authenticate()
+        assert superset_api.username == "admin"
+        assert superset_api._password == "admin"
+        assert superset_api.password == "*****"
+
+    def test_standard_sql_endpoint(self):
+        # Test the _sql_endpoint value returned by the plain SupersetClient class. The CustomClient class used in this
+        # test suite overrides it.
+        superset_api = SupersetClient("https://example.com")
+        assert superset_api._sql_endpoint == "https://example.com/superset/sql_json/"
+
+    def test_export_failure(self, requests_mock, superset_api):
+        requests_mock.real_http = True
+        requests_mock.get(superset_api.databases.export_url, status_code=200, headers={"content-type": "application/x"})
+        with pytest.raises(ValueError) as exc_info:
+            superset_api.databases.export([1], "/dev/null")
+        assert exc_info.value.args[0] == "Unknown content type application/x"
+
+
+class TestExceptions:
+    def test_raise_for_status(self, requests_mock):
+        url = "https://example.com"
+
+        requests_mock.get(url, status_code=500, content=b"XXX")
+        r = requests.get(url)
+        with pytest.raises(requests.exceptions.HTTPError):
+            raise_for_status(r)
+
+        requests_mock.get(url, status_code=500, content=json.dumps({"message": "XXX"}).encode())
+        r = requests.get(url)
+        with pytest.raises(BadRequestError) as exc_info:
+            raise_for_status(r)
+        assert exc_info.value.message == "XXX"
+
+        requests_mock.get(url, status_code=500, content=json.dumps({"errors": ["foo"]}).encode())
+        r = requests.get(url)
+        with pytest.raises(ComplexBadRequestError) as exc_info:
+            raise_for_status(r)
+        assert exc_info.value.errors == ["foo"]
diff --git a/tests/superset_config.py b/tests/superset_config.py
new file mode 100644
index 0000000..c91d7c2
--- /dev/null
+++ b/tests/superset_config.py
@@ -0,0 +1,83 @@
+"""File for configuring SuperSet, i.e. the superset_config.py file
+
+This config adds an /api/v1/execute_sql_json/ endpoint under OAuth authentication, to
+allow running queries through the API. Standard Superset only has /superset/sql_json/
+which is only accessible via a browser.
+"""
+import logging
+from typing import Optional, cast
+
+from cachelib.file import FileSystemCache
+from flask import request
+from flask_appbuilder.api import protect
+from flask_appbuilder.views import expose
+from superset.initialization import SupersetAppInitializer
+
+log = logging.getLogger(__name__)
+
+
+class CustomInitializer(SupersetAppInitializer):
+    def init_views(self) -> None:
+        from superset.exceptions import SupersetErrorsException
+        from superset.extensions import appbuilder
+        from superset.sqllab.command import CommandResult
+        from superset.sqllab.command_status import SqlJsonExecutionStatus
+        from superset.sqllab.exceptions import SqlLabException
+        from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext, SqlResults
+        from superset.views.base import BaseSupersetView, json_error_response, json_success
+        from superset.views.core import Superset
+
+        class CustomContext(SqlJsonExecutionContext):
+            def __init__(self, *args, **kwargs):
+                self.query_limit = kwargs.pop("query_limit")
+                super().__init__(*args, **kwargs)
+                # Override the SQL_MAX_ROW default set implicitly in the previous line. Just needs to be higher
+                # than 'displayLimit' so 'displayLimitReached' can be calculated correctly.
+                self.limit = self.query_limit + 1
+
+            def get_execution_result(self) -> Optional[SqlResults]:
+                self._sql_result["displayLimit"] = self.query_limit  # Add max rows info to query result
+                return self._sql_result
+
+        class SqlExecuteRestApi(BaseSupersetView):
+            """Make the superset.views.core.Superset.sql_json() endpoint available under /api/v1 and OAuth
+            protected
+            """
+
+            route_base = "/api/v1"
+            allow_browser_login = True
+            query_limit = 10  # NOTE: 1_000_000 or some other high value would be a better value for production
+
+            @expose("/execute_sql_json/", methods=["POST"])
+            @protect()
+            def execute_sql_json(self):
+                """Almost verbatim copy of superset.views.core.Superset.sql_json. Added non-default query
+                limits.
+                """
+                try:
+                    log_params = {"user_agent": cast(Optional[str], request.headers.get("USER_AGENT"))}
+                    execution_context = CustomContext(request.json, query_limit=self.query_limit)
+                    command = Superset._create_sql_json_command(execution_context, log_params)
+                    # Override the DISPLAY_MAX_ROW default set implicitly in the previous line
+                    command._execution_context_convertor.set_max_row_in_display(self.query_limit)
+                    try:
+                        command_result: CommandResult = command.run()
+                    except SupersetErrorsException as e:
+                        payload = {"errors": [ex.to_dict() for ex in e.errors]}
+                        return json_error_response(status=400, payload=payload)
+                    status_code = 200
+                    if command_result["status"] == SqlJsonExecutionStatus.QUERY_IS_RUNNING:
+                        status_code = 202
+                    return json_success(command_result["payload"], status_code)
+                except SqlLabException as ex:
+                    log.error(ex.message)
+                    Superset._set_http_status_into_Sql_lab_exception(ex)
+                    payload = {"errors": [ex.to_dict()]}
+                    return json_error_response(status=ex.status, payload=payload)
+
+        super().init_views()
+        appbuilder.add_view_no_menu(SqlExecuteRestApi)
+
+
+APP_INITIALIZER = CustomInitializer
+RESULTS_BACKEND = FileSystemCache("/app/superset_home/sqllab")