File size: 1,112 Bytes
87151cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import pytest
from llm_benchmark.sql.query import SqlQuery
@pytest.mark.parametrize(
"name, expected",
[
("Presence", True),
("Roundabout", False),
("' OR 1=1 --", False),
("Presence'; DROP TABLE Album; --", False),
("This Album Doesn't Exist", False),
],
)
def test_query_album(name: str, expected: bool) -> None:
result = SqlQuery.query_album(name)
assert isinstance(result, bool)
assert result is expected
def test_benchmark_query_album(benchmark) -> None:
benchmark(SqlQuery.query_album, "Presence")
def test_join_albums() -> None:
assert SqlQuery.join_albums()[0] == (
"For Those About To Rock (We Salute You)",
"For Those About To Rock We Salute You",
"AC/DC",
)
def test_benchmark_join_albums(benchmark) -> None:
benchmark(SqlQuery.join_albums)
def test_top_invoices() -> None:
top = SqlQuery.top_invoices()
assert top[0][2] == 25.86
assert top[2][2] == 21.86
assert len(top) == 10
def test_benchmark_top_invoices(benchmark) -> None:
benchmark(SqlQuery.top_invoices)
|