remote tests are functional!

This commit is contained in:
Ashwin Bharambe 2024-11-12 19:17:58 -08:00
parent 8b7be87bec
commit 0121114a5d
5 changed files with 93 additions and 89 deletions

View file

@ -35,8 +35,8 @@ def remote_stack_fixture() -> ProviderFixture:
return ProviderFixture(
providers=[
Provider(
provider_id="remote",
provider_type="remote",
provider_id="test::remote",
provider_type="test::remote",
config=config.model_dump(),
)
],

View file

@ -17,11 +17,25 @@ from llama_stack.distribution.build import print_pip_install_help
from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
from llama_stack.distribution.distribution import get_provider_registry
from llama_stack.distribution.request_headers import set_request_provider_data
from llama_stack.distribution.resolver import resolve_impls
from llama_stack.distribution.resolver import resolve_impls, resolve_remote_stack_impls
from llama_stack.distribution.stack import construct_stack
from llama_stack.providers.utils.kvstore import SqliteKVStoreConfig
async def construct_stack_for_test(run_config: StackRunConfig):
remote_config = remote_provider_config(run_config)
if not remote_config:
return await construct_stack(run_config)
impls = await resolve_remote_stack_impls(remote_config, run_config.apis)
# we don't register resources for a remote stack as part of the fixture setup
# because the stack is already "up". if a test needs to register resources, it
# can do so manually always.
return impls
async def resolve_impls_for_test_v2(
apis: List[Api],
providers: Dict[str, List[Provider]],
@ -49,7 +63,7 @@ async def resolve_impls_for_test_v2(
)
run_config = parse_and_maybe_upgrade_config(run_config)
try:
impls = await construct_stack(run_config)
impls = await construct_stack_for_test(run_config)
except ModuleNotFoundError as e:
print_pip_install_help(providers)
raise e
@ -62,6 +76,24 @@ async def resolve_impls_for_test_v2(
return impls
def remote_provider_config(
run_config: StackRunConfig,
) -> Optional[RemoteProviderConfig]:
remote_config = None
has_non_remote = False
for api_providers in run_config.providers.values():
for provider in api_providers:
if provider.provider_type == "test::remote":
remote_config = RemoteProviderConfig(**provider.config)
else:
has_non_remote = True
if remote_config:
assert not has_non_remote, "Remote stack cannot have non-remote providers"
return remote_config
async def resolve_impls_for_test(api: Api, deps: List[Api] = None):
if "PROVIDER_CONFIG" not in os.environ:
raise ValueError(