def test_repo_double_init(tmp_path: str) -> None:
"""explicitly constructing another repository should fail"""
_ = Repo.init(tmp_path)
os.mkdir(os.path.join(tmp_path, ZENML_DIR_NAME))
with pytest.raises(Exception):
Repository(str(tmp_path)).init_repo(repo_path=tmp_path)
def test_initializing_repository_without_git_repo_does_not_raise_error(
tmp_path: str, ) -> None:
"""Check initializing repository without git repository does not raise
error"""
with does_not_raise():
repo = Repository(str(tmp_path))
assert repo.git_wrapper is None
def test_get_pipeline_returns_none_if_non_existent(tmp_path: str) -> None:
"""Check get_pipeline returns None if it doesn't exist"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
repo.set_active_stack("local_stack")
our_pipeline = repo.get_pipeline("not_a_pipeline")
assert our_pipeline is None
def list_stacks() -> None:
"""List all available stacks from service."""
service = Repository().get_service()
cli_utils.title("Stacks:")
# TODO [ENG-144]: once there is a common superclass for Stack/ArtifactStore etc.,
# remove the mypy ignore
cli_utils.echo_component_list(service.stacks) # type: ignore[arg-type]
def register_metadata_store(metadata_store_name: str, metadata_store_type: str,
args: List[str]) -> None:
"""Register a metadata store."""
try:
parsed_args = cli_utils.parse_unknown_options(args)
except AssertionError as e:
cli_utils.error(str(e))
return
repo: Repository = Repository()
try:
# TODO [ENG-187]: Remove when we rework the registry logic
from zenml.core.component_factory import metadata_store_factory
comp = metadata_store_factory.get_single_component(metadata_store_type)
except AssertionError as e:
cli_utils.error(str(e))
return
metadata_store = comp(**parsed_args)
service = repo.get_service()
service.register_metadata_store(metadata_store_name, metadata_store)
cli_utils.declare(
f"Metadata Store `{metadata_store_name}` successfully registered!")
def test_getting_active_stack_returns_local_stack(tmp_path: str, ) -> None:
"""Check getting the active stack"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
repo.set_active_stack("local_stack")
assert repo.get_active_stack() == LocalService().get_stack("local_stack")
assert isinstance(repo.get_active_stack(), BaseStack)
def test_init_repo_creates_a_zen_folder(tmp_path: str) -> None:
"""Check initializing repository creates a ZenML folder"""
_ = Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_stack = LocalService().get_stack("local_stack")
repo.init_repo(repo_path=tmp_path, stack=local_stack)
assert os.path.exists(os.path.join(tmp_path, ZENML_DIR_NAME))
def is_inside_repository(file_path: str) -> bool:
"""Returns whether a file is inside a zenml repository."""
from zenml.core.repo import Repository
repo_path = pathlib.Path(Repository().path).resolve()
absolute_file_path = pathlib.Path(file_path).resolve()
return repo_path in absolute_file_path.parents
def test_get_artifact_store_returns_artifact_store(tmp_path: str) -> None:
"""Test get_artifact_store returns artifact store."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
artifact_store = local_service.get_artifact_store(LOCAL_ARTIFACT_STORE_NAME)
assert artifact_store is not None
assert isinstance(artifact_store, BaseArtifactStore)
def test_get_pipelines_returns_list(tmp_path: str) -> None:
"""Check get_pipelines returns a list"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
repo.set_active_stack("local_stack")
our_pipelines = repo.get_pipelines()
assert our_pipelines is not None
assert isinstance(our_pipelines, list)
def test_get_metadata_store_returns_metadata_store(tmp_path: str) -> None:
"""Test get_metadata_store returns metadata store."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
metadata_store = local_service.get_metadata_store(LOCAL_METADATA_STORE_NAME)
assert metadata_store is not None
assert isinstance(metadata_store, BaseMetadataStore)
def test_initializing_repo_with_git_repo_present_sets_git_wrapper(
tmp_path: str, ) -> None:
"""Check initializing repository with git repository sets git wrapper"""
git_repo_instance = Repo.init(tmp_path)
repo = Repository(str(tmp_path))
assert repo.git_wrapper is not None
assert repo.git_wrapper.repo_path == str(tmp_path)
assert repo.git_wrapper.git_repo == git_repo_instance
def test_get_orchestrator_returns_orchestrator(tmp_path: str) -> None:
"""Test get_orchestrator returns orchestrator"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
orchestrator = local_service.get_orchestrator(LOCAL_ORCHESTRATOR_NAME)
assert orchestrator is not None
assert isinstance(orchestrator, BaseOrchestrator)
def test_getting_the_active_service_returns_local_service(
tmp_path: str, ) -> None:
"""Check getting the active service"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
assert repo.get_service() is not None
assert isinstance(repo.get_service(), BaseComponent)
assert isinstance(repo.get_service(), LocalService)
assert repo.get_service() == repo.service
def test_get_stack_raises_exception_when_key_does_not_exist(
tmp_path: str,
) -> None:
"""Test get_stack raises exception when key does not exist."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
with pytest.raises(DoesNotExistException):
local_service.get_stack("made_up_stack")
def test_register_orchestrator_works_as_expected(tmp_path: str) -> None:
"""Test register_orchestrator method registers an orchestrator as expected."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
orchestrator = LocalOrchestrator()
local_service.register_orchestrator("local_orchestrator_2", orchestrator)
assert local_service.get_orchestrator("local_orchestrator_2") is not None
local_service.delete_orchestrator("local_orchestrator_2")
def test_get_pipeline_returns_same_when_stack_specified(tmp_path: str) -> None:
"""Check get_pipeline returns the same if stack specified"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
repo.set_active_stack("local_stack")
our_pipeline_default = repo.get_pipeline("pipeline_1")
our_pipeline_local = repo.get_pipeline("pipeline_1",
stack_key="local_stack")
assert our_pipeline_default == our_pipeline_local
def test_register_stack_raises_exception_when_key_already_exists(
tmp_path: str,
) -> None:
"""Test register_stack raises exception when key already exists."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
local_stack1 = local_service.get_stack("local_stack")
with pytest.raises(AlreadyExistsException):
local_service.register_stack("local_stack", local_stack1)
def test_register_orchestrator_with_existing_key_fails(tmp_path: str) -> None:
"""Test register_orchestrator with existing key fails."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
with pytest.raises(AlreadyExistsException):
local_service.register_orchestrator(
LOCAL_ORCHESTRATOR_NAME,
local_service.get_orchestrator(LOCAL_ORCHESTRATOR_NAME),
)
def test_delete_orchestrator_works(tmp_path: str) -> None:
"""Test delete_orchestrator works as expected."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
orchestrator = LocalOrchestrator()
local_service.register_orchestrator("local_orchestrator_2", orchestrator)
local_service.delete_orchestrator("local_orchestrator_2")
with pytest.raises(DoesNotExistException):
local_service.get_orchestrator("local_orchestrator_2")
def test_register_metadata_store_with_existing_key_fails(tmp_path: str) -> None:
"""Test register_metadata_store with existing key fails."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
with pytest.raises(AlreadyExistsException):
local_service.register_metadata_store(
LOCAL_METADATA_STORE_NAME,
local_service.get_metadata_store(LOCAL_METADATA_STORE_NAME),
)
def test_local_service_can_access_orchestrators(tmp_path: str) -> None:
"""Local Service can access orchestrators."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
assert local_service.orchestrators is not None
assert isinstance(local_service.orchestrators, dict)
assert isinstance(
local_service.orchestrators[LOCAL_ORCHESTRATOR_NAME],
BaseOrchestrator,
)
def test_local_service_can_access_artifact_stores(tmp_path: str) -> None:
"""Local Service can access artifact stores."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
assert local_service.artifact_stores is not None
assert isinstance(local_service.artifact_stores, dict)
assert isinstance(
local_service.artifact_stores[LOCAL_ARTIFACT_STORE_NAME],
BaseArtifactStore,
)
def test_local_service_can_access_metadata_stores(tmp_path: str) -> None:
"""Local Service can access metadata stores."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
assert local_service.metadata_stores is not None
assert isinstance(local_service.metadata_stores, dict)
assert isinstance(
local_service.metadata_stores[LOCAL_METADATA_STORE_NAME],
BaseMetadataStore,
)
def test_get_pipelines_returns_same_list_when_stack_specified(
tmp_path) -> None:
"""Check get_pipelines returns the same list when stack specified"""
# TODO [MEDIUM]: update test once we have custom environments being created
# to check with actual pipelines
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
repo.set_active_stack("local_stack")
our_pipelines_default = repo.get_pipelines()
our_pipelines_local = repo.get_pipelines(stack_key="local_stack")
assert our_pipelines_default == our_pipelines_local
def test_delete_stack_deletes_the_stack(tmp_path: str) -> None:
"""Test delete_stack deletes the stack."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
local_stack1 = local_service.get_stack("local_stack")
local_service.register_stack("local_stack_2", local_stack1)
local_stack2 = local_service.get_stack("local_stack_2")
assert local_stack2 is not None
local_service.delete_stack("local_stack_2")
with pytest.raises(DoesNotExistException):
local_service.get_stack("local_stack_2")
def test_delete_artifact_store_works(tmp_path: str) -> None:
"""Test delete_artifact_store works as expected."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
artifact_store_dir = os.path.join(tmp_path, "test_store")
local_artifact_store = LocalArtifactStore(path=artifact_store_dir)
local_service.register_artifact_store(
"local_artifact_store_2", local_artifact_store
)
local_service.delete_artifact_store("local_artifact_store_2")
with pytest.raises(DoesNotExistException):
local_service.get_artifact_store("local_artifact_store_2")
def test_register_stack_works_as_expected(tmp_path: str) -> None:
"""Test register_stack method registers a stack as expected."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
local_stack1 = local_service.get_stack("local_stack")
local_service.register_stack("local_stack_2", local_stack1)
local_stack2 = local_service.get_stack("local_stack_2")
assert local_stack2 is not None
assert local_stack2.orchestrator == local_stack1.orchestrator
assert local_stack2.artifact_store == local_stack1.artifact_store
assert local_stack2.metadata_store == local_stack1.metadata_store
# TODO [MEDIUM]: rework this as a fixture to be run after the test completes
local_service.delete_stack("local_stack_2")
def test_delete_metadata_store_works(tmp_path: str) -> None:
"""Test delete_metadata_store works as expected"""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
artifact_store_dir = os.path.join(tmp_path, "test_store")
metadata_file = os.path.join(artifact_store_dir, "metadata.db")
metadata_store = SQLiteMetadataStore(uri=metadata_file)
local_service.register_metadata_store(
"local_metadata_store_2", metadata_store
)
local_service.delete_metadata_store("local_metadata_store_2")
with pytest.raises(DoesNotExistException):
local_service.get_metadata_store("local_metadata_store_2")
def test_register_artifact_store_works_as_expected(tmp_path: str) -> None:
"""Test register_artifact_store method registers an artifact store as expected."""
Repo.init(tmp_path)
repo = Repository(str(tmp_path))
local_service = repo.get_service()
artifact_store_dir = os.path.join(tmp_path, "test_store")
local_artifact_store = LocalArtifactStore(path=artifact_store_dir)
local_service.register_artifact_store(
"local_artifact_store_2", local_artifact_store
)
assert (
local_service.get_artifact_store("local_artifact_store_2") is not None
)
local_service.delete_artifact_store("local_artifact_store_2")
