129 lines
3.8 KiB
Python
129 lines
3.8 KiB
Python
import sqlite3
|
|
import tempfile
|
|
|
|
from pathlib import Path
|
|
from sqlite3 import Connection
|
|
from dodgeballtracker.typing import TrackedDetectedObject
|
|
|
|
import pandas as pd
|
|
import pandera as pa
|
|
import numpy as np
|
|
import pytest
|
|
from contextlib import closing, contextmanager
|
|
|
|
from dodgeballtracker.data_dumper import SQLiteDataDumper
|
|
|
|
|
|
def list_tables(conn: Connection) -> set[str]:
|
|
cur = conn.cursor()
|
|
res = cur.execute('''
|
|
SELECT
|
|
name
|
|
FROM
|
|
sqlite_schema
|
|
WHERE
|
|
type ='table' AND
|
|
name NOT LIKE 'sqlite_%';
|
|
''')
|
|
return {row[0] for row in res}
|
|
|
|
@pytest.fixture
|
|
def tmp_dir():
|
|
with tempfile.TemporaryDirectory() as dir_:
|
|
yield Path(dir_)
|
|
|
|
|
|
def number_of_items_in_table(table: str, conn: Connection) -> int:
|
|
cur = conn.cursor()
|
|
res = cur.execute(f'SELECT COUNT(*) FROM {table}')
|
|
return res.fetchone()[0]
|
|
|
|
|
|
# @pytest.fixture
|
|
@contextmanager
|
|
def make_conn(tmp_dir: Path):
|
|
db_file= tmp_dir / 'db.sqlite'
|
|
conn_ = sqlite3.connect(db_file)
|
|
with closing(conn_) as open_conn:
|
|
open_conn.execute('PRAGMA foreign_keys = 1')
|
|
yield open_conn
|
|
|
|
|
|
@pytest.fixture
|
|
@pa.check_output(TrackedDetectedObject)
|
|
def mock_tracked_detected_object():
|
|
rd = np.random.RandomState(seed=109)
|
|
data = rd.rand(10, len(TrackedDetectedObject.columns))
|
|
return pd.DataFrame(data, columns=TrackedDetectedObject.columns)
|
|
|
|
|
|
def test_start(tmp_dir: Path):
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('foo/bar.MOV'), 'aaabbb')
|
|
assert dumper._video_id == 1
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
expt = {
|
|
'tracked_detected_object',
|
|
'video',
|
|
'video_frame'
|
|
}
|
|
assert expt == list_tables(conn)
|
|
|
|
|
|
def test_restart(tmp_dir: Path):
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('foo/bar.MOV'), 'aaabbb')
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('foo/bar.MOV'), 'aaabbb')
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
expt = {
|
|
'tracked_detected_object',
|
|
'video',
|
|
'video_frame'
|
|
}
|
|
assert expt == list_tables(conn)
|
|
|
|
|
|
def test_write_boxes_simple(tmp_dir: Path, mock_tracked_detected_object: pd.DataFrame):
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('/foo/toto.MOV'), 'aabbcc')
|
|
dumper.write_boxes(mock_tracked_detected_object, 40018, 34)
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
assert 10 == number_of_items_in_table('tracked_detected_object', conn)
|
|
|
|
|
|
def test_write_boxes_multiple_times(tmp_dir: Path, mock_tracked_detected_object: pd.DataFrame):
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('/foo/toto.MOV'), 'aabbcc')
|
|
dumper.write_boxes(mock_tracked_detected_object, 40018, 34)
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('/foo/toto.MOV'), 'aabbcc')
|
|
dumper.write_boxes(mock_tracked_detected_object, 40018, 34)
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
assert 20 == number_of_items_in_table('tracked_detected_object', conn)
|
|
|
|
|
|
|
|
def test_write_boxes_with_empty_frame(tmp_dir: Path,
|
|
mock_tracked_detected_object: pd.DataFrame):
|
|
empty_tracked_detected_object = mock_tracked_detected_object[0:0]
|
|
with make_conn(tmp_dir) as conn:
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(Path('/foo/toto.MOV'), 'aabbcc')
|
|
dumper.write_boxes(empty_tracked_detected_object, 40018, 34)
|
|
|
|
with make_conn(tmp_dir) as conn:
|
|
assert 0 == number_of_items_in_table('tracked_detected_object', conn)
|
|
assert 0 == number_of_items_in_table('video_frame', conn) |