diff --git a/timely_beliefs/tests/test_belief_query.py b/timely_beliefs/tests/test_belief_query.py index 8dc0d346..e51a5d10 100644 --- a/timely_beliefs/tests/test_belief_query.py +++ b/timely_beliefs/tests/test_belief_query.py @@ -6,6 +6,7 @@ import pandas as pd import pytest from pytz import utc +from sqlalchemy import select import timely_beliefs.beliefs.queries as query_utils import timely_beliefs.beliefs.utils as belief_utils @@ -407,8 +408,8 @@ def test_select_most_recent_probabilistic_beliefs( ], ) def test_query_unchanged_beliefs(event_values, expected_unchanged_event_values): - sensor = session.query(DBSensor).first() - source = session.query(DBBeliefSource).first() + sensor = session.execute(select(DBSensor).limit(1)).scalar() + source = session.execute(select(DBBeliefSource).limit(1)).scalar() beliefs = [ DBTimedBelief( sensor=sensor, @@ -435,12 +436,12 @@ def test_query_unchanged_beliefs(event_values, expected_unchanged_event_values): ] ) session.add_all(beliefs) - all_beliefs_query = session.query(DBTimedBelief).filter( + all_beliefs_query = select(DBTimedBelief).filter( DBTimedBelief.sensor == sensor, DBTimedBelief.source == source ) q = query_utils.query_unchanged_beliefs( session=session, query=all_beliefs_query, ) - unchanged_beliefs = BeliefsDataFrame(q.all()) + unchanged_beliefs = BeliefsDataFrame(session.scalars(q).all()) pd.testing.assert_frame_equal(unchanged_beliefs, expected_unchanged_beliefs) diff --git a/timely_beliefs/tests/test_db_subclassing.py b/timely_beliefs/tests/test_db_subclassing.py index dcbf50a2..3fe956b6 100644 --- a/timely_beliefs/tests/test_db_subclassing.py +++ b/timely_beliefs/tests/test_db_subclassing.py @@ -1,7 +1,7 @@ from datetime import datetime, timedelta from pytz import timezone -from sqlalchemy import Column, Float, ForeignKey, Integer +from sqlalchemy import Column, Float, ForeignKey, Integer, func, select from sqlalchemy.ext.declarative import declared_attr from sqlalchemy.orm import backref, relationship @@ -30,14 +30,28 @@ def test_subclassing_source(db): session.add(RatedSource(name="test_source", rating=5)) session.add(RatedSource(name="dummy")) - q = session.query(RatedSource) - print(q.all()) + q = select(RatedSource) + print(session.scalars(q).all()) print(db.tables.keys()) - assert session.query(RatedSource).filter(RatedSource.rating == 5).count() == 1 + assert ( + session.execute( + select(func.count()) + .select_from(RatedSource) + .filter(RatedSource.rating == 5) + ).scalar() + == 1 + ) # We made one with default rating (0) and in conftest three are made in advance - assert session.query(RatedSource).filter(RatedSource.rating == 0).count() == 4 + assert ( + session.execute( + select(func.count()) + .select_from(RatedSource) + .filter(RatedSource.rating == 0) + ).scalar() + == 4 + ) class RatedSourceInCustomTable(Base, BeliefSourceDBMixin): @@ -56,9 +70,11 @@ def test_custom_source_with_mixin(db): session.add(RatedSourceInCustomTable(name="dummy")) assert ( - session.query(RatedSourceInCustomTable) - .filter(RatedSourceInCustomTable.rating == 7) - .count() + session.execute( + select(func.count()) + .select_from(RatedSourceInCustomTable) + .filter(RatedSourceInCustomTable.rating == 7) + ).scalar() == 1 ) @@ -116,17 +132,29 @@ def test_custom_source_and_beliefs_with_mixin(db): ) session.add(belief) - q = session.query(RatedSourceInCustomTable).filter( - RatedSourceInCustomTable.rating == 7 + q = select(RatedSourceInCustomTable).filter(RatedSourceInCustomTable.rating == 7) + assert ( + session.execute( + select(func.count()) + .select_from(RatedSourceInCustomTable) + .filter_by(rating=7) + ).scalar() + == 1 ) - assert q.count() == 1 - assert q.first().rating == 7 + assert session.scalar(q.limit(1)).rating == 7 - q = session.query(JoyfulBeliefInCustomTable).filter( + q = select(JoyfulBeliefInCustomTable).filter( JoyfulBeliefInCustomTable.happiness == 3 ) - assert q.count() == 1 - the_belief = q.first() + assert ( + session.execute( + select(func.count()) + .select_from(JoyfulBeliefInCustomTable) + .filter_by(happiness=3) + ).scalar() + == 1 + ) + the_belief = session.scalar(q.limit(1)) assert the_belief.event_value == belief.event_value assert the_belief.sensor.__class__ == DBSensor assert the_belief.source.__class__ == RatedSourceInCustomTable diff --git a/timely_beliefs/tests/test_sensor_query.py b/timely_beliefs/tests/test_sensor_query.py index 2a3b39b7..105dd86c 100644 --- a/timely_beliefs/tests/test_sensor_query.py +++ b/timely_beliefs/tests/test_sensor_query.py @@ -1,6 +1,8 @@ +from sqlalchemy import select + from timely_beliefs import DBSensor from timely_beliefs.tests import session def test_persist_sensor(db): - assert session.query(DBSensor).first() + assert session.scalar(select(DBSensor).limit(1))