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!")