mirror of
https://github.com/Dispatcharr/Dispatcharr.git
synced 2026-07-18 17:16:26 +00:00
4fc12bca4a
- Improved logging in the Stream model for better debugging of profile evaluations. - Introduced a new method `_stream_assignment_is_reusable` in the Channel model to determine if existing stream assignments can be reused, enhancing efficiency. - Updated the release logic in `release_profile_slot` to utilize stored credential keys, reducing unnecessary database lookups. - Simplified error handling in the `get_stream_info_for_switch` function to ensure proper stream release on exceptions. - Enhanced tests for connection pool management and error handling in the ServerGroupsTable component to improve reliability and user feedback.
594 lines
21 KiB
Python
594 lines
21 KiB
Python
"""Tests for shared ServerGroup connection pools (#1137)."""
|
|
|
|
from django.test import TestCase
|
|
from unittest.mock import patch
|
|
|
|
from apps.m3u.connection_pool import (
|
|
extract_credentials_from_stream_url,
|
|
get_credential_connection_count,
|
|
get_enforced_server_group_for_profile,
|
|
get_profile_connection_count,
|
|
get_profile_credential_fingerprint,
|
|
group_has_capacity_for_profile,
|
|
pool_has_capacity_for_profile,
|
|
profile_has_capacity_for_selection,
|
|
profile_connections_key,
|
|
profile_credential_release_key,
|
|
release_profile_slot,
|
|
reserve_profile_slot,
|
|
server_group_connections_key,
|
|
)
|
|
from apps.m3u.models import M3UAccount, M3UAccountProfile, ServerGroup
|
|
|
|
|
|
class FakeRedis:
|
|
"""Minimal in-memory Redis stand-in for counter tests."""
|
|
|
|
def __init__(self):
|
|
self._data = {}
|
|
|
|
def get(self, key):
|
|
val = self._data.get(key)
|
|
if val is None:
|
|
return None
|
|
if isinstance(val, str):
|
|
return val.encode()
|
|
return str(val).encode()
|
|
|
|
def set(self, key, value, ex=None):
|
|
try:
|
|
self._data[key] = int(value)
|
|
except (ValueError, TypeError):
|
|
self._data[key] = value
|
|
|
|
def incr(self, key):
|
|
self._data[key] = self._data.get(key, 0) + 1
|
|
return self._data[key]
|
|
|
|
def decr(self, key):
|
|
self._data[key] = self._data.get(key, 0) - 1
|
|
return self._data[key]
|
|
|
|
def delete(self, key):
|
|
self._data.pop(key, None)
|
|
|
|
def pipeline(self):
|
|
return FakeRedisPipeline(self)
|
|
|
|
|
|
class FakeRedisPipeline:
|
|
def __init__(self, redis):
|
|
self.redis = redis
|
|
self._ops = []
|
|
|
|
def decr(self, key):
|
|
self._ops.append(("decr", key))
|
|
return self
|
|
|
|
def incr(self, key):
|
|
self._ops.append(("incr", key))
|
|
return self
|
|
|
|
def set(self, key, value):
|
|
self._ops.append(("set", key, value))
|
|
return self
|
|
|
|
def execute(self):
|
|
for op in self._ops:
|
|
if op[0] == "decr":
|
|
self.redis.decr(op[1])
|
|
elif op[0] == "incr":
|
|
self.redis.incr(op[1])
|
|
elif op[0] == "set":
|
|
self.redis.set(op[1], op[2])
|
|
self._ops = []
|
|
|
|
|
|
class ExtractCredentialsTests(TestCase):
|
|
def test_extract_credentials_from_xc_style_url(self):
|
|
url = "http://example.com/live/alice/secret123/99999.ts"
|
|
user, password = extract_credentials_from_stream_url(url)
|
|
self.assertEqual(user, "alice")
|
|
self.assertEqual(password, "secret123")
|
|
|
|
|
|
class ManualServerGroupTests(TestCase):
|
|
def test_group_enforced_when_account_assigned(self):
|
|
group = ServerGroup.objects.create(name="provider-a")
|
|
account = M3UAccount.objects.create(
|
|
name="Account A",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_group=group,
|
|
)
|
|
profile = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
|
|
self.assertEqual(get_enforced_server_group_for_profile(profile), group)
|
|
|
|
def test_accounts_in_same_group_share_credential_counter(self):
|
|
group = ServerGroup.objects.create(name="shared")
|
|
account1 = M3UAccount.objects.create(
|
|
name="XC Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
account2 = M3UAccount.objects.create(
|
|
name="M3U Account",
|
|
account_type="STD",
|
|
username="user",
|
|
password="pass",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
profile1 = M3UAccountProfile.objects.get(m3u_account=account1, is_default=True)
|
|
profile2 = M3UAccountProfile.objects.get(m3u_account=account2, is_default=True)
|
|
profile1.max_streams = 1
|
|
profile1.save()
|
|
profile2.max_streams = 1
|
|
profile2.save()
|
|
|
|
redis = FakeRedis()
|
|
reserved1, _, _ = reserve_profile_slot(profile1, redis)
|
|
self.assertTrue(reserved1)
|
|
|
|
reserved2, _, _ = reserve_profile_slot(profile2, redis)
|
|
self.assertFalse(reserved2)
|
|
self.assertFalse(group_has_capacity_for_profile(profile2, redis))
|
|
|
|
def test_profile_rotation_when_default_profile_full(self):
|
|
account = M3UAccount.objects.create(
|
|
name="Multi-profile",
|
|
account_type="XC",
|
|
max_streams=1,
|
|
)
|
|
default = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
default.max_streams = 1
|
|
default.save()
|
|
|
|
alt = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt_profile",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
redis = FakeRedis()
|
|
reserved, _, _ = reserve_profile_slot(default, redis)
|
|
self.assertTrue(reserved)
|
|
self.assertFalse(profile_has_capacity_for_selection(default, redis))
|
|
self.assertTrue(profile_has_capacity_for_selection(alt, redis))
|
|
|
|
|
|
class PoolEnforcementTests(TestCase):
|
|
def setUp(self):
|
|
self.redis = FakeRedis()
|
|
self.group = ServerGroup.objects.create(name="test-pool")
|
|
self.account = M3UAccount.objects.create(
|
|
name="Test Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=self.group,
|
|
max_streams=5,
|
|
)
|
|
self.profile = M3UAccountProfile.objects.get(
|
|
m3u_account=self.account, is_default=True
|
|
)
|
|
self.profile.max_streams = 1
|
|
self.profile.save()
|
|
|
|
def test_group_has_capacity_reads_credential_counter_directly(self):
|
|
cred_key = server_group_connections_key(
|
|
self.group.id,
|
|
get_profile_credential_fingerprint(self.profile),
|
|
)
|
|
self.redis.set(cred_key, 1)
|
|
|
|
self.assertFalse(group_has_capacity_for_profile(self.profile, self.redis))
|
|
|
|
self.redis.set(cred_key, 0)
|
|
self.assertTrue(group_has_capacity_for_profile(self.profile, self.redis))
|
|
|
|
def test_reserve_and_release_both_counters(self):
|
|
reserved, count, _ = reserve_profile_slot(self.profile, self.redis)
|
|
self.assertTrue(reserved)
|
|
self.assertEqual(count, 1)
|
|
|
|
cred_key = server_group_connections_key(
|
|
self.group.id,
|
|
get_profile_credential_fingerprint(self.profile),
|
|
)
|
|
profile_key = profile_connections_key(self.profile.id)
|
|
self.assertEqual(self.redis._data[cred_key], 1)
|
|
self.assertEqual(self.redis._data[profile_key], 1)
|
|
|
|
release_profile_slot(self.profile.id, self.redis)
|
|
self.assertEqual(self.redis._data[cred_key], 0)
|
|
self.assertEqual(self.redis._data[profile_key], 0)
|
|
|
|
def test_same_credential_capped_at_profile_max(self):
|
|
"""Shared credential counter is capped by each profile's max_streams."""
|
|
account2 = M3UAccount.objects.create(
|
|
name="Second Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=self.group,
|
|
max_streams=5,
|
|
)
|
|
profile2 = M3UAccountProfile.objects.get(m3u_account=account2, is_default=True)
|
|
profile2.max_streams = 1
|
|
profile2.save()
|
|
|
|
self.assertTrue(reserve_profile_slot(self.profile, self.redis)[0])
|
|
self.assertFalse(reserve_profile_slot(profile2, self.redis)[0])
|
|
|
|
fp = get_profile_credential_fingerprint(self.profile)
|
|
cred_key = server_group_connections_key(self.group.id, fp)
|
|
self.assertEqual(self.redis._data[cred_key], 1)
|
|
|
|
def test_different_logins_both_stream_in_same_group(self):
|
|
"""Different provider logins keep separate credential counters in one group."""
|
|
account = M3UAccount.objects.create(
|
|
name="Grouped multi-login",
|
|
account_type="XC",
|
|
username="login_a",
|
|
password="pass_a",
|
|
server_group=self.group,
|
|
max_streams=5,
|
|
)
|
|
default = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
default.max_streams = 1
|
|
default.save()
|
|
|
|
alt = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt_login",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
fp_a = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
|
fp_b = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
|
|
|
|
with patch(
|
|
"apps.m3u.connection_pool.get_profile_credential_fingerprint",
|
|
side_effect=lambda profile: fp_a if profile.id == default.id else fp_b,
|
|
):
|
|
self.assertTrue(reserve_profile_slot(default, self.redis)[0])
|
|
self.assertTrue(reserve_profile_slot(alt, self.redis)[0])
|
|
|
|
key_a = server_group_connections_key(self.group.id, fp_a)
|
|
key_b = server_group_connections_key(self.group.id, fp_b)
|
|
self.assertEqual(self.redis._data[key_a], 1)
|
|
self.assertEqual(self.redis._data[key_b], 1)
|
|
|
|
def test_no_fingerprint_skips_credential_counter(self):
|
|
account = M3UAccount.objects.create(
|
|
name="No creds",
|
|
account_type="STD",
|
|
server_group=self.group,
|
|
max_streams=5,
|
|
)
|
|
profile = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
profile.max_streams = 1
|
|
profile.save()
|
|
|
|
with patch(
|
|
"apps.m3u.connection_pool.get_profile_credential_fingerprint",
|
|
return_value=None,
|
|
):
|
|
self.assertTrue(reserve_profile_slot(profile, self.redis)[0])
|
|
self.assertEqual(get_credential_connection_count(profile, self.redis), 0)
|
|
self.assertEqual(get_credential_connection_count(profile, self.redis), 0)
|
|
|
|
def test_release_when_profile_row_deleted(self):
|
|
profile_id = self.profile.id
|
|
fp = get_profile_credential_fingerprint(self.profile)
|
|
cred_key = server_group_connections_key(self.group.id, fp)
|
|
|
|
reserved, _, failure_reason = reserve_profile_slot(
|
|
self.profile, self.redis
|
|
)
|
|
self.assertTrue(reserved)
|
|
self.assertIsNone(failure_reason)
|
|
self.profile.delete()
|
|
|
|
release_profile_slot(profile_id, self.redis)
|
|
|
|
self.assertEqual(self.redis._data[profile_connections_key(profile_id)], 0)
|
|
self.assertEqual(self.redis._data[cred_key], 0)
|
|
self.assertNotIn(profile_credential_release_key(profile_id), self.redis._data)
|
|
|
|
def test_release_uses_stored_credential_key_without_db_lookup(self):
|
|
profile_id = self.profile.id
|
|
cred_key = server_group_connections_key(
|
|
self.group.id,
|
|
get_profile_credential_fingerprint(self.profile),
|
|
)
|
|
|
|
self.assertTrue(reserve_profile_slot(self.profile, self.redis)[0])
|
|
|
|
with patch("apps.m3u.models.M3UAccountProfile.objects.get") as mock_get:
|
|
release_profile_slot(profile_id, self.redis)
|
|
mock_get.assert_not_called()
|
|
|
|
self.assertEqual(self.redis._data[profile_connections_key(profile_id)], 0)
|
|
self.assertEqual(self.redis._data[cred_key], 0)
|
|
|
|
def test_reserve_returns_failure_reason_without_extra_checks(self):
|
|
account2 = M3UAccount.objects.create(
|
|
name="Reason Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=self.group,
|
|
max_streams=5,
|
|
)
|
|
profile2 = M3UAccountProfile.objects.get(m3u_account=account2, is_default=True)
|
|
profile2.max_streams = 1
|
|
profile2.save()
|
|
|
|
self.assertTrue(reserve_profile_slot(self.profile, self.redis)[0])
|
|
reserved, _count, reason = reserve_profile_slot(profile2, self.redis)
|
|
self.assertFalse(reserved)
|
|
self.assertEqual(reason, "credential_full")
|
|
|
|
|
|
class StaleAssignmentTests(TestCase):
|
|
def test_stale_assignment_releases_counters(self):
|
|
from apps.channels.models import Channel, Stream
|
|
|
|
redis = FakeRedis()
|
|
group = ServerGroup.objects.create(name="stale-group")
|
|
account = M3UAccount.objects.create(
|
|
name="Stale Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
profile = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
profile.max_streams = 1
|
|
profile.save()
|
|
|
|
stream = Stream.objects.create(name="Stale Stream", m3u_account=account)
|
|
channel = Channel.objects.create(channel_number=502, name="Stale Channel")
|
|
channel.streams.add(stream)
|
|
|
|
reserve_profile_slot(profile, redis)
|
|
redis.set(f"channel_stream:{channel.id}", stream.id)
|
|
redis.set(f"stream_profile:{stream.id}", profile.id)
|
|
|
|
profile_key = profile_connections_key(profile.id)
|
|
cred_key = server_group_connections_key(
|
|
group.id, get_profile_credential_fingerprint(profile)
|
|
)
|
|
self.assertEqual(redis._data[profile_key], 1)
|
|
self.assertEqual(redis._data[cred_key], 1)
|
|
|
|
with patch("core.utils.RedisClient.get_client", return_value=redis):
|
|
channel._release_stale_stream_assignment(redis, stream.id)
|
|
|
|
self.assertNotIn(f"channel_stream:{channel.id}", redis._data)
|
|
self.assertNotIn(f"stream_profile:{stream.id}", redis._data)
|
|
self.assertEqual(redis._data[profile_key], 0)
|
|
self.assertEqual(redis._data[cred_key], 0)
|
|
|
|
|
|
class UpdateStreamProfileTests(TestCase):
|
|
def test_switch_updates_profile_counters_when_group_assigned(self):
|
|
from apps.channels.models import Channel, Stream
|
|
|
|
redis = FakeRedis()
|
|
group = ServerGroup.objects.create(name="switch-group")
|
|
account = M3UAccount.objects.create(
|
|
name="Switch Account",
|
|
account_type="XC",
|
|
username="user",
|
|
password="pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
profile_a = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
profile_a.max_streams = 1
|
|
profile_a.save()
|
|
profile_b = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
stream = Stream.objects.create(name="Test Stream", m3u_account=account)
|
|
channel = Channel.objects.create(channel_number=501, name="Switch Channel")
|
|
channel.streams.add(stream)
|
|
|
|
reserve_profile_slot(profile_a, redis)
|
|
redis.set(f"channel_stream:{channel.id}", stream.id)
|
|
redis.set(f"stream_profile:{stream.id}", profile_a.id)
|
|
|
|
cred_key = server_group_connections_key(
|
|
group.id, get_profile_credential_fingerprint(profile_a)
|
|
)
|
|
cred_before = redis._data[cred_key]
|
|
|
|
with patch("core.utils.RedisClient.get_client", return_value=redis):
|
|
self.assertTrue(channel.update_stream_profile(profile_b.id))
|
|
|
|
self.assertEqual(int(redis.get(f"stream_profile:{stream.id}")), profile_b.id)
|
|
self.assertEqual(redis._data[profile_connections_key(profile_a.id)], 0)
|
|
self.assertEqual(redis._data[profile_connections_key(profile_b.id)], 1)
|
|
self.assertEqual(redis._data[cred_key], cred_before)
|
|
|
|
def test_switch_moves_credential_counter_when_login_changes(self):
|
|
from apps.channels.models import Channel, Stream
|
|
|
|
redis = FakeRedis()
|
|
group = ServerGroup.objects.create(name="cred-switch-group")
|
|
account = M3UAccount.objects.create(
|
|
name="Cred Switch Account",
|
|
account_type="XC",
|
|
username="login_a",
|
|
password="pass_a",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
profile_a = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
profile_a.max_streams = 1
|
|
profile_a.save()
|
|
profile_b = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt_login",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
fp_a = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
|
fp_b = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
|
|
key_a = server_group_connections_key(group.id, fp_a)
|
|
key_b = server_group_connections_key(group.id, fp_b)
|
|
|
|
stream = Stream.objects.create(name="Cred Switch Stream", m3u_account=account)
|
|
channel = Channel.objects.create(channel_number=503, name="Cred Switch Channel")
|
|
channel.streams.add(stream)
|
|
|
|
with patch(
|
|
"apps.m3u.connection_pool.get_profile_credential_fingerprint",
|
|
side_effect=lambda profile: fp_a if profile.id == profile_a.id else fp_b,
|
|
):
|
|
reserve_profile_slot(profile_a, redis)
|
|
redis.set(f"channel_stream:{channel.id}", stream.id)
|
|
redis.set(f"stream_profile:{stream.id}", profile_a.id)
|
|
self.assertEqual(redis._data[key_a], 1)
|
|
self.assertNotIn(key_b, redis._data)
|
|
|
|
with patch("core.utils.RedisClient.get_client", return_value=redis):
|
|
self.assertTrue(channel.update_stream_profile(profile_b.id))
|
|
|
|
self.assertEqual(redis._data[key_a], 0)
|
|
self.assertEqual(redis._data[key_b], 1)
|
|
|
|
|
|
class VodProfileSelectionTests(TestCase):
|
|
def test_get_m3u_profile_skips_default_when_profile_full(self):
|
|
from apps.proxy.vod_proxy.views import _get_m3u_profile
|
|
|
|
account = M3UAccount.objects.create(
|
|
name="VOD multi-profile",
|
|
account_type="XC",
|
|
username="xc_user_a",
|
|
password="xc_pass_a",
|
|
server_url="http://xc.example.com",
|
|
max_streams=1,
|
|
)
|
|
default = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
default.max_streams = 1
|
|
default.save()
|
|
|
|
alt = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt_profile",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
redis = FakeRedis()
|
|
self.assertTrue(reserve_profile_slot(default, redis)[0])
|
|
|
|
with patch("core.utils.RedisClient.get_client", return_value=redis):
|
|
result = _get_m3u_profile(account, None, None)
|
|
|
|
self.assertIsNotNone(result)
|
|
selected, _connections = result
|
|
self.assertEqual(selected.id, alt.id)
|
|
|
|
def test_get_m3u_profile_skips_default_when_credential_pool_full(self):
|
|
from apps.proxy.vod_proxy.views import _get_m3u_profile
|
|
|
|
group = ServerGroup.objects.create(name="vod-cred-pool")
|
|
account = M3UAccount.objects.create(
|
|
name="VOD pooled",
|
|
account_type="XC",
|
|
username="shared_user",
|
|
password="shared_pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
default = M3UAccountProfile.objects.get(m3u_account=account, is_default=True)
|
|
default.max_streams = 1
|
|
default.save()
|
|
|
|
alt = M3UAccountProfile.objects.create(
|
|
m3u_account=account,
|
|
name="alt_login",
|
|
is_default=False,
|
|
is_active=True,
|
|
max_streams=1,
|
|
search_pattern="",
|
|
replace_pattern="",
|
|
)
|
|
|
|
other_account = M3UAccount.objects.create(
|
|
name="Other pooled account",
|
|
account_type="XC",
|
|
username="shared_user",
|
|
password="shared_pass",
|
|
server_url="http://xc.example.com",
|
|
server_group=group,
|
|
max_streams=5,
|
|
)
|
|
other_profile = M3UAccountProfile.objects.get(
|
|
m3u_account=other_account, is_default=True
|
|
)
|
|
other_profile.max_streams = 1
|
|
other_profile.save()
|
|
|
|
fp_shared = "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"
|
|
fp_alt = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd"
|
|
|
|
redis = FakeRedis()
|
|
with patch(
|
|
"apps.m3u.connection_pool.get_profile_credential_fingerprint",
|
|
side_effect=lambda profile: (
|
|
fp_shared
|
|
if profile.id in (default.id, other_profile.id)
|
|
else fp_alt
|
|
),
|
|
):
|
|
self.assertTrue(reserve_profile_slot(other_profile, redis)[0])
|
|
|
|
with patch("core.utils.RedisClient.get_client", return_value=redis):
|
|
result = _get_m3u_profile(account, None, None)
|
|
|
|
self.assertIsNotNone(result)
|
|
selected, _connections = result
|
|
self.assertEqual(selected.id, alt.id)
|