126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
|
|
from crawlee.sessions._cookies import SessionCookies
|
|
from crawlee.sessions._session import Session
|
|
|
|
|
|
@pytest.fixture
|
|
def session() -> Session:
|
|
return Session(
|
|
id='test_session',
|
|
max_age=timedelta(minutes=30),
|
|
user_data={'user_key': 'user_value'},
|
|
max_error_score=3.0,
|
|
error_score_decrement=0.5,
|
|
created_at=datetime.now(timezone.utc),
|
|
usage_count=0,
|
|
max_usage_count=10,
|
|
error_score=0.0,
|
|
cookies={'cookie_key': 'cookie_value'},
|
|
blocked_status_codes=[401, 403, 429],
|
|
)
|
|
|
|
|
|
def test_session_init(session: Session) -> None:
|
|
"""Verify that the session initializes correctly with the expected properties."""
|
|
assert session.id == 'test_session'
|
|
assert session.user_data == {'user_key': 'user_value'}
|
|
assert session.cookies == SessionCookies({'cookie_key': 'cookie_value'})
|
|
assert session.expires_at >= datetime.now(timezone.utc)
|
|
assert not session.is_blocked
|
|
assert not session.is_expired
|
|
assert not session.is_max_usage_count_reached
|
|
assert session.is_usable
|
|
|
|
|
|
def test_session_get_state(session: Session) -> None:
|
|
"""Check if the session state is correctly retrievable in both dict and model forms."""
|
|
session_state_dict = session.get_state(as_dict=True)
|
|
assert session_state_dict['id'] == 'test_session'
|
|
|
|
session_state_model = session.get_state(as_dict=False)
|
|
assert session_state_model.id == 'test_session'
|
|
|
|
session_2 = Session.from_model(session_state_model)
|
|
assert session_2.id == 'test_session'
|
|
|
|
|
|
def test_mark_good(session: Session) -> None:
|
|
"""Test the mark_good method increases usage count and potentially decreases error score."""
|
|
initial_usage_count = session.usage_count
|
|
session.mark_good()
|
|
assert session.usage_count == initial_usage_count + 1
|
|
assert session.error_score == 0
|
|
|
|
|
|
def test_mark_bad(session: Session) -> None:
|
|
"""Test the mark_bad method affects the session's error score and usage."""
|
|
initial_error_score = session.error_score
|
|
session.mark_bad()
|
|
assert session.error_score == initial_error_score + 1
|
|
|
|
|
|
def test_multiple_marks(session: Session) -> None:
|
|
"""Test the mark_good and mark_bad methods in sequence."""
|
|
initial_usage_count = session.usage_count
|
|
session.mark_bad()
|
|
session.mark_bad()
|
|
assert session.error_score == initial_usage_count + 2
|
|
session.mark_good()
|
|
session.mark_good()
|
|
assert session.error_score == initial_usage_count + 1
|
|
session.mark_bad()
|
|
session.mark_bad()
|
|
session.mark_good()
|
|
assert session.is_blocked
|
|
assert not session.is_usable
|
|
|
|
|
|
def test_retire_method(session: Session) -> None:
|
|
"""Test that retire method properly sets the session as unusable."""
|
|
initial_usage_count = session.usage_count
|
|
session.retire()
|
|
assert not session.is_usable
|
|
assert session.error_score == 3.0
|
|
assert session.usage_count == initial_usage_count
|
|
|
|
|
|
def test_mark_good_at_usage_limit_no_double_increment() -> None:
|
|
"""Test that mark_good at max usage count does not double-increment usage_count via retire."""
|
|
session = Session(max_usage_count=5, usage_count=4)
|
|
session.mark_good()
|
|
assert session.usage_count == 5
|
|
assert not session.is_usable
|
|
|
|
|
|
def test_mark_bad_at_usage_limit_no_double_increment() -> None:
|
|
"""Test that mark_bad at max usage count does not double-increment usage_count via retire."""
|
|
session = Session(max_usage_count=5, usage_count=4)
|
|
session.mark_bad()
|
|
assert session.usage_count == 5
|
|
assert not session.is_usable
|
|
|
|
|
|
def test_retire_on_blocked_status_code(session: Session) -> None:
|
|
"""Test retiring the session based on specific HTTP status codes."""
|
|
status_code = 403
|
|
result = session.is_blocked_status_code(status_code=status_code)
|
|
assert result is True
|
|
|
|
|
|
def test_not_retire_on_not_block_status_code(session: Session) -> None:
|
|
"""Test that the session is not retired on a non-blocked status code."""
|
|
status_code = 200
|
|
result = session.is_blocked_status_code(status_code=status_code)
|
|
assert result is False
|
|
|
|
|
|
def test_session_expiration() -> None:
|
|
"""Test the expiration logic of the session."""
|
|
session = Session(created_at=datetime.now(timezone.utc) - timedelta(hours=1))
|
|
assert session.is_expired
|