diff --git a/tests/test_geo_file.py b/tests/test_geo_file.py index d142edc..17c195f 100644 --- a/tests/test_geo_file.py +++ b/tests/test_geo_file.py @@ -1,16 +1,40 @@ -from dzgui.config.ipdb import get_ipdb -from pathlib import Path import os +import sqlite3 import pytest import tempfile +from pathlib import Path + +from dzgui.const.endpoints import DB_IP +from dzgui.config.ipdb import ( + create_table, + find_download_url, + serialize, + strip_ipv6, + unzip, +) + + +def get_db_len(path: Path) -> int: + conn = sqlite3.connect(path) + c = conn.cursor() + c.execute("SELECT COUNT(*) from ips") + db_len = c.fetchone()[0] + return db_len + + +@pytest.fixture +def csv() -> Path: + csv = Path(tempfile.NamedTemporaryFile().name) + url = find_download_url(DB_IP) + tmp = serialize(url) + unzip(tmp, csv) + strip_ipv6(csv) + return csv @pytest.mark.slow -def test_stripping(): - d = tempfile.TemporaryDirectory() - path = Path(d.name).joinpath("ips.csv") - get_ipdb(path) - with open(path, "rb") as f: +def test_stripping(csv): + with open(csv, "rb") as f: first = f.readline().decode().rstrip() try: f.seek(-2, os.SEEK_END) @@ -25,3 +49,16 @@ def test_stripping(): assert first_els[0] == "0.0.0.0" assert last_els[0] == "224.0.0.0" + + +@pytest.mark.slow +def test_row_parity(csv) -> None: + rows = csv.read_text().splitlines() + + db = Path(tempfile.NamedTemporaryFile().name) + create_table(db, csv) + + csv_len = len(rows) + db_len = get_db_len(db) + + assert csv_len == db_len