Repository: ploomber/jupysql (python)

Issue:
better argument validation in `%sqlplot`
In sqlplot we have `nargs="*"`: https://github.com/ploomber/jupysql/blob/2e847539185ebe99175e9eb89b35b4da5f516f85/src/sql/magic_plot.py#L29

this causes the parser to take n arguments, changing it (like we have it in `%sqlcmd`: https://github.com/ploomber/jupysql/blob/2e847539185ebe99175e9eb89b35b4da5f516f85/src/sql/magic_cmd.py#L47) will allow us to better validate the input arguments.



Pull request:
Remove nargs from plotting
## Describe your changes

## Issue number

Closes #595 

## Checklist before requesting a review

- [ ] Performed a self-review of my code
- [ ] Formatted my code with [`pkgmt format`](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#linting-formatting)
- [ ] Added [tests](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#testing) (when necessary).
- [ ] Added [docstring](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#documenting-changes-and-new-features) documentation and update the [changelog](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#changelog) (when needed)



<!-- readthedocs-preview jupysql start -->
----
:books: Documentation preview :books:: https://jupysql--839.org.readthedocs.build/en/839/

<!-- readthedocs-preview jupysql end -->

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9b44d24..d2d69c3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## 0.10.2dev
 
+* [Feature] Improved messages when loading configurations from `pyproject.toml` file.
+* [Feature] Add `--schema/-s` for `%sqlcmd` commands that support `--table/-t` and ensure `--table schema.table` works (#519)
+
 ## 0.10.1 (2023-08-30)
 
 * [Feature] Automatically connect if the `dsn_filename` (defaults to `~/.jupysql/connections.ini`) contains a `default` section
@@ -431,4 +434,4 @@ Converted from an IPython Plugin to an Extension for 1.0 compatibility
 
 *Release date: 21-Mar-2013*
 
-* Initial release
+* Initial release
\ No newline at end of file
diff --git a/src/sql/cmd/explore.py b/src/sql/cmd/explore.py
index 15daee0..3a089d1 100644
--- a/src/sql/cmd/explore.py
+++ b/src/sql/cmd/explore.py
@@ -1,5 +1,4 @@
 from sql.widgets import TableWidget
-from IPython.display import display
 from sql.cmd.cmd_utils import CmdParser
 
 
@@ -18,7 +17,7 @@ def explore(others):
     """
     parser = CmdParser()
     parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+    parser.add_argument("-s", "--schema", type=str, help="Schema name", required=False)
     args = parser.parse_args(others)
-
-    table_widget = TableWidget(args.table)
-    display(table_widget)
+    table_widget = TableWidget(args.table, args.schema)
+    return table_widget
diff --git a/src/sql/display.py b/src/sql/display.py
index ae0df3d..8837171 100644
--- a/src/sql/display.py
+++ b/src/sql/display.py
@@ -94,7 +94,7 @@ def message_success(message):
 
 
 def message_warning(message):
-    """Display a success message"""
+    """Display a warning message"""
     display(
         Message(
             message,
diff --git a/src/sql/inspect.py b/src/sql/inspect.py
index 62dad30..c21051d 100644
--- a/src/sql/inspect.py
+++ b/src/sql/inspect.py
@@ -180,6 +180,8 @@ class Columns(DatabaseInspection):
         # this returns a list of dictionaries. e.g.,
         # [{"name": "column_a", "type": "INT"}
         #  {"name": "column_b", "type": "FLOAT"}]
+        if not schema and "." in name:
+            schema, name = name.split(".")
         columns = inspector.get_columns(name, schema) or []
 
         self._table = PrettyTable()
diff --git a/src/sql/magic.py b/src/sql/magic.py
index 4734ed9..557b3ff 100644
--- a/src/sql/magic.py
+++ b/src/sql/magic.py
@@ -647,7 +647,7 @@ class SqlMagic(Magics, Configurable):
 def set_configs(ip, file_path):
     """Set user defined SqlMagic configuration settings"""
     sql = ip.find_cell_magic("sql").__self__
-    user_configs = util.get_user_configs(file_path, ["tool", "jupysql", "SqlMagic"])
+    user_configs = util.get_user_configs(file_path)
     default_configs = util.get_default_configs(sql)
     table_rows = []
     for config, value in user_configs.items():
diff --git a/src/sql/magic_plot.py b/src/sql/magic_plot.py
index 33f2445..1d01bea 100644
--- a/src/sql/magic_plot.py
+++ b/src/sql/magic_plot.py
@@ -1,8 +1,4 @@
-from IPython.core.magic import (
-    Magics,
-    line_magic,
-    magics_class,
-)
+from IPython.core.magic import Magics, line_magic, magics_class
 from IPython.core.magic_arguments import argument, magic_arguments
 from ploomber_core.exceptions import modify_exceptions
 
@@ -14,7 +10,6 @@ except ModuleNotFoundError:
 
 from sql import plot
 from sql.command import SQLPlotCommand
-from sql import exceptions
 from sql import util
 from sql.inspect import is_table_exists
 from sql.store import is_saved_snippet
@@ -28,7 +23,12 @@ class SqlPlotMagic(Magics, Configurable):
 
     @line_magic("sqlplot")
     @magic_arguments()
-    @argument("line", default="", nargs="*", type=str, help="Plot name")
+    @argument(
+        "plot_name",
+        type=str,
+        help="Plot name",
+        choices=["histogram", "hist", "boxplot", "box", "bar", "pie"],
+    )
     @argument("-t", "--table", type=str, help="Table to use", required=True)
     @argument(
         "-c", "--column", type=str, nargs="+", help="Column(s) to use", required=True
@@ -87,20 +87,6 @@ class SqlPlotMagic(Magics, Configurable):
         else:
             column = cmd.args.column
 
-        if not cmd.args.line:
-            plot_str = util.pretty_print(SUPPORTED_PLOTS, last_delimiter="or")
-            raise exceptions.UsageError(
-                "Missing the first argument, must be any of: "
-                f"{plot_str}\n"
-                "Example: %sqlplot histogram"
-            )
-
-        if cmd.args.line[0] not in SUPPORTED_PLOTS + ["hist", "box"]:
-            plot_str = util.pretty_print(SUPPORTED_PLOTS, last_delimiter="or")
-            raise exceptions.UsageError(
-                f"Unknown plot {cmd.args.line[0]!r}. Must be any of: " f"{plot_str}"
-            )
-
         column = util.sanitize_identifier(column)
         table = util.sanitize_identifier(cmd.args.table)
 
@@ -109,7 +95,7 @@ class SqlPlotMagic(Magics, Configurable):
         else:
             with_ = self._check_table_exists(table)
 
-        if cmd.args.line[0] in {"box", "boxplot"}:
+        if cmd.args.plot_name in {"box", "boxplot"}:
             return plot.boxplot(
                 table=table,
                 column=column,
@@ -117,7 +103,7 @@ class SqlPlotMagic(Magics, Configurable):
                 orient=cmd.args.orient,
                 conn=None,
             )
-        elif cmd.args.line[0] in {"hist", "histogram"}:
+        elif cmd.args.plot_name in {"hist", "histogram"}:
             # to avoid passing bins default value when breaks or binwidth is specified
             bin_specified = " --bins " in line or " -b " in line
             breaks_specified = " --breaks " in line or " -B " in line
@@ -135,7 +121,7 @@ class SqlPlotMagic(Magics, Configurable):
                 breaks=cmd.args.breaks,
                 binwidth=cmd.args.binwidth,
             )
-        elif cmd.args.line[0] in {"bar"}:
+        elif cmd.args.plot_name in {"bar"}:
             return plot.bar(
                 table=table,
                 column=column,
@@ -144,7 +130,7 @@ class SqlPlotMagic(Magics, Configurable):
                 show_num=cmd.args.show_numbers,
                 conn=None,
             )
-        elif cmd.args.line[0] in {"pie"}:
+        elif cmd.args.plot_name in {"pie"}:
             return plot.pie(
                 table=table,
                 column=column,
diff --git a/src/sql/util.py b/src/sql/util.py
index 6ea59fc..7cdbfcf 100644
--- a/src/sql/util.py
+++ b/src/sql/util.py
@@ -18,6 +18,8 @@ except ModuleNotFoundError:
 SINGLE_QUOTE = "'"
 DOUBLE_QUOTE = '"'
 
+CONFIGURATION_DOCS_STR = "https://jupysql.ploomber.io/en/latest/api/configuration.html#loading-from-pyproject-toml"  # noqa
+
 
 def sanitize_identifier(identifier):
     if (identifier[0] == SINGLE_QUOTE and identifier[-1] == SINGLE_QUOTE) or (
@@ -164,7 +166,6 @@ def find_path_from_root(file_name):
             return None
 
         current = current.parent
-    display.message(f"Found {file_name} from '{current}'")
 
     return str(Path(current, file_name))
 
@@ -230,22 +231,24 @@ def load_toml(file_path):
 def parse_toml_error(e, file_path):
     eline, ekey, evalue = get_line_content_from_toml(file_path, e.lineno)
     if "Duplicate keys!" in str(e):
-        return exceptions.ConfigurationError(f"Duplicate key found : '{ekey}'")
+        return exceptions.ConfigurationError(
+            f"Duplicate key found: '{ekey}' in {file_path}"
+        )
     elif "Only all lowercase booleans" in str(e):
         return exceptions.ConfigurationError(
-            f"Invalid value '{evalue}' in '{eline}'. "
+            f"Invalid value '{evalue}' in '{eline}' in {file_path}. "
             "Valid boolean values: true, false"
         )
     elif "invalid literal for int()" in str(e):
         return exceptions.ConfigurationError(
-            f"Invalid value '{evalue}' in '{eline}'. "
+            f"Invalid value '{evalue}' in '{eline}' in {file_path}. "
             "To use str value, enclose it with ' or \"."
         )
     else:
         return e
 
 
-def get_user_configs(file_path, section_names):
+def get_user_configs(file_path):
     """
     Returns saved configuration settings in a toml file from given file_path
 
@@ -253,9 +256,6 @@ def get_user_configs(file_path, section_names):
     ----------
     file_path : str
         file path to a toml file
-    section_names : list
-        section names that contains the configuration settings
-        (e.g., ["tool", "jupysql", "SqlMagic"])
 
     Returns
     -------
@@ -263,18 +263,40 @@ def get_user_configs(file_path, section_names):
         saved configuration settings
     """
     data = load_toml(file_path)
+    section_names = ["tool", "jupysql", "SqlMagic"]
     while section_names:
         section_to_find, sections_from_user = section_names.pop(0), data.keys()
         if section_to_find not in sections_from_user:
             close_match = difflib.get_close_matches(section_to_find, sections_from_user)
             if not close_match:
+                MESSAGE_PREFIX = (
+                    f"Tip: You may define configurations in "
+                    f"{file_path}. Please review our "
+                )
+                display.message_html(
+                    f"{MESSAGE_PREFIX}<a href='{CONFIGURATION_DOCS_STR}'>"
+                    "configuration guideline</a>."
+                )
                 return {}
             else:
                 raise exceptions.ConfigurationError(
-                    f"{pretty_print(close_match)} is an invalid section name. "
+                    f"{pretty_print(close_match)} is an invalid section "
+                    f"name in {file_path}. "
                     f"Did you mean '{section_to_find}'?"
                 )
         data = data[section_to_find]
+    if not data:
+        if section_to_find == "SqlMagic":
+            MESSAGE_PREFIX = (
+                f"[tool.jupysql.SqlMagic] present in {file_path} but empty. "
+                f"Please review our "
+            )
+            display.message_html(
+                f"{MESSAGE_PREFIX}<a href='{CONFIGURATION_DOCS_STR}'>"
+                "configuration guideline</a>."
+            )
+    else:
+        display.message(f"Loading configurations from {file_path}")
     return data
 
 
diff --git a/src/sql/widgets/table_widget/table_widget.py b/src/sql/widgets/table_widget/table_widget.py
index 04cc85d..17da7ec 100644
--- a/src/sql/widgets/table_widget/table_widget.py
+++ b/src/sql/widgets/table_widget/table_widget.py
@@ -16,7 +16,7 @@ BASE_DIR = os.path.dirname(__file__)
 
 class TableWidget:
     @telemetry.log_call("TableWidget-init")
-    def __init__(self, table):
+    def __init__(self, table, schema=None):
         """
         Creates an HTML table element and populates it with SQL table
 
@@ -28,13 +28,13 @@ class TableWidget:
 
         self.html = ""
 
-        is_table_exists(table)
+        is_table_exists(table, schema)
 
         # load css
         html_style = utils.load_css(f"{BASE_DIR}/css/tableWidget.css")
         self.add_to_html(html_style)
 
-        self.create_table(table)
+        self.create_table(table, schema)
 
         # register listener for jupyter lab
         self.register_comm()
@@ -48,17 +48,22 @@ class TableWidget:
     def add_to_html(self, html):
         self.html += html
 
-    def create_table(self, table):
+    def create_table(self, table, schema):
         """
         Creates an HTML table with default data
         """
+        if schema:
+            table_ = f"{schema}.{table}"
+        else:
+            table_ = table
+
         rows_per_page = 10
-        rows, columns = fetch_sql_with_pagination(table, 0, rows_per_page)
+        rows, columns = fetch_sql_with_pagination(table_, 0, rows_per_page)
         rows = parse_sql_results_to_json(rows, columns)
 
-        query = f"SELECT count(*) FROM {table}"
+        query = f"SELECT count(*) FROM {table_}"
         n_total = ConnectionManager.current.raw_execute(query).fetchone()[0]
-        table_name = table.strip('"').strip("'")
+        table_name = table_.strip('"').strip("'")
 
         n_pages = math.ceil(n_total / rows_per_page)
 
@@ -81,7 +86,7 @@ class TableWidget:
                     n_total=n_total,
                     table_name=table_name,
                     table_container_id=table_container_id,
-                    table=table,
+                    table=table_,
                     initialRows=rows,
                 ),
             ]


Test patch:
diff --git a/src/sql/cmd/test.py b/src/sql/cmd/test.py
index be79c74..0390ed9 100644
--- a/src/sql/cmd/test.py
+++ b/src/sql/cmd/test.py
@@ -23,7 +23,11 @@ def return_test_results(args, conn, query):
 
 
 def run_each_individually(args, conn):
-    base_query = select("*").from_(args.table)
+    if args.schema:
+        table_ = f"{args.schema}.{args.table}"
+    else:
+        table_ = args.table
+    base_query = select("*").from_(table_)
 
     storage = {}
 
@@ -99,6 +103,7 @@ def test(others):
     parser = CmdParser()
 
     parser.add_argument("-t", "--table", type=str, help="Table name", required=True)
+    parser.add_argument("-s", "--schema", type=str, help="Schema name", required=False)
     parser.add_argument("-c", "--column", type=str, help="Column name", required=False)
     parser.add_argument(
         "-g",
diff --git a/src/tests/test_config.py b/src/tests/test_config.py
index 52312bb..ac07227 100644
--- a/src/tests/test_config.py
+++ b/src/tests/test_config.py
@@ -4,10 +4,12 @@ from pathlib import Path
 
 import pytest
 import sys
+from unittest.mock import Mock
 
 from sql.magic import load_ipython_extension
 from sql.connection import ConnectionManager
-from sql.util import get_default_configs
+from sql.util import get_default_configs, CONFIGURATION_DOCS_STR
+from sql import display
 from IPython.core.error import UsageError
 
 
@@ -146,87 +148,107 @@ dsn_filename = myconnections.ini
     assert "Could not load configuration file" in captured.out
 
 
-@pytest.mark.parametrize(
-    "file_content, expect, config_expected",
-    [
-        (
-            """
+def test_loading_valid_pyproject_toml_shows_feedback_and_modifies_config(
+    tmp_empty, ip_no_magics, capsys
+):
+    Path("pyproject.toml").write_text(
+        """
 [tool.jupysql.SqlMagic]
 autocommit = false
 autolimit = 1
 style = "RANDOM"
-""",
-            [
-                "Found pyproject.toml from '%s'",
-                "Settings changed:",
-                r"autocommit\s*\|\s*False",
-                r"autolimit\s*\|\s*1",
-                r"style\s*\|\s*RANDOM",
-            ],
-            {"autocommit": False, "autolimit": 1, "style": "RANDOM"},
-        ),
+"""
+    )
+
+    expect = [
+        "Loading configurations from {path}",
+        "Settings changed:",
+        r"autocommit\s*\|\s*False",
+        r"autolimit\s*\|\s*1",
+        r"style\s*\|\s*RANDOM",
+    ]
+
+    config_expected = {"autocommit": False, "autolimit": 1, "style": "RANDOM"}
+
+    toml_path = str(Path(os.getcwd()).joinpath("pyproject.toml"))
+
+    os.mkdir("sub")
+    os.chdir("sub")
+
+    load_ipython_extension(ip_no_magics)
+    magic = ip_no_magics.find_magic("sql").__self__
+    combined = {**get_default_testing_configs(magic), **config_expected}
+    out, _ = capsys.readouterr()
+    expect[0] = expect[0].format(path=re.escape(toml_path))
+    assert all(re.search(substring, out) for substring in expect)
+    assert get_current_configs(magic) == combined
+
+
+@pytest.mark.parametrize(
+    "file_content, param",
+    [
         (
             """
 [tool.jupysql.SqlMagic]
 """,
-            ["Found pyproject.toml from '%s'"],
-            {},
+            "[tool.jupysql.SqlMagic] present in {path} but empty.",
         ),
+        ("", "Tip: You may define configurations in {path}."),
+    ],
+    ids=["empty_sqlmagic_key", "missing_sqlmagic_key"],
+)
+def test_loading_pyproject_toml_display_configuration_docs_link(
+    tmp_empty, ip_no_magics, file_content, param, monkeypatch
+):
+    Path("pyproject.toml").write_text(file_content)
+    toml_path = str(Path(os.getcwd()).joinpath("pyproject.toml"))
+
+    os.mkdir("sub")
+    os.chdir("sub")
+    mock = Mock()
+    monkeypatch.setattr(display, "message_html", mock)
+
+    load_ipython_extension(ip_no_magics)
+    param = (
+        f"{param.format(path=toml_path)} Please review our "
+        f"<a href='{CONFIGURATION_DOCS_STR}'>configuration guideline</a>."
+    )
+    mock.assert_called_once_with(param)
+
+
+@pytest.mark.parametrize(
+    "file_content",
+    [
         (
             """
 [test]
 github = "ploomber/jupysql"
-""",
-            ["Found pyproject.toml from '%s'"],
-            {},
+"""
         ),
         (
             """
 [tool.pkgmt]
 github = "ploomber/jupysql"
-""",
-            ["Found pyproject.toml from '%s'"],
-            {},
+"""
         ),
         (
             """
 [tool.jupysql.test]
 github = "ploomber/jupysql"
-""",
-            ["Found pyproject.toml from '%s'"],
-            {},
-        ),
-        (
-            "",
-            ["Found pyproject.toml from '%s'"],
-            {},
+"""
         ),
     ],
 )
-def test_loading_valid_pyproject_toml_shows_feedback_and_modifies_config(
-    tmp_empty,
-    ip_no_magics,
-    capsys,
-    file_content,
-    expect,
-    config_expected,
+def test_load_pyproject_toml_user_configurations_not_specified(
+    tmp_empty, ip_no_magics, capsys, file_content
 ):
     Path("pyproject.toml").write_text(file_content)
-    toml_dir = os.getcwd()
-
     os.mkdir("sub")
     os.chdir("sub")
 
     load_ipython_extension(ip_no_magics)
-
-    magic = ip_no_magics.find_magic("sql").__self__
-
-    combined = {**get_default_testing_configs(magic), **config_expected}
     out, _ = capsys.readouterr()
-
-    expect[0] = expect[0] % (re.escape(toml_dir))
-    assert all(re.search(substring, out) for substring in expect)
-    assert get_current_configs(magic) == combined
+    assert "Loading configurations from" not in out
 
 
 @pytest.mark.parametrize(
@@ -238,14 +260,14 @@ def test_loading_valid_pyproject_toml_shows_feedback_and_modifies_config(
 autocommit = true
 autocommit = true
 """,
-            "Duplicate key found : 'autocommit'",
+            "Duplicate key found: 'autocommit' in {path}",
         ),
         (
             """
 [tool.jupySql.SqlMagic]
 autocommit = true
 """,
-            "'jupySql' is an invalid section name. Did you mean 'jupysql'?",
+            "'jupySql' is an invalid section name in {path}. Did you mean 'jupysql'?",
         ),
         (
             """
@@ -253,7 +275,7 @@ autocommit = true
 autocommit = True
 """,
             (
-                "Invalid value 'True' in 'autocommit = True'. "
+                "Invalid value 'True' in 'autocommit = True' in {path}. "
                 "Valid boolean values: true, false"
             ),
         ),
@@ -263,7 +285,7 @@ autocommit = True
 autocommit = invalid
 """,
             (
-                "Invalid value 'invalid' in 'autocommit = invalid'. "
+                "Invalid value 'invalid' in 'autocommit = invalid' in {path}. "
                 "To use str value, enclose it with ' or \"."
             ),
         ),
@@ -273,8 +295,7 @@ def test_error_on_toml_parsing(
     tmp_empty, ip_no_magics, capsys, file_content, error_msg
 ):
     Path("pyproject.toml").write_text(file_content)
-    toml_dir = os.getcwd()
-    found_statement = "Found pyproject.toml from '%s'" % (toml_dir)
+    toml_path = str(Path(os.getcwd()).joinpath("pyproject.toml"))
     os.makedirs("sub")
     os.chdir("sub")
 
@@ -283,9 +304,8 @@ def test_error_on_toml_parsing(
 
     out, _ = capsys.readouterr()
 
-    assert out.strip() == found_statement
     assert excinfo.value.error_type == "ConfigurationError"
-    assert str(excinfo.value) == error_msg
+    assert str(excinfo.value) == error_msg.format(path=toml_path)
 
 
 def test_valid_and_invalid_configs(tmp_empty, ip_no_magics, capsys):
@@ -299,14 +319,14 @@ invalid = false
 displaycon = false
 """
     )
-    toml_dir = os.getcwd()
+    toml_path = str(Path(os.getcwd()).joinpath("pyproject.toml"))
     os.makedirs("sub")
     os.chdir("sub")
 
     load_ipython_extension(ip_no_magics)
     out, _ = capsys.readouterr()
     expect = [
-        "Found pyproject.toml from '%s'" % (re.escape(toml_dir)),
+        f"Loading configurations from {re.escape(toml_path)}",
         "'autocomm' is an invalid configuration. Did you mean 'autocommit'?",
         (
             "'autop' is an invalid configuration. "
diff --git a/src/tests/test_magic_cmd.py b/src/tests/test_magic_cmd.py
index 4e8f7a1..d093eda 100644
--- a/src/tests/test_magic_cmd.py
+++ b/src/tests/test_magic_cmd.py
@@ -7,6 +7,7 @@ from sqlalchemy import create_engine
 from sql.connection import SQLAlchemyConnection
 from sql.inspect import _is_numeric
 from sql.display import Table, Message
+from sql.widgets import TableWidget
 from jupysql_plugin.widgets import ConnectorWidget
 import duckdb
 import sqlite3
@@ -81,6 +82,19 @@ def test_snippet_ip(ip):
     yield ip
 
 
+@pytest.fixture
+def sample_schema_with_table(ip_empty):
+    ip_empty.run_cell("%sql duckdb://")
+    ip_empty.run_cell(
+        """%%sql
+CREATE SCHEMA schema1;
+CREATE TABLE schema1.table1 (x INT, y TEXT);
+INSERT INTO schema1.table1 VALUES (1, 'one');
+INSERT INTO schema1.table1 VALUES (2, 'two');
+"""
+    )
+
+
 @pytest.mark.parametrize(
     "cell, error_message",
     [
@@ -210,7 +224,10 @@ def test_columns(ip, cmd, cols):
     assert all(col in out for col in cols)
 
 
-def test_columns_with_schema(ip, tmp_empty):
+@pytest.mark.parametrize(
+    "arguments", ["--table numbers --schema some_schema", "--table some_schema.numbers"]
+)
+def test_columns_with_schema(ip, tmp_empty, arguments):
     conn = SQLAlchemyConnection(engine=create_engine("sqlite:///my.db"))
     conn.execute("CREATE TABLE numbers (some_number FLOAT)")
 
@@ -220,9 +237,7 @@ ATTACH DATABASE 'my.db' AS some_schema
 """
     )
 
-    out = ip.run_cell(
-        "%sqlcmd columns --table numbers --schema some_schema"
-    ).result._repr_html_()
+    out = ip.run_cell(f"%sqlcmd columns {arguments}").result._repr_html_()
 
     assert "some_number" in out
 
@@ -345,7 +360,10 @@ def test_table_profile_with_stdev(ip_with_connections, tmp_empty, conn):
     assert "position: sticky;" in out._table_html
 
 
-def test_table_schema_profile(ip, tmp_empty):
+@pytest.mark.parametrize(
+    "arguments", ["--table t --schema b_schema", "--table b_schema.t"]
+)
+def test_table_schema_profile(ip, tmp_empty, arguments):
     ip.run_cell("%sql sqlite:///a.db")
     ip.run_cell("%sql CREATE TABLE t (n FLOAT)")
     ip.run_cell("%sql INSERT INTO t VALUES (1)")
@@ -379,7 +397,7 @@ def test_table_schema_profile(ip, tmp_empty):
         "top": [math.nan],
     }
 
-    out = ip.run_cell("%sqlcmd profile -t t --schema b_schema").result
+    out = ip.run_cell(f"%sqlcmd profile {arguments}").result
 
     stats_table = out._table
 
@@ -392,7 +410,11 @@ def test_table_schema_profile(ip, tmp_empty):
             assert cell == str(expected[profile_metric][0])
 
 
-def test_sqlcmd_profile_with_schema_argument_and_dbapi(ip_empty, tmp_empty):
+@pytest.mark.parametrize(
+    "arguments",
+    ["--table sample_table --schema test_schema", "--table test_schema.sample_table"],
+)
+def test_sqlcmd_profile_with_schema_argument_and_dbapi(ip_empty, tmp_empty, arguments):
     sqlite_dbapi_testdb_conn = sqlite3.connect("test.db")
     ip_empty.push({"sqlite_dbapi_testdb_conn": sqlite_dbapi_testdb_conn})
 
@@ -423,9 +445,7 @@ INSERT INTO sample_table VALUES (33);
         "top": [math.nan],
     }
 
-    out = ip_empty.run_cell(
-        "%sqlcmd profile --table sample_table --schema test_schema"
-    ).result
+    out = ip_empty.run_cell(f"%sqlcmd profile {arguments}").result
 
     stats_table = out._table
 
@@ -569,6 +589,26 @@ def test_test_error(ip, cell, error_message):
     assert str(excinfo.value) == error_message
 
 
+@pytest.mark.parametrize(
+    "arguments", ["--table schema1.table1", "--table table1 --schema schema1"]
+)
+def test_failing_test_with_schema(ip_empty, sample_schema_with_table, arguments):
+    expected_error_message = "The above values do not match your test requirements."
+
+    with pytest.raises(UsageError) as excinfo:
+        ip_empty.run_cell(f"%sqlcmd test {arguments} --column x --less-than 2")
+
+    assert expected_error_message in str(excinfo.value)
+
+
+@pytest.mark.parametrize(
+    "arguments", ["--table schema1.table1", "--table table1 --schema schema1"]
+)
+def test_passing_test_with_schema(ip_empty, sample_schema_with_table, arguments):
+    out = ip_empty.run_cell(f"%sqlcmd test {arguments} --column x --less-than 3").result
+    assert out is True
+
+
 @pytest.mark.parametrize(
     "cmds, result",
     [
@@ -749,6 +789,17 @@ def test_delete_invalid_snippet(arg, ip_snippets):
     assert str(excinfo.value) == "No such saved snippet found : non_existent_snippet"
 
 
+@pytest.mark.parametrize(
+    "arguments", ["--table schema1.table1", "--table table1 --schema schema1"]
+)
+def test_explore_with_schema(ip_empty, sample_schema_with_table, arguments):
+    expected_rows = ['"x": 1', '"y": "one"', '"x": 2', '"y": "two"']
+
+    out = ip_empty.run_cell(f"%sqlcmd explore {arguments}").result
+    assert isinstance(out, TableWidget)
+    assert [row in out._repr_html_() for row in expected_rows]
+
+
 @pytest.mark.parametrize(
     "file_content, stored_conns",
     [
diff --git a/src/tests/test_magic_plot.py b/src/tests/test_magic_plot.py
index e6f9fec..0cf0a91 100644
--- a/src/tests/test_magic_plot.py
+++ b/src/tests/test_magic_plot.py
@@ -45,19 +45,21 @@ WHERE x > 2
     [
         [
             "%sqlplot someplot -t a -c b",
-            f"Unknown plot 'someplot'. Must be any of: {plot_str}",
+            "argument plot_name: invalid choice: 'someplot' "
+            "(choose from 'histogram', 'hist', 'boxplot', 'box', 'bar', 'pie')",
         ],
         [
             "%sqlplot -t a -c b",
-            f"Missing the first argument, must be any of: {plot_str}",
+            "the following arguments are required: plot_name",
         ],
     ],
+    ids=["invalid_plot_name", "missing_plot_name"],
 )
 def test_validate_plot_name(tmp_empty, ip, cell, error_message):
     with pytest.raises(UsageError) as excinfo:
         ip.run_cell(cell)
 
-    assert excinfo.value.error_type == "UsageError"
+    assert excinfo.typename == "UsageError"
     assert str(error_message).lower() in str(excinfo.value).lower()