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 init(
repo_path: Optional[str],
) -> None:
"""Initialize ZenML on given path.
Args:
repo_path: Path to the repository.
Raises:
InvalidGitRepositoryError: If repo is not a git repo.
AssertionError
"""
if sys.version_info.minor == 6:
warning(
"ZenML support for Python 3.6 will be deprecated soon. Please "
"consider upgrading your Python version to ensure ZenML works "
"properly in the future."
)
if repo_path is None:
repo_path = os.getcwd()
declare(f"Initializing at {repo_path}")
try:
Repository.init_repo(repo_path=repo_path)
declare(f"ZenML repo initialized at {repo_path}")
except git.InvalidGitRepositoryError: # type: ignore[attr-defined]
error(
f"{repo_path} is not a valid git repository! Please "
f"initialize ZenML within a git repository using "
f"`git init `"
)
except AssertionError as e:
error(f"{e}")
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 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_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_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_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_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_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_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 pytest_sessionstart(session):
"""Called after the Session object has been created and
before performing collection and entering the run test loop.
"""
os.environ[ENV_ZENML_DEBUG] = "true"
try:
Repository.init_repo()
except AssertionError:
# already initialized
logging.info("Repo already initialized for testing.")
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_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_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_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_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")
def test_register_metadata_store_works_as_expected(tmp_path: str) -> None:
"""Test register_metadata_store method registers an metadata 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")
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
)
assert (
local_service.get_metadata_store("local_metadata_store_2") is not None
)
local_service.delete_metadata_store("local_metadata_store_2")
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_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_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 register_stack(
stack_name: str,
metadata_store: str,
artifact_store: str,
orchestrator: str,
) -> None:
"""Register a stack."""
service = Repository().get_service()
stack = BaseStack(
artifact_store_name=artifact_store,
orchestrator_name=orchestrator,
metadata_store_name=metadata_store,
)
service.register_stack(stack_name, stack)
cli_utils.declare(f"Stack `{stack_name}` successfully registered!")
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!")