from __future__ import annotations

from type_aliases import ConnDB

def test_iteration_list(conn_db_in_mem: ConnDB) -> None:

conn, _ = conn_db_in_mem

conn.execute("CREATE NODE TABLE person(name STRING, age INT64, PRIMARY KEY(name));")

conn.execute("CREATE (:person {name: 'Alice', age: 30});")

conn.execute("CREATE (:person {name: 'Bob', age: 40});")

result = conn.execute("MATCH (p:person) RETURN p.*")

assert result.get_num_tuples() == 2

rows = list(result)

assert rows[0] == ["Alice", 30]

assert rows[1] == ["Bob", 40]

def test_iteration_loop(conn_db_in_mem: ConnDB) -> None:

conn, _ = conn_db_in_mem

conn.execute("CREATE NODE TABLE person(name STRING, age INT64, PRIMARY KEY(name));")

conn.execute("CREATE (:person {name: 'Alice', age: 30});")

conn.execute("CREATE (:person {name: 'Bob', age: 40});")

result = conn.execute("MATCH (p:person) RETURN p.*")

assert result.get_num_tuples() == 2

rows = []

for row in result:

rows.append(row)

assert rows[0] == ["Alice", 30]

assert rows[1] == ["Bob", 40]

def test_get_all(conn_db_in_mem: ConnDB) -> None:

conn, _ = conn_db_in_mem

conn.execute("CREATE NODE TABLE person(name STRING, age INT64, PRIMARY KEY(name));")

conn.execute("CREATE (:person {name: 'Alice', age: 30});")

conn.execute("CREATE (:person {name: 'Bob', age: 40});")

result = conn.execute("MATCH (p:person) RETURN p.*")

assert result.get_num_tuples() == 2

rows = result.get_all()

assert rows[0] == ["Alice", 30]

assert rows[1] == ["Bob", 40]

def test_get_n(conn_db_in_mem: ConnDB) -> None:

conn, _ = conn_db_in_mem

conn.execute("CREATE NODE TABLE person(name STRING, age INT64, PRIMARY KEY(name));")

conn.execute("CREATE (:person {name: 'Alice', age: 30});")

conn.execute("CREATE (:person {name: 'Bob', age: 40});")

conn.execute("CREATE (:person {name: 'Cole', age: 20});")

# Basic

result = conn.execute("MATCH (p:person) RETURN p.*")

assert result.get_num_tuples() == 3

rows = result.get_n(2)

assert len(rows) == 2

assert rows[0] == ["Alice", 30]

assert rows[1] == ["Bob", 40]

rows = result.get_n(1)

assert len(rows) == 1

assert rows[0] == ["Cole", 20]

rows = result.get_n(3)

assert len(rows) == 0

# Empty rows on excess count

result = conn.execute("MATCH (p:person) RETURN p.*")

assert result.get_num_tuples() == 3

rows = result.get_n(10)

assert len(rows) == 3

assert rows[0] == ["Alice", 30]

assert rows[1] == ["Bob", 40]

assert rows[2] == ["Cole", 20]