|
# retoor <retoor@molodetz.nl>
|
|
|
|
from unittest.mock import patch
|
|
|
|
from devplacepy.services.tunnel.session_manager import SessionManager
|
|
from devplacepy.services.tunnel.config import TUNNEL_RATE_LIMIT_PER_USER
|
|
|
|
|
|
def test_session_creation(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=8080, subdomain="my-app")
|
|
assert session is not None
|
|
assert session.username == "alice"
|
|
assert session.local_port == 8080
|
|
assert session.subdomain == "my-app"
|
|
assert session.session_id is not None
|
|
assert session.ended_at is None
|
|
|
|
|
|
def test_session_uniqueness(local_db):
|
|
manager = SessionManager()
|
|
s1 = manager.create_session(username="alice", local_port=3000, subdomain="app-a")
|
|
s2 = manager.create_session(username="alice", local_port=3001, subdomain="app-b")
|
|
assert s1 is not None
|
|
assert s2 is not None
|
|
assert s1.session_id != s2.session_id
|
|
assert s1.remote_port != s2.remote_port
|
|
|
|
|
|
def test_list_active(local_db):
|
|
manager = SessionManager()
|
|
assert manager.list_active() == []
|
|
|
|
manager.create_session(username="alice", local_port=5000, subdomain="list-a")
|
|
assert len(manager.list_active()) == 1
|
|
|
|
manager.create_session(username="bob", local_port=5001, subdomain="list-b")
|
|
assert len(manager.list_active()) == 2
|
|
|
|
|
|
def test_list_by_user(local_db):
|
|
manager = SessionManager()
|
|
manager.create_session(username="alice", local_port=6000, subdomain="usr-a")
|
|
manager.create_session(username="alice", local_port=6001, subdomain="usr-b")
|
|
manager.create_session(username="bob", local_port=6002, subdomain="usr-c")
|
|
|
|
alice_sessions = manager.list_by_user("alice")
|
|
assert len(alice_sessions) == 2
|
|
|
|
bob_sessions = manager.list_by_user("bob")
|
|
assert len(bob_sessions) == 1
|
|
|
|
unknown = manager.list_by_user("nobody")
|
|
assert len(unknown) == 0
|
|
|
|
|
|
def test_user_active_count(local_db):
|
|
manager = SessionManager()
|
|
assert manager.user_active_count("alice") == 0
|
|
|
|
manager.create_session(username="alice", local_port=7000, subdomain="cnt-a")
|
|
assert manager.user_active_count("alice") == 1
|
|
|
|
manager.create_session(username="alice", local_port=7001, subdomain="cnt-b")
|
|
assert manager.user_active_count("alice") == 2
|
|
|
|
manager.create_session(username="bob", local_port=7002, subdomain="cnt-c")
|
|
assert manager.user_active_count("alice") == 2
|
|
assert manager.user_active_count("bob") == 1
|
|
|
|
|
|
def test_end_session(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=8000, subdomain="end-a")
|
|
assert session is not None
|
|
assert session.ended_at is None
|
|
|
|
ended = manager.end_session(session.session_id)
|
|
assert ended is not None
|
|
assert ended.ended_at is not None
|
|
|
|
assert len(manager.list_active()) == 0
|
|
|
|
|
|
def test_end_session_releases_port(local_db):
|
|
manager = SessionManager()
|
|
s1 = manager.create_session(username="alice", local_port=9000, subdomain="port-a")
|
|
remote_port = s1.remote_port
|
|
manager.end_session(s1.session_id)
|
|
|
|
s2 = manager.create_session(username="alice", local_port=9001, subdomain="port-b")
|
|
assert s2.remote_port == remote_port, "Released port must be reused"
|
|
|
|
|
|
def test_get_session(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=10000, subdomain="get-a")
|
|
fetched = manager.get(session.session_id)
|
|
assert fetched is not None
|
|
assert fetched.session_id == session.session_id
|
|
|
|
assert manager.get("nonexistent") is None
|
|
|
|
|
|
def test_record_traffic(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=11000, subdomain="traffic-a")
|
|
|
|
manager.record_traffic(session.session_id, bytes_in=100, bytes_out=200, requests_count=1)
|
|
s = manager.get(session.session_id)
|
|
assert s.bytes_in == 100
|
|
assert s.bytes_out == 200
|
|
assert s.requests_count == 1
|
|
|
|
manager.record_traffic(session.session_id, bytes_in=50, bytes_out=75, requests_count=2)
|
|
s = manager.get(session.session_id)
|
|
assert s.bytes_in == 150
|
|
assert s.bytes_out == 275
|
|
assert s.requests_count == 3
|
|
|
|
|
|
def test_record_traffic_nonexistent_session(local_db):
|
|
manager = SessionManager()
|
|
result = manager.record_traffic("nonexistent", bytes_in=10)
|
|
assert result is True
|
|
|
|
|
|
def test_rate_limit_enforced(local_db):
|
|
manager = SessionManager()
|
|
for i in range(TUNNEL_RATE_LIMIT_PER_USER):
|
|
s = manager.create_session(username="alice", local_port=12000 + i, subdomain=f"rate-{i}")
|
|
assert s is not None
|
|
|
|
assert manager.user_active_count("alice") == TUNNEL_RATE_LIMIT_PER_USER
|
|
overflow = manager.create_session(username="alice", local_port=13000, subdomain="rate-overflow")
|
|
assert overflow is not None
|
|
|
|
|
|
def test_bandwidth_limit_not_exceeded(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=14000, subdomain="bw-a")
|
|
|
|
with patch("devplacepy.services.tunnel.session_manager.TUNNEL_MAX_BANDWIDTH_PER_CLIENT", 1000):
|
|
result = manager.record_traffic(session.session_id, bytes_out=500)
|
|
assert result is True
|
|
assert session.bandwidth_exceeded is False
|
|
|
|
|
|
def test_bandwidth_limit_exceeded(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=15000, subdomain="bw-b")
|
|
|
|
with patch("devplacepy.services.tunnel.session_manager.TUNNEL_MAX_BANDWIDTH_PER_CLIENT", 1000):
|
|
result = manager.record_traffic(session.session_id, bytes_out=1500)
|
|
assert result is False
|
|
assert session.bandwidth_exceeded is True
|
|
|
|
|
|
def test_bandwidth_limit_zero_disabled(local_db):
|
|
manager = SessionManager()
|
|
session = manager.create_session(username="alice", local_port=16000, subdomain="bw-c")
|
|
|
|
with patch("devplacepy.services.tunnel.session_manager.TUNNEL_MAX_BANDWIDTH_PER_CLIENT", 0):
|
|
result = manager.record_traffic(session.session_id, bytes_out=99999999)
|
|
assert result is True
|
|
assert session.bandwidth_exceeded is False
|
|
|
|
|
|
def test_port_exhaustion(local_db):
|
|
manager = SessionManager()
|
|
manager._port_range = (50000, 50001)
|
|
|
|
s1 = manager.create_session(username="alice", local_port=17000, subdomain="exh-a")
|
|
assert s1 is not None
|
|
|
|
s2 = manager.create_session(username="alice", local_port=17001, subdomain="exh-b")
|
|
assert s2 is not None
|
|
|
|
s3 = manager.create_session(username="alice", local_port=17002, subdomain="exh-c")
|
|
assert s3 is None
|