102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
"""Tests for pet and wildlife query functions."""
|
|
|
|
import time
|
|
|
|
from vigilar.storage.queries import (
|
|
insert_pet,
|
|
get_pet,
|
|
get_all_pets,
|
|
insert_pet_sighting,
|
|
get_pet_sightings,
|
|
get_pet_last_location,
|
|
insert_wildlife_sighting,
|
|
get_wildlife_sightings,
|
|
insert_training_image,
|
|
get_training_images,
|
|
get_unlabeled_sightings,
|
|
label_sighting,
|
|
)
|
|
|
|
|
|
class TestPetCRUD:
|
|
def test_insert_and_get_pet(self, test_db):
|
|
pet_id = insert_pet(test_db, name="Angel", species="cat", breed="DSH",
|
|
color_description="black")
|
|
pet = get_pet(test_db, pet_id)
|
|
assert pet is not None
|
|
assert pet["name"] == "Angel"
|
|
assert pet["species"] == "cat"
|
|
assert pet["training_count"] == 0
|
|
|
|
def test_get_all_pets(self, test_db):
|
|
insert_pet(test_db, name="Angel", species="cat")
|
|
insert_pet(test_db, name="Milo", species="dog")
|
|
all_pets = get_all_pets(test_db)
|
|
assert len(all_pets) == 2
|
|
|
|
|
|
class TestPetSightings:
|
|
def test_insert_and_query(self, test_db):
|
|
pet_id = insert_pet(test_db, name="Angel", species="cat")
|
|
insert_pet_sighting(test_db, pet_id=pet_id, species="cat",
|
|
camera_id="kitchen", confidence=0.92)
|
|
sightings = get_pet_sightings(test_db, limit=10)
|
|
assert len(sightings) == 1
|
|
assert sightings[0]["camera_id"] == "kitchen"
|
|
|
|
def test_last_location(self, test_db):
|
|
pet_id = insert_pet(test_db, name="Angel", species="cat")
|
|
insert_pet_sighting(test_db, pet_id=pet_id, species="cat",
|
|
camera_id="kitchen", confidence=0.9)
|
|
insert_pet_sighting(test_db, pet_id=pet_id, species="cat",
|
|
camera_id="living_room", confidence=0.95)
|
|
loc = get_pet_last_location(test_db, pet_id)
|
|
assert loc is not None
|
|
assert loc["camera_id"] == "living_room"
|
|
|
|
def test_unlabeled_sightings(self, test_db):
|
|
insert_pet_sighting(test_db, pet_id=None, species="cat",
|
|
camera_id="kitchen", confidence=0.6, crop_path="/tmp/crop.jpg")
|
|
unlabeled = get_unlabeled_sightings(test_db, limit=10)
|
|
assert len(unlabeled) == 1
|
|
assert unlabeled[0]["labeled"] == 0
|
|
|
|
def test_label_sighting(self, test_db):
|
|
pet_id = insert_pet(test_db, name="Angel", species="cat")
|
|
insert_pet_sighting(test_db, pet_id=None, species="cat",
|
|
camera_id="kitchen", confidence=0.6)
|
|
sightings = get_unlabeled_sightings(test_db, limit=10)
|
|
sighting_id = sightings[0]["id"]
|
|
label_sighting(test_db, sighting_id, pet_id)
|
|
updated = get_pet_sightings(test_db, pet_id=pet_id)
|
|
assert len(updated) == 1
|
|
assert updated[0]["labeled"] == 1
|
|
|
|
|
|
class TestWildlifeSightings:
|
|
def test_insert_and_query(self, test_db):
|
|
insert_wildlife_sighting(test_db, species="bear", threat_level="PREDATOR",
|
|
camera_id="front", confidence=0.88)
|
|
sightings = get_wildlife_sightings(test_db, limit=10)
|
|
assert len(sightings) == 1
|
|
assert sightings[0]["threat_level"] == "PREDATOR"
|
|
|
|
def test_filter_by_threat(self, test_db):
|
|
insert_wildlife_sighting(test_db, species="bear", threat_level="PREDATOR",
|
|
camera_id="front", confidence=0.88)
|
|
insert_wildlife_sighting(test_db, species="deer", threat_level="PASSIVE",
|
|
camera_id="back", confidence=0.75)
|
|
predators = get_wildlife_sightings(test_db, threat_level="PREDATOR")
|
|
assert len(predators) == 1
|
|
|
|
|
|
class TestTrainingImages:
|
|
def test_insert_and_query(self, test_db):
|
|
pet_id = insert_pet(test_db, name="Angel", species="cat")
|
|
insert_training_image(test_db, pet_id=pet_id,
|
|
image_path="/var/vigilar/pets/training/angel/001.jpg",
|
|
source="upload")
|
|
images = get_training_images(test_db, pet_id)
|
|
assert len(images) == 1
|
|
assert images[0]["source"] == "upload"
|