Skip to content

Commit

Permalink
Refactor test pool to be less optional (#75)
Browse files Browse the repository at this point in the history
  • Loading branch information
mmcgr authored Aug 18, 2023
1 parent bff097c commit 6105e57
Show file tree
Hide file tree
Showing 6 changed files with 161 additions and 187 deletions.
9 changes: 1 addition & 8 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -72,14 +72,7 @@ Run multiple tests concurrently with a custom engine provider.
using RAITest
using Test
# Always use `my_existing_engine_name` as the engine name. This must already be created as
# RAITest will not create the engine itself
set_engine_name_provider!()->"my_existing_engine_name")
# Releasing the engine does nothing. Controlling the engine use and cleanup is not handled
# by RAITest
set_engine_name_releaser!(name::String)->return)
set_engine_creater!((name)->create_default_engine(name, "M"))
@testset "My tests" begin
for i in 1:10
query = "def output = $i ic { output = $i }"
Expand Down
2 changes: 0 additions & 2 deletions src/RAITest.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@ export add_test_engine!

export set_context!

export set_engine_name_provider!
export set_engine_name_releaser!
export set_engine_creater!

include("code-util.jl")
Expand Down
99 changes: 14 additions & 85 deletions src/engines.jl
Original file line number Diff line number Diff line change
@@ -1,12 +1,3 @@
mutable struct TestEngineProvision
# Returns the name of a provisioned, currently valid, engine
provider::Function
# Sends a notification that the engine represented by a string name is no longer in use
releaser::Function
# Create an engine. This is expected to be used by the provider as needed.
creater::Function
end

function _wait_till_provisioned(engine_name, max_wait_time_s=600)
start_time = time()
# This should be a rare event, so a coarse-grained period is acceptable
Expand All @@ -32,96 +23,34 @@ function _wait_till_provisioned(engine_name, max_wait_time_s=600)
end

"""
create_default_engine(name::String)
create_default_engine(name::String; size::String)
Create an XS engine with default settings and the provided name.
Create an engine with the default configuration.
If the engine already exists, return immediately. If not, create the engine then
return once the provisioning process is complete, or failed.
"""
function create_default_engine(name::String)
size = "XS"
function create_default_engine(name::String; size::String="XS")
max_wait_time_s = 600

try
get_engine(get_context(), name; readtimeout=max_wait_time_s)
# The engine already exists so return it immediately
return name
catch
# Engine does not exist yet so we'll need to create it
create_engine(get_context(), name; size, readtimeout=max_wait_time_s)
catch e
# If the status code is 409 then the engine already exists and we can wait for it
# to be ready
if !(e isa HTTPError) || e.status_code != 409
rethrow()
end
end

# Request engine creation
create_engine(get_context(), name; size=size, readtimeout=max_wait_time_s)

# Wait for engine to be provisioned
return _wait_till_provisioned(name, max_wait_time_s)
end

# Get test engine. If a name is provided, the corresponding engine will be provided.
get_test_engine()::String = TEST_ENGINE_PROVISION.provider()
# Get test engine.
get_test_engine()::String = get_free_test_engine_name()

# Release test engine. Notifies the provider that this engine is no longer in use.
release_test_engine(name::String) = TEST_ENGINE_PROVISION.releaser(name)

"""
set_engine_name_provider!(provider::Function)
Set a provider for test engine names.
The provider is called by each test to select an engine to run the test with. The default
provider selects a name from a pool of available test engines.
# Examples
```
set_engine_name_provider!() -> "MyEngine")
set_engine_name_provider!() -> my_custom_engine_selector())
```
"""
function set_engine_name_provider!(provider::Function)
return TEST_ENGINE_PROVISION.provider = provider
end

"""
set_engine_name_releaser!(releaser::Function)
Set a releaser for test engine names.
The releaser will be called after a test has been run. The default releaser
marks an engine in the test engine pool as available for use by another test.
# Examples
```
set_engine_name_releaser!(::String) -> nothing)
set_engine_name_releaser!(name::String) -> my_custom_engine_releaser(name))
```
"""
function set_engine_name_releaser!(releaser::Function)
return TEST_ENGINE_PROVISION.releaser = releaser
end

"""
set_engine_creater!(creater::Function)
Set a function used to create engines.
# Examples
```
set_engine_creater!(create_default_engine)
```
"""
function set_engine_creater!(creater::Function)
return TEST_ENGINE_PROVISION.creater = creater
end

TEST_ENGINE_POOL = TestEnginePool(Dict{String, Int64}(), 0, get_next_engine_name)
release_test_engine(name::String) = release_pooled_test_engine(name)

TEST_ENGINE_PROVISION = TestEngineProvision(
get_pooled_test_engine,
release_pooled_test_engine,
create_default_engine,
)
TEST_ENGINE_POOL = TestEnginePool()
Loading

0 comments on commit 6105e57

Please sign in to comment.