Coverage for app/repositories/probe_repository.py: 84%

1from typing import Sequence 

2from uuid import UUID 

3from sqlalchemy import select 

4from sqlalchemy.ext.asyncio import AsyncSession 

5from sqlalchemy.orm import selectinload 

6 

7from app.core.logging import Logger 

8from app.models.Grid import Grid 

9from app.models.Probe import Probe 

10 

11 

12class ProbeRepository: 

13 def __init__(self, session: AsyncSession): 

14 self.session = session 

15 

16 async def is_ready(self) -> bool: 

17 try: 

18 await self.session.execute(select(1)) 

19 return True 

20 except Exception as e: 

21 Logger.error(f"Database readiness check failed: {e}", exc_info=True) 

22 return False 

23 

24 async def setup( 

25 self, 

26 position: Probe, 

27 ) -> Probe: 

28 probe = Probe(x=0, y=0, direction=position.direction) 

29 grid = Grid(x=position.x, y=position.y) 

30 probe.grid = grid 

31 

32 self.session.add(probe) 

33 

34 await self.session.commit() 

35 await self.session.refresh(probe) 

36 

37 return probe 

38 

39 async def save(self, probe: Probe) -> Probe: 

40 merged = await self.session.merge(probe) 

41 await self.session.commit() 

42 await self.session.refresh(merged) 

43 

44 return merged 

45 

46 async def find_by_id(self, id: UUID) -> Probe | None: 

47 statement = ( 

48 select(Probe).options(selectinload(Probe.grid)).where(Probe.id == id) 

49 ) 

50 

51 result = await self.session.execute(statement) 

52 

53 return result.scalar_one_or_none() 

54 

55 async def find_all(self) -> Sequence[Probe]: 

56 probes = await self.session.scalars(select(Probe)) 

57 return probes.all()