703 lines
21 KiB
Python
703 lines
21 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
import ray
|
|
from ray import serve
|
|
from ray._common.constants import HEAD_NODE_RESOURCE_NAME
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.serve._private.constants import SERVE_NAMESPACE
|
|
from ray.serve._private.utils import (
|
|
Semaphore,
|
|
calculate_remaining_timeout,
|
|
get_active_placement_group_ids,
|
|
get_all_live_placement_group_names,
|
|
get_current_actor_id,
|
|
get_head_node_id,
|
|
is_running_in_asyncio_loop,
|
|
merge_dict,
|
|
msgpack_deserialize,
|
|
msgpack_serialize,
|
|
override_runtime_envs_except_env_vars,
|
|
serve_encoders,
|
|
snake_to_camel_case,
|
|
)
|
|
from ray.serve.tests.common.remote_uris import (
|
|
TEST_DAG_PINNED_URI,
|
|
TEST_DEPLOY_GROUP_PINNED_URI,
|
|
)
|
|
|
|
|
|
def test_serialize():
|
|
data = msgpack_serialize(5)
|
|
obj = msgpack_deserialize(data)
|
|
assert 5 == obj
|
|
|
|
|
|
def test_merge_dict():
|
|
dict1 = {"pending": 1, "running": 1, "finished": 1}
|
|
dict2 = {"pending": 4, "finished": 1}
|
|
merge = merge_dict(dict1, dict2)
|
|
assert merge["pending"] == 5
|
|
assert merge["running"] == 1
|
|
assert merge["finished"] == 2
|
|
dict1 = None
|
|
merge = merge_dict(dict1, dict2)
|
|
assert merge["pending"] == 4
|
|
assert merge["finished"] == 1
|
|
try:
|
|
assert merge["running"] == 1
|
|
assert False
|
|
except KeyError:
|
|
assert True
|
|
dict2 = None
|
|
merge = merge_dict(dict1, dict2)
|
|
assert merge is None
|
|
|
|
|
|
def test_bytes_encoder():
|
|
data_before = {"inp": {"nest": b"bytes"}}
|
|
data_after = {"inp": {"nest": "bytes"}}
|
|
assert json.loads(json.dumps(jsonable_encoder(data_before))) == data_after
|
|
|
|
|
|
def test_numpy_encoding():
|
|
data = [1, 2]
|
|
floats = np.array(data).astype(np.float32)
|
|
ints = floats.astype(np.int32)
|
|
uints = floats.astype(np.uint32)
|
|
list_of_uints = [np.int64(1), np.int64(2)]
|
|
|
|
for np_data in [floats, ints, uints, list_of_uints]:
|
|
assert (
|
|
json.loads(
|
|
json.dumps(jsonable_encoder(np_data, custom_encoder=serve_encoders))
|
|
)
|
|
== data
|
|
)
|
|
nested = {"a": np.array([1, 2])}
|
|
assert json.loads(
|
|
json.dumps(jsonable_encoder(nested, custom_encoder=serve_encoders))
|
|
) == {"a": [1, 2]}
|
|
|
|
|
|
@serve.deployment
|
|
def decorated_f(*args):
|
|
return "reached decorated_f"
|
|
|
|
|
|
@serve.deployment
|
|
class DecoratedActor:
|
|
def __call__(self, *args):
|
|
return "reached decorated_actor"
|
|
|
|
|
|
def gen_func():
|
|
@serve.deployment
|
|
def f():
|
|
pass
|
|
|
|
return f
|
|
|
|
|
|
def gen_class():
|
|
@serve.deployment
|
|
class A:
|
|
pass
|
|
|
|
return A
|
|
|
|
|
|
class TestOverrideRuntimeEnvsExceptEnvVars:
|
|
def test_merge_empty(self):
|
|
assert {"env_vars": {}} == override_runtime_envs_except_env_vars({}, {})
|
|
|
|
def test_merge_empty_parent(self):
|
|
child = {"env_vars": {"test1": "test_val"}, "working_dir": "."}
|
|
assert child == override_runtime_envs_except_env_vars({}, child)
|
|
|
|
def test_merge_empty_child(self):
|
|
parent = {"env_vars": {"test1": "test_val"}, "working_dir": "."}
|
|
assert parent == override_runtime_envs_except_env_vars(parent, {})
|
|
|
|
@pytest.mark.parametrize("invalid_env", [None, 0, "runtime_env", set()])
|
|
def test_invalid_type(self, invalid_env):
|
|
with pytest.raises(TypeError):
|
|
override_runtime_envs_except_env_vars(invalid_env, {})
|
|
with pytest.raises(TypeError):
|
|
override_runtime_envs_except_env_vars({}, invalid_env)
|
|
with pytest.raises(TypeError):
|
|
override_runtime_envs_except_env_vars(invalid_env, invalid_env)
|
|
|
|
def test_basic_merge(self):
|
|
parent = {
|
|
"py_modules": ["http://test.com/test0.zip", "s3://path/test1.zip"],
|
|
"working_dir": "gs://path/test2.zip",
|
|
"env_vars": {"test": "val", "trial": "val2"},
|
|
"pip": ["pandas", "numpy"],
|
|
"excludes": ["my_file.txt"],
|
|
}
|
|
original_parent = parent.copy()
|
|
|
|
child = {
|
|
"py_modules": [],
|
|
"working_dir": "s3://path/test1.zip",
|
|
"env_vars": {"test": "val", "trial": "val2"},
|
|
"pip": ["numpy"],
|
|
}
|
|
original_child = child.copy()
|
|
|
|
merged = override_runtime_envs_except_env_vars(parent, child)
|
|
|
|
assert original_parent == parent
|
|
assert original_child == child
|
|
assert merged == {
|
|
"py_modules": [],
|
|
"working_dir": "s3://path/test1.zip",
|
|
"env_vars": {"test": "val", "trial": "val2"},
|
|
"pip": ["numpy"],
|
|
"excludes": ["my_file.txt"],
|
|
}
|
|
|
|
def test_merge_deep_copy(self):
|
|
"""Check that the env values are actually deep-copied."""
|
|
|
|
parent_env_vars = {"parent": "pval"}
|
|
child_env_vars = {"child": "cval"}
|
|
|
|
parent = {"env_vars": parent_env_vars}
|
|
child = {"env_vars": child_env_vars}
|
|
original_parent = parent.copy()
|
|
original_child = child.copy()
|
|
|
|
merged = override_runtime_envs_except_env_vars(parent, child)
|
|
assert merged["env_vars"] == {"parent": "pval", "child": "cval"}
|
|
assert original_parent == parent
|
|
assert original_child == child
|
|
|
|
def test_merge_empty_env_vars(self):
|
|
env_vars = {"test": "val", "trial": "val2"}
|
|
|
|
non_empty = {"env_vars": {"test": "val", "trial": "val2"}}
|
|
empty = {}
|
|
|
|
assert (
|
|
env_vars
|
|
== override_runtime_envs_except_env_vars(non_empty, empty)["env_vars"]
|
|
)
|
|
assert (
|
|
env_vars
|
|
== override_runtime_envs_except_env_vars(empty, non_empty)["env_vars"]
|
|
)
|
|
assert {} == override_runtime_envs_except_env_vars(empty, empty)["env_vars"]
|
|
|
|
def test_merge_env_vars(self):
|
|
parent = {
|
|
"py_modules": ["http://test.com/test0.zip", "s3://path/test1.zip"],
|
|
"working_dir": "gs://path/test2.zip",
|
|
"env_vars": {"parent": "pval", "override": "old"},
|
|
"pip": ["pandas", "numpy"],
|
|
"excludes": ["my_file.txt"],
|
|
}
|
|
|
|
child = {
|
|
"py_modules": [],
|
|
"working_dir": "s3://path/test1.zip",
|
|
"env_vars": {"child": "cval", "override": "new"},
|
|
"pip": ["numpy"],
|
|
}
|
|
|
|
merged = override_runtime_envs_except_env_vars(parent, child)
|
|
assert merged == {
|
|
"py_modules": [],
|
|
"working_dir": "s3://path/test1.zip",
|
|
"env_vars": {"parent": "pval", "child": "cval", "override": "new"},
|
|
"pip": ["numpy"],
|
|
"excludes": ["my_file.txt"],
|
|
}
|
|
|
|
def test_inheritance_regression(self):
|
|
"""Check if the general Ray runtime_env inheritance behavior matches.
|
|
|
|
override_runtime_envs_except_env_vars should match the general Ray
|
|
runtime_env inheritance behavior. This test checks if that behavior
|
|
has changed, which would indicate a regression in
|
|
override_runtime_envs_except_env_vars. If the runtime_env inheritance
|
|
behavior changes, override_runtime_envs_except_env_vars should also
|
|
change to match.
|
|
"""
|
|
|
|
with ray.init(
|
|
runtime_env={
|
|
"py_modules": [TEST_DAG_PINNED_URI],
|
|
"env_vars": {"var1": "hello"},
|
|
}
|
|
):
|
|
|
|
@ray.remote
|
|
def check_module():
|
|
# Check that Ray job's py_module loaded correctly
|
|
from conditional_dag import serve_dag # noqa: F401
|
|
|
|
return os.getenv("var1")
|
|
|
|
assert ray.get(check_module.remote()) == "hello"
|
|
|
|
@ray.remote(
|
|
runtime_env={
|
|
"py_modules": [TEST_DEPLOY_GROUP_PINNED_URI],
|
|
"env_vars": {"var2": "world"},
|
|
}
|
|
)
|
|
def test_task():
|
|
with pytest.raises(ImportError):
|
|
# Check that Ray job's py_module was overwritten
|
|
from conditional_dag import serve_dag # noqa: F401
|
|
|
|
from test_env.shallow_import import ShallowClass
|
|
|
|
if ShallowClass()() == "Hello shallow world!":
|
|
return os.getenv("var1") + " " + os.getenv("var2")
|
|
|
|
assert ray.get(test_task.remote()) == "hello world"
|
|
|
|
|
|
class TestSnakeToCamelCase:
|
|
def test_empty(self):
|
|
assert "" == snake_to_camel_case("")
|
|
|
|
def test_single_word(self):
|
|
assert snake_to_camel_case("oneword") == "oneword"
|
|
|
|
def test_multiple_words(self):
|
|
assert (
|
|
snake_to_camel_case("there_are_multiple_words_in_this_phrase")
|
|
== "thereAreMultipleWordsInThisPhrase"
|
|
)
|
|
|
|
def test_single_char_words(self):
|
|
assert snake_to_camel_case("this_is_a_test") == "thisIsATest"
|
|
|
|
def test_leading_capitalization(self):
|
|
"""If the leading character is already capitalized, leave it capitalized."""
|
|
|
|
assert snake_to_camel_case("Leading_cap") == "LeadingCap"
|
|
|
|
def test_leading_alphanumeric(self):
|
|
assert (
|
|
snake_to_camel_case("check_@lphanum3ric_©har_behavior")
|
|
== "check@lphanum3ric©harBehavior"
|
|
)
|
|
|
|
def test_embedded_capitalization(self):
|
|
assert snake_to_camel_case("check_eMbEDDed_caPs") == "checkEMbEDDedCaPs"
|
|
|
|
def test_mixed_caps_alphanumeric(self):
|
|
assert (
|
|
snake_to_camel_case("check_3Mb3DD*d_©a!s_behAvior_Here_wIth_MIxed_cAPs")
|
|
== "check3Mb3DD*d©a!sBehAviorHereWIthMIxedCAPs"
|
|
)
|
|
|
|
def test_leading_underscore(self):
|
|
"""Should strip leading underscores."""
|
|
|
|
assert snake_to_camel_case("_leading_underscore") == "leadingUnderscore"
|
|
|
|
def test_trailing_underscore(self):
|
|
"""Should strip trailing underscores."""
|
|
|
|
assert snake_to_camel_case("trailing_underscore_") == "trailingUnderscore"
|
|
|
|
def test_leading_and_trailing_underscores(self):
|
|
"""Should strip leading and trailing underscores"""
|
|
|
|
assert snake_to_camel_case(f"{'_' * 5}hello__world{'_' * 10}") == "helloWorld"
|
|
|
|
def test_double_underscore(self):
|
|
"""Should treat repeated underscores as single underscore."""
|
|
|
|
assert snake_to_camel_case("double__underscore") == "doubleUnderscore"
|
|
assert snake_to_camel_case(f"many{'_' * 30}underscore") == "manyUnderscore"
|
|
|
|
|
|
def test_get_head_node_id():
|
|
"""Test get_head_node_id() returning the correct head node id.
|
|
|
|
When there are woker node, dead head node, and other alive head nodes,
|
|
get_head_node_id() should return the node id of the first alive head node.
|
|
When there are no alive head nodes, get_head_node_id() should raise assertion error.
|
|
"""
|
|
nodes = [
|
|
{"NodeID": "worker_node1", "Alive": True, "Resources": {"CPU": 1}},
|
|
{
|
|
"NodeID": "dead_head_node1",
|
|
"Alive": False,
|
|
"Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0},
|
|
},
|
|
{
|
|
"NodeID": "alive_head_node1",
|
|
"Alive": True,
|
|
"Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0},
|
|
},
|
|
{
|
|
"NodeID": "alive_head_node2",
|
|
"Alive": True,
|
|
"Resources": {"CPU": 1, HEAD_NODE_RESOURCE_NAME: 1.0},
|
|
},
|
|
]
|
|
with patch("ray.nodes", return_value=nodes):
|
|
assert get_head_node_id() == "alive_head_node1"
|
|
|
|
with patch("ray.nodes", return_value=[]):
|
|
with pytest.raises(AssertionError):
|
|
get_head_node_id()
|
|
|
|
|
|
def test_calculate_remaining_timeout():
|
|
# Always return `None` or negative value.
|
|
assert (
|
|
calculate_remaining_timeout(
|
|
timeout_s=None,
|
|
start_time_s=100,
|
|
curr_time_s=101,
|
|
)
|
|
is None
|
|
)
|
|
|
|
assert (
|
|
calculate_remaining_timeout(
|
|
timeout_s=-1,
|
|
start_time_s=100,
|
|
curr_time_s=101,
|
|
)
|
|
== -1
|
|
)
|
|
|
|
# Return delta from start.
|
|
assert (
|
|
calculate_remaining_timeout(
|
|
timeout_s=10,
|
|
start_time_s=100,
|
|
curr_time_s=101,
|
|
)
|
|
== 9
|
|
)
|
|
|
|
assert (
|
|
calculate_remaining_timeout(
|
|
timeout_s=100,
|
|
start_time_s=100,
|
|
curr_time_s=101.1,
|
|
)
|
|
== 98.9
|
|
)
|
|
|
|
# Never return a negative timeout once it has elapsed.
|
|
assert (
|
|
calculate_remaining_timeout(
|
|
timeout_s=10,
|
|
start_time_s=100,
|
|
curr_time_s=111,
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
def test_get_all_live_placement_group_names(ray_instance):
|
|
"""Test the logic to parse the Ray placement group table.
|
|
|
|
The test contains three cases:
|
|
- A named placement group that was created and is still alive ("pg2").
|
|
- A named placement group that's still being scheduled ("pg3").
|
|
- An unnamed placement group.
|
|
|
|
Only "pg2" and "pg3" should be returned as live placement group names.
|
|
"""
|
|
|
|
# Named placement group that's been removed (should not be returned).
|
|
pg1 = ray.util.placement_group([{"CPU": 0.1}], name="pg1")
|
|
ray.util.remove_placement_group(pg1)
|
|
|
|
# Named, detached placement group that's been removed (should not be returned).
|
|
pg2 = ray.util.placement_group([{"CPU": 0.1}], name="pg2", lifetime="detached")
|
|
ray.util.remove_placement_group(pg2)
|
|
|
|
# Named placement group that's still alive (should be returned).
|
|
pg3 = ray.util.placement_group([{"CPU": 0.1}], name="pg3")
|
|
assert pg3.wait()
|
|
|
|
# Named, detached placement group that's still alive (should be returned).
|
|
pg4 = ray.util.placement_group([{"CPU": 0.1}], name="pg4", lifetime="detached")
|
|
assert pg4.wait()
|
|
|
|
# Named placement group that's being scheduled (should be returned).
|
|
pg5 = ray.util.placement_group([{"CPU": 1000}], name="pg5")
|
|
assert not pg5.wait(timeout_seconds=0.001)
|
|
|
|
# Named, detached placement group that's being scheduled (should be returned).
|
|
pg6 = ray.util.placement_group([{"CPU": 1000}], name="pg6", lifetime="detached")
|
|
assert not pg6.wait(timeout_seconds=0.001)
|
|
|
|
# Unnamed placement group (should not be returned).
|
|
pg7 = ray.util.placement_group([{"CPU": 0.1}])
|
|
assert pg7.wait()
|
|
|
|
# Unnamed, detached placement group (should not be returned).
|
|
pg8 = ray.util.placement_group([{"CPU": 0.1}], lifetime="detached")
|
|
assert pg8.wait()
|
|
|
|
# Use wait_for_condition to allow GCS state to propagate after
|
|
# remove_placement_group(). Removal is async — the scheduling state may not
|
|
# have flipped to REMOVED by the time we query.
|
|
def check_live_placement_group_names():
|
|
assert set(get_all_live_placement_group_names()) == {
|
|
"pg3",
|
|
"pg4",
|
|
"pg5",
|
|
"pg6",
|
|
}
|
|
return True
|
|
|
|
wait_for_condition(check_live_placement_group_names, timeout=10)
|
|
|
|
|
|
def test_get_active_placement_group_ids(ray_instance):
|
|
"""Test that get_active_placement_group_ids returns PG IDs for alive actors only.
|
|
|
|
Cases covered:
|
|
- A PG with a live Serve actor scheduled on it (should be returned).
|
|
- A PG with no actors (should NOT be returned).
|
|
- An actor without a PG (should NOT cause its absence to break anything).
|
|
- A PG whose actor has been killed (should NOT be returned).
|
|
- A PG with a live non-Serve actor (should NOT be returned).
|
|
"""
|
|
|
|
# PG with a live actor scheduled on it
|
|
pg_with_actor = ray.util.placement_group([{"CPU": 0.1}], name="pg_with_actor")
|
|
assert pg_with_actor.wait()
|
|
|
|
@ray.remote(num_cpus=0.1)
|
|
class DummyActor:
|
|
def ready(self):
|
|
return True
|
|
|
|
actor_on_pg = DummyActor.options(
|
|
namespace=SERVE_NAMESPACE,
|
|
placement_group=pg_with_actor,
|
|
).remote()
|
|
ray.get(actor_on_pg.ready.remote())
|
|
|
|
# PG with no actors scheduled on it
|
|
pg_no_actor = ray.util.placement_group([{"CPU": 0.1}], name="pg_no_actor")
|
|
assert pg_no_actor.wait()
|
|
|
|
# Actor without any PG
|
|
actor_no_pg = DummyActor.options(namespace=SERVE_NAMESPACE).remote()
|
|
ray.get(actor_no_pg.ready.remote())
|
|
|
|
# PG whose actor will be killed
|
|
pg_killed = ray.util.placement_group([{"CPU": 0.1}], name="pg_killed")
|
|
assert pg_killed.wait()
|
|
actor_killed = DummyActor.options(
|
|
namespace=SERVE_NAMESPACE,
|
|
placement_group=pg_killed,
|
|
).remote()
|
|
ray.get(actor_killed.ready.remote())
|
|
ray.kill(actor_killed)
|
|
|
|
# PG with a live actor in a non-Serve namespace (should be excluded)
|
|
pg_non_serve = ray.util.placement_group([{"CPU": 0.1}], name="pg_non_serve")
|
|
assert pg_non_serve.wait()
|
|
actor_non_serve = DummyActor.options(
|
|
namespace="non_serve",
|
|
placement_group=pg_non_serve,
|
|
).remote()
|
|
ray.get(actor_non_serve.ready.remote())
|
|
|
|
pg_with_actor_id = pg_with_actor.id.hex()
|
|
pg_no_actor_id = pg_no_actor.id.hex()
|
|
pg_killed_id = pg_killed.id.hex()
|
|
pg_non_serve_id = pg_non_serve.id.hex()
|
|
|
|
# Use wait_for_condition to allow GCS state to propagate after ray.kill().
|
|
# ray.kill() is async — the actor's ALIVE -> DEAD transition in GCS may lag.
|
|
def check_active_placement_group_ids():
|
|
active_ids = get_active_placement_group_ids()
|
|
assert pg_with_actor_id in active_ids
|
|
assert pg_no_actor_id not in active_ids
|
|
assert pg_killed_id not in active_ids
|
|
assert pg_non_serve_id not in active_ids
|
|
return True
|
|
|
|
wait_for_condition(check_active_placement_group_ids, timeout=10)
|
|
|
|
|
|
def test_is_running_in_asyncio_loop_false():
|
|
assert is_running_in_asyncio_loop() is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_running_in_asyncio_loop_true():
|
|
assert is_running_in_asyncio_loop() is True
|
|
|
|
async def check():
|
|
return is_running_in_asyncio_loop()
|
|
|
|
# Verify that it also works in a task.
|
|
assert await asyncio.ensure_future(check()) is True
|
|
|
|
|
|
def test_get_current_actor_id(ray_instance):
|
|
@ray.remote
|
|
class A:
|
|
def call_get_current_actor_id(self):
|
|
return get_current_actor_id()
|
|
|
|
a = A.remote()
|
|
actor_id = ray.get(a.call_get_current_actor_id.remote())
|
|
assert len(actor_id) > 0
|
|
assert actor_id != "DRIVER"
|
|
|
|
assert get_current_actor_id() == "DRIVER"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_semaphore():
|
|
"""Test core Semaphore functionality."""
|
|
max_value = 2
|
|
sema = Semaphore(get_value_fn=lambda: max_value)
|
|
|
|
# Test get_max_value functionality
|
|
assert sema.get_max_value() == max_value
|
|
|
|
# Initially, semaphore should not be locked and should allow acquisitions
|
|
assert not sema.locked()
|
|
|
|
# Acquire one
|
|
await sema.acquire()
|
|
assert not sema.locked()
|
|
assert sema._value == 1
|
|
|
|
# Acquire one
|
|
await sema.acquire()
|
|
assert sema.locked() # Should now be locked (2 out of 2)
|
|
assert sema._value == 2
|
|
|
|
# Release one
|
|
sema.release()
|
|
assert not sema.locked() # Should not be locked anymore (1 out of 2)
|
|
assert sema._value == 1
|
|
|
|
# Acquire one
|
|
await sema.acquire()
|
|
assert sema.locked()
|
|
assert sema._value == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_semaphore_waiters_and_single_release():
|
|
"""Test that release() wakes up exactly one waiter."""
|
|
max_value = 1
|
|
sema = Semaphore(get_value_fn=lambda: max_value)
|
|
|
|
# Fill the semaphore to capacity
|
|
await sema.acquire()
|
|
assert sema.locked()
|
|
assert sema._value == 1
|
|
|
|
# Create multiple waiters
|
|
waiters_completed = []
|
|
|
|
async def waiter(waiter_id):
|
|
await sema.acquire()
|
|
waiters_completed.append(waiter_id)
|
|
|
|
# Start 2 waiters that will all block
|
|
waiter_tasks = [
|
|
asyncio.create_task(waiter(1)),
|
|
asyncio.create_task(waiter(2)),
|
|
]
|
|
|
|
# Yield the event loop
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Verify they are all waiting
|
|
assert len(waiters_completed) == 0
|
|
assert sema.locked()
|
|
assert len(sema._waiters) == 2
|
|
|
|
# Release once - this should wake up exactly ONE waiter
|
|
sema.release()
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Verify exactly one waiter was woken up and completed
|
|
assert len(waiters_completed) == 1
|
|
assert sema._value == 1
|
|
assert sema.locked()
|
|
assert len(sema._waiters) == 1
|
|
|
|
# Release again - should wake up exactly one more waiter
|
|
sema.release()
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Verify exactly one more waiter was woken up
|
|
assert len(waiters_completed) == 2
|
|
assert sema._value == 1
|
|
assert sema.locked()
|
|
assert len(sema._waiters) == 0
|
|
|
|
assert len(await asyncio.gather(*waiter_tasks)) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_semaphore_dynamic_max_value():
|
|
"""Test that Semaphore respects dynamic changes to max_value."""
|
|
current_max = 2
|
|
|
|
def get_dynamic_max():
|
|
return current_max
|
|
|
|
sema = Semaphore(get_value_fn=get_dynamic_max)
|
|
|
|
# Initially max is 2
|
|
assert sema.get_max_value() == 2
|
|
|
|
# Acquire up to the limit
|
|
await sema.acquire()
|
|
await sema.acquire()
|
|
assert sema.locked()
|
|
|
|
# Increase the max value dynamically
|
|
current_max = 3
|
|
assert sema.get_max_value() == 3
|
|
assert not sema.locked()
|
|
|
|
# Should be able to acquire one more
|
|
await sema.acquire()
|
|
assert sema.locked()
|
|
assert sema._value == 3
|
|
|
|
# Decrease the max value
|
|
current_max = 1
|
|
assert sema.get_max_value() == 1
|
|
assert sema.locked()
|
|
|
|
# Release to get back within limits
|
|
sema.release()
|
|
sema.release()
|
|
assert sema.locked()
|
|
sema.release()
|
|
assert not sema.locked()
|
|
assert sema._value == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|