diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 4f2b906..d40951e 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -410,6 +410,13 @@ class Table: elif len(pks) > 1: pk_names = pks last_pk = pk_values + if len(pk_names) != len(pk_values): + raise NotFoundError( + "Need {} primary key value{}".format( + len(pk_names), "" if len(pk_names) == 1 else "s" + ) + ) + wheres = ["[{}] = ?".format(pk_name) for pk_name in pk_names] rows = self.rows_where(" and ".join(wheres), pk_values) try: diff --git a/tests/test_get.py b/tests/test_get.py index bd79490..ae5ffdf 100644 --- a/tests/test_get.py +++ b/tests/test_get.py @@ -1,4 +1,5 @@ import pytest +from sqlite_utils.db import NotFoundError def test_get_rowid(fresh_db): @@ -11,10 +12,25 @@ def test_get_rowid(fresh_db): def test_get_primary_key(fresh_db): dogs = fresh_db["dogs"] cleo = {"name": "Cleo", "age": 4, "id": 5} - row_id = dogs.insert(cleo, pk="id").last_pk + last_pk = dogs.insert(cleo, pk="id").last_pk + assert 5 == last_pk assert cleo == dogs.get(5) +@pytest.mark.parametrize( + "argument,expected_msg", + [(100, None), (None, None), ((1, 2), "Need 1 primary key value"), ("2", None)], +) +def test_get_not_found(argument, expected_msg, fresh_db): + fresh_db["dogs"].insert( + {"id": 1, "name": "Cleo", "age": 4, "is_good": True}, pk="id" + ) + with pytest.raises(NotFoundError) as excinfo: + fresh_db["dogs"].get(argument) + if expected_msg is not None: + assert expected_msg == excinfo.value.args[0] + + @pytest.mark.parametrize( "where,where_args,expected_ids", [