#!/usr/bin/env python
"""Collection of tests for :mod:`orion.storage`."""

import copy
import datetime
import logging
import os
import pickle
import time

import pytest

import orion.core
from orion.core.io.database import DuplicateKeyError
from orion.core.io.database.pickleddb import PickledDB
from orion.core.worker.trial import Trial
from orion.storage.base import (
    FailedUpdate,
    LockAcquisitionTimeout,
    LockedAlgorithmState,
    MissingArguments,
    setup_storage,
)
from orion.storage.legacy import Legacy
from orion.storage.track import HAS_TRACK, REASON
from orion.testing import OrionState, base_experiment

log = logging.getLogger(__name__)
log.setLevel(logging.WARNING)

storage_backends = [None]  # defaults to legacy with PickleDB

if not HAS_TRACK:
    log.warning("Track is not tested because: %s!", REASON)
else:
    storage_backends.append({"type": "track", "uri": "file://${file}?objective=loss"})


base_trial = {
    "experiment": "default_name",
    "status": "new",  # new, reserved, suspended, completed, broken
    "worker": None,
    "submit_time": "2017-11-23T02:00:00",
    "start_time": None,
    "end_time": None,
    "heartbeat": None,
    "results": [
        {"name": "loss", "type": "objective", "value": 2}  # objective, constraint
    ],
    "params": [
        {"name": "/encoding_layer", "type": "categorical", "value": "rnn"},
        {
            "name": "/decoding_layer",
            "type": "categorical",
            "value": "lstm_with_attention",
        },
    ],
}


def _generate(obj, *args, value):
    if obj is None:
        return None

    obj = copy.deepcopy(obj)
    data = obj

    for arg in args[:-1]:
        data = data[arg]

    data[args[-1]] = value
    return obj


def make_lost_trial(delay=2):
    """Make a lost trial"""
    obj = copy.deepcopy(base_trial)
    obj["status"] = "reserved"
    obj["heartbeat"] = datetime.datetime.utcnow() - datetime.timedelta(
        seconds=orion.core.config.worker.heartbeat * delay
    )
    obj["params"].append(
        {"name": "/index", "type": "categorical", "value": f"lost_trial_{delay}"}
    )
    return obj


all_status = ["completed", "broken", "reserved", "interrupted", "suspended", "new"]


# TODO: Reuse the one from orion.testing
def generate_trials(status=None, heartbeat=None):
    """Generate Trials with different configurations"""
    if status is None:
        status = all_status

    new_trials = [_generate(base_trial, "status", value=s) for s in status]

    if heartbeat:
        for trial in new_trials:
            trial["heartbeat"] = heartbeat

    # make each trial unique
    for i, trial in enumerate(new_trials):
        trial["params"].append({"name": "/index", "type": "categorical", "value": i})

    return new_trials


def generate_experiments():
    """Generate a set of experiments"""
    users = ["a", "b", "c"]
    exps = [_generate(base_experiment, "metadata", "user", value=u) for u in users]
    return [_generate(exp, "name", value=str(i)) for i, exp in enumerate(exps)]


@pytest.mark.usefixtures("orionstate")
def test_setup_storage_default():
    """Test that storage is setup using default config"""

    storage = setup_storage()
    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)


def test_setup_storage_bad():
    """Test how setup fails when configuring with non-existent backends"""

    with pytest.raises(NotImplementedError) as exc:
        setup_storage({"type": "idontexist"})

    assert exc.match("idontexist")


def test_setup_storage_custom():
    """Test setup with local configuration"""

    storage = setup_storage(
        {"type": "legacy", "database": {"type": "pickleddb", "host": "test.pkl"}}
    )
    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)
    assert storage._db.host == os.path.abspath("test.pkl")


def test_setup_storage_custom_type_missing():
    """Test setup with local configuration with type missing"""

    storage = setup_storage({"database": {"type": "pickleddb", "host": "test.pkl"}})

    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)
    assert storage._db.host == os.path.abspath("test.pkl")


@pytest.mark.usefixtures("orionstate")
def test_setup_storage_custom_legacy_emtpy():
    """Test setup with local configuration with legacy but no config"""

    storage = setup_storage({"type": "legacy"})

    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)
    assert storage._db.host == orion.core.config.storage.database.host


def test_setup_storage_bad_override():
    """Test setup with different type than existing singleton"""

    storage = setup_storage(
        {"type": "legacy", "database": {"type": "pickleddb", "host": "test.pkl"}}
    )
    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)


def test_setup_storage_stateless():
    """Test that passed configuration dictionary is not modified by the function"""

    config = {"database": {"type": "pickleddb", "host": "test.pkl"}}
    passed_config = copy.deepcopy(config)
    setup_storage(passed_config)
    assert config == passed_config


@pytest.mark.usefixtures("version_XYZ")
@pytest.mark.parametrize("storage", storage_backends)
class TestStorage:
    """Test all storage backend"""

    def test_create_experiment(self, storage):
        """Test create experiment"""
        with OrionState(experiments=[], storage=storage) as cfg:
            storage = cfg.storage

            storage.create_experiment(base_experiment)

            experiments = storage.fetch_experiments({})
            assert len(experiments) == 1, "Only one experiment in the database"

            experiment = experiments[0]
            assert base_experiment == experiment, "Local experiment and DB should match"

            # Insert it again
            with pytest.raises(DuplicateKeyError):
                storage.create_experiment(base_experiment)

    def test_fetch_experiments(self, storage, name="0", user="a"):
        """Test fetch experiments"""
        with OrionState(experiments=generate_experiments(), storage=storage) as cfg:
            storage = cfg.storage

            experiments = storage.fetch_experiments({})
            assert len(experiments) == len(cfg.experiments)

            experiments = storage.fetch_experiments(
                {"name": name, "metadata.user": user}
            )
            assert len(experiments) == 1

            experiment = experiments[0]
            assert experiment["name"] == name, "name should match query"
            assert (
                experiment["metadata"]["user"] == user
            ), "user name should match query"

            experiments = storage.fetch_experiments(
                {"name": "-1", "metadata.user": user}
            )
            assert len(experiments) == 0

    def test_update_experiment(self, monkeypatch, storage, name="0", user="a"):
        """Test fetch experiments"""
        with OrionState(experiments=generate_experiments(), storage=storage) as cfg:
            storage = cfg.storage

            class _Dummy:
                pass

            experiment = cfg.experiments[0]
            mocked_experiment = _Dummy()
            mocked_experiment.id = experiment["_id"]

            storage.update_experiment(mocked_experiment, test=True)
            experiments = storage.fetch_experiments({"_id": experiment["_id"]})
            assert len(experiments) == 1

            fetched_experiment = experiments[0]
            assert fetched_experiment["test"]

            assert (
                "test"
                not in storage.fetch_experiments({"_id": cfg.experiments[1]["_id"]})[0]
            )

            storage.update_experiment(uid=experiment["_id"], test2=True)
            assert storage.fetch_experiments({"_id": experiment["_id"]})[0]["test2"]
            assert (
                "test2"
                not in storage.fetch_experiments({"_id": cfg.experiments[1]["_id"]})[0]
            )

            with pytest.raises(MissingArguments):
                storage.update_experiment()

            with pytest.raises(AssertionError):
                storage.update_experiment(experiment=mocked_experiment, uid="123")

    def test_delete_experiment(self, storage):
        """Test delete one experiment"""
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support deletion yet.")

        with OrionState(experiments=generate_experiments(), storage=storage) as cfg:
            storage = cfg.storage

            n_experiments = len(storage.fetch_experiments({}))
            storage.delete_experiment(uid=cfg.experiments[0]["_id"])
            experiments = storage.fetch_experiments({})
            assert len(experiments) == n_experiments - 1
            assert cfg.experiments[0]["_id"] not in [exp["_id"] for exp in experiments]

    def test_register_trial(self, storage):
        """Test register trial"""
        with OrionState(experiments=[base_experiment], storage=storage) as cfg:
            storage = cfg.storage
            trial1 = storage.register_trial(Trial(**base_trial))
            trial2 = storage.get_trial(trial1)

            assert (
                trial1.to_dict() == trial2.to_dict()
            ), "Trials should match after insert"

    def test_register_duplicate_trial(self, storage):
        """Test register trial"""
        with OrionState(
            experiments=[base_experiment], trials=[base_trial], storage=storage
        ) as cfg:
            storage = cfg.storage

            with pytest.raises(DuplicateKeyError):
                storage.register_trial(Trial(**base_trial))

    def test_update_trials(self, storage):
        """Test update many trials"""
        with OrionState(
            experiments=[base_experiment],
            trials=generate_trials(status=["completed", "reserved", "reserved"]),
            storage=storage,
        ) as cfg:
            storage = cfg.storage

            class _Dummy:
                pass

            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.fetch_trials(experiment)
            assert sum(trial.status == "reserved" for trial in trials) == 2
            count = storage.update_trials(
                experiment, where={"status": "reserved"}, status="interrupted"
            )
            assert count == 2
            trials = storage.fetch_trials(experiment)
            assert sum(trial.status == "interrupted" for trial in trials) == 2

    def test_update_trial(self, storage):
        """Test update one trial"""
        with OrionState(
            experiments=[base_experiment], trials=[base_trial], storage=storage
        ) as cfg:
            storage = cfg.storage

            trial = Trial(**cfg.trials[0])

            assert trial.status != "interrupted"
            storage.update_trial(trial, status="interrupted")
            assert storage.get_trial(trial).status == "interrupted"

    def test_reserve_trial_success(self, storage):
        """Test reserve trial"""
        with OrionState(
            experiments=[base_experiment], trials=[base_trial], storage=storage
        ) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            trial = storage.reserve_trial(experiment)

            assert trial is not None
            assert trial.status == "reserved"

    def test_reserve_trial_fail(self, storage):
        """Test reserve trial"""
        with OrionState(
            experiments=[base_experiment],
            trials=generate_trials(status=["completed", "reserved"]),
            storage=storage,
        ) as cfg:

            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            trial = storage.reserve_trial(experiment)
            assert trial is None

    def test_fetch_trials(self, storage):
        """Test fetch experiment trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            trials1 = storage.fetch_trials(experiment=experiment)
            trials2 = storage.fetch_trials(uid=experiment.id)

            with pytest.raises(MissingArguments):
                storage.fetch_trials()

            with pytest.raises(AssertionError):
                storage.fetch_trials(experiment=experiment, uid="123")

            assert len(trials1) == len(cfg.trials), "trial count should match"
            assert len(trials2) == len(cfg.trials), "trial count should match"

    def test_fetch_trials_with_query(self, storage):
        """Test fetch experiment trials with queries"""
        with OrionState(
            experiments=[base_experiment],
            trials=generate_trials(status=["completed", "reserved", "reserved"]),
            storage=storage,
        ) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            trials_all = storage.fetch_trials(experiment=experiment)
            trials_completed = storage.fetch_trials(
                experiment=experiment, where={"status": "completed"}
            )
            trials_reserved = storage.fetch_trials(
                experiment=experiment, where={"status": "reserved"}
            )

            assert len(trials_all) == len(cfg.trials), "trial count should match"
            assert len(trials_completed) == 1, "trial count should match"
            assert len(trials_reserved) == 2, "trial count should match"

    def test_delete_all_trials(self, storage):
        """Test delete all trials of an experiment"""
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support deletion yet.")

        trials = generate_trials()
        trial_from_other_exp = copy.deepcopy(trials[0])
        trial_from_other_exp["experiment"] = "other"
        trials.append(trial_from_other_exp)
        with OrionState(
            experiments=[base_experiment], trials=trials, storage=storage
        ) as cfg:
            storage = cfg.storage

            # Make sure we have sufficient trials to test deletion
            trials = storage.fetch_trials(uid="default_name")
            assert len(trials) > 2

            count = storage.delete_trials(uid="default_name")
            assert count == len(trials)
            assert storage.fetch_trials(uid="default_name") == []

            # Make sure trials from other experiments were not deleted
            assert len(storage.fetch_trials(uid="other")) == 1

    def test_delete_trials_with_query(self, storage):
        """Test delete experiment trials matching a query"""
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support deletion yet.")

        trials = generate_trials()
        trial_from_other_exp = copy.deepcopy(trials[0])
        trial_from_other_exp["experiment"] = "other"
        trials.append(trial_from_other_exp)
        with OrionState(
            experiments=[base_experiment], trials=trials, storage=storage
        ) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name")

            # Make sure we have sufficient trials to test deletion
            status = trials[0]["status"]
            trials = storage.fetch_trials(experiment)
            trials_with_status = storage.fetch_trials_by_status(experiment, status)
            assert len(trials_with_status) > 0
            assert len(trials) > len(trials_with_status)

            # Test deletion
            count = storage.delete_trials(uid="default_name", where={"status": status})
            assert count == len(trials_with_status)
            assert storage.fetch_trials_by_status(experiment, status) == []
            assert len(storage.fetch_trials(experiment)) == len(trials) - len(
                trials_with_status
            )

            # Make sure trials from other experiments were not deleted
            assert len(storage.fetch_trials(uid="other")) == 1

    def test_get_trial(self, storage):
        """Test get trial"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage = cfg.storage

            trial_dict = cfg.trials[0]

            trial1 = storage.get_trial(trial=Trial(**trial_dict))
            trial2 = storage.get_trial(uid=trial1.id, experiment_uid=trial1.experiment)

            with pytest.raises(MissingArguments):
                storage.get_trial()

            with pytest.raises(MissingArguments):
                storage.get_trial(uid=trial1.id)

            with pytest.raises(AssertionError):
                storage.get_trial(trial=trial1, uid="123")

            assert trial1.to_dict() == trial_dict
            assert trial2.to_dict() == trial_dict

    def test_fetch_lost_trials(self, storage):
        """Test update heartbeat"""
        lost_trials = [
            make_lost_trial(2),  # Is not lost for long enough to be caught
            make_lost_trial(10),  # Is lost for long enough to be caught
        ]
        # Force recent heartbeat to avoid mixing up with lost trials.
        trials = lost_trials + generate_trials(heartbeat=datetime.datetime.utcnow())
        with OrionState(
            experiments=[base_experiment], trials=trials, storage=storage
        ) as cfg:
            storage = cfg.storage

            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.fetch_lost_trials(experiment)

            assert len(trials) == 1

    def test_change_status_success(self, storage):
        """Change the status of a Trial"""

        def check_status_change(new_status):
            with OrionState(
                experiments=[base_experiment], trials=generate_trials(), storage=storage
            ) as cfg:
                trial = setup_storage().get_trial(cfg.get_trial(0))
                assert trial is not None, "was not able to retrieve trial for test"

                setup_storage().set_trial_status(trial, status=new_status)
                assert (
                    trial.status == new_status
                ), "Trial status should have been updated locally"

                trial = setup_storage().get_trial(trial)
                assert (
                    trial.status == new_status
                ), "Trial status should have been updated in the storage"

        check_status_change("completed")
        check_status_change("broken")
        check_status_change("reserved")
        check_status_change("interrupted")
        check_status_change("suspended")
        check_status_change("new")

    def test_change_status_invalid(self, storage):
        """Attempt to change the status of a Trial with an invalid one"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            trial = setup_storage().get_trial(cfg.get_trial(0))
            assert trial is not None, "Was not able to retrieve trial for test"

            with pytest.raises(ValueError) as exc:
                setup_storage().set_trial_status(trial, status="moo")

            assert exc.match("Given status `moo` not one of")

    def test_change_status_failed_update(self, storage):
        """Change the status of a Trial"""

        def check_status_change(new_status):
            with OrionState(
                experiments=[base_experiment], trials=generate_trials(), storage=storage
            ) as cfg:
                trial = cfg.storage.get_trial(cfg.get_trial(1))
                assert trial is not None, "Was not able to retrieve trial for test"
                assert trial.status != new_status

                if trial.status == new_status:
                    return

                with pytest.raises(FailedUpdate):
                    trial.status = new_status
                    setup_storage().set_trial_status(trial, status=new_status)

        check_status_change("completed")
        check_status_change("broken")
        check_status_change("reserved")
        check_status_change("interrupted")
        check_status_change("suspended")

    def test_change_status_success_thanks_to_was(self, storage):
        """Change the status of a Trial requesting the correct previous state, although local
        one is not up-to-date.
        """

        def check_status_change(new_status):
            with OrionState(
                experiments=[base_experiment], trials=generate_trials(), storage=storage
            ) as cfg:
                trial = cfg.storage.get_trial(cfg.get_trial(1))
                assert trial is not None, "Was not able to retrieve trial for test"
                assert trial.status != new_status

                if trial.status == new_status:
                    return

                correct_status = trial.status
                trial.status = "broken"
                assert correct_status != "broken"
                with pytest.raises(FailedUpdate):
                    setup_storage().set_trial_status(trial, status=new_status)

                setup_storage().set_trial_status(
                    trial, status=new_status, was=correct_status
                )

        check_status_change("completed")
        check_status_change("broken")
        check_status_change("reserved")
        check_status_change("interrupted")
        check_status_change("suspended")

    def test_change_status_failed_update_because_of_was(self, storage):
        """Change the status of a Trial requesting the wrong previous state."""

        def check_status_change(new_status):
            with OrionState(
                experiments=[base_experiment], trials=generate_trials(), storage=storage
            ) as cfg:
                trial = cfg.storage.get_trial(cfg.get_trial(1))
                assert trial is not None, "Was not able to retrieve trial for test"
                assert trial.status != new_status

                if trial.status == new_status:
                    return

                with pytest.raises(FailedUpdate):
                    setup_storage().set_trial_status(
                        trial, status=new_status, was=new_status
                    )

        check_status_change("completed")
        check_status_change("broken")
        check_status_change("reserved")
        check_status_change("interrupted")
        check_status_change("suspended")

    def test_fetch_pending_trials(self, storage):
        """Test fetch pending trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage = cfg.storage

            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.fetch_pending_trials(experiment)

            count = 0
            for trial in cfg.trials:
                if trial["status"] in {"new", "suspended", "interrupted"}:
                    count += 1

            assert len(trials) == count
            for trial in trials:
                assert trial.status in {"new", "suspended", "interrupted"}

    def test_fetch_noncompleted_trials(self, storage):
        """Test fetch non completed trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage = cfg.storage

            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.fetch_noncompleted_trials(experiment)

            count = 0
            for trial in cfg.trials:
                if trial["status"] != "completed":
                    count += 1

            for trial in trials:
                assert trial.status != "completed"

            assert len(trials) == count

    def test_fetch_trials_by_status(self, storage):
        """Test fetch completed trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            count = 0
            for trial in cfg.trials:
                if trial["status"] == "completed":
                    count += 1

            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.fetch_trials_by_status(experiment, "completed")

            assert len(trials) == count
            for trial in trials:
                assert trial.status == "completed", trial

    def test_count_completed_trials(self, storage):
        """Test count completed trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            count = 0
            for trial in cfg.trials:
                if trial["status"] == "completed":
                    count += 1

            storage = cfg.storage

            experiment = cfg.get_experiment("default_name", version=None)
            trials = storage.count_completed_trials(experiment)
            assert trials == count

    def test_count_broken_trials(self, storage):
        """Test count broken trials"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            count = 0
            for trial in cfg.trials:
                if trial["status"] == "broken":
                    count += 1

            storage = cfg.storage

            experiment = cfg.get_experiment("default_name", version=None)

            trials = storage.count_broken_trials(experiment)

            assert trials == count

    def test_update_heartbeat(self, storage):
        """Test update heartbeat"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage_name = storage
            storage = cfg.storage

            exp = cfg.get_experiment("default_name")
            trial1 = storage.fetch_trials_by_status(exp, status="reserved")[0]
            trial1b = copy.deepcopy(trial1)

            storage.update_heartbeat(trial1)

            trial2 = storage.get_trial(trial1)

            # this check that heartbeat is the correct type and that it was updated prior to now
            assert trial1b.heartbeat is None
            assert trial1.heartbeat is None
            assert trial2.heartbeat is not None

            # Sleep a bit, because fast CPUs make this test fail
            time.sleep(0.1)
            assert trial2.heartbeat < datetime.datetime.utcnow()

            if storage_name is None:
                trial3 = storage.fetch_trials_by_status(exp, status="completed")[0]
                storage.update_heartbeat(trial3)

                assert (
                    trial3.heartbeat is None
                ), "Legacy does not update trials with a status different from reserved"

    def test_serializable(self, storage):
        """Test storage can be serialized"""
        with OrionState(
            experiments=[base_experiment], trials=generate_trials(), storage=storage
        ) as cfg:
            storage = cfg.storage
            serialized = pickle.dumps(storage)
            deserialized = pickle.loads(serialized)
            assert storage.fetch_experiments({}) == deserialized.fetch_experiments({})

    def test_get_algorithm_lock_info(self, storage):
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support algorithm lock yet.")

        with OrionState(experiments=generate_experiments(), storage=storage) as cfg:
            storage = cfg.storage

            experiments = storage.fetch_experiments({})

            algo_state_lock = storage.get_algorithm_lock_info(uid=experiments[0]["_id"])
            assert isinstance(algo_state_lock, LockedAlgorithmState)
            assert algo_state_lock.state is None
            assert algo_state_lock.configuration == experiments[0]["algorithm"]

    def test_delete_algorithm_lock(self, storage):
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support algorithm lock yet.")

        with OrionState(experiments=generate_experiments(), storage=storage) as cfg:
            storage = cfg.storage

            experiments = storage.fetch_experiments({})

            assert storage.delete_algorithm_lock(uid=experiments[0]["_id"]) == 1
            assert storage.get_algorithm_lock_info(uid=experiments[0]["_id"]) is None

    def test_acquire_algorithm_lock_successful(self, storage):
        if storage and storage["type"] == "track":
            pytest.xfail("Track does not support algorithm lock yet.")

        with OrionState(experiments=[base_experiment], storage=storage) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            with storage.acquire_algorithm_lock(
                experiment, timeout=0.1
            ) as locked_algo_state:
                assert locked_algo_state.state is None
                locked_algo_state.set_state("my new state")

            with storage.acquire_algorithm_lock(experiment) as locked_algo_state:
                assert locked_algo_state.state == "my new state"

    def test_acquire_algorithm_lock_timeout(self, storage, mocker):
        with OrionState(experiments=[base_experiment], storage=storage) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)

            with storage.acquire_algorithm_lock(experiment) as locked_algo_state:

                retry_interval = 0.2
                sleep_mock = mocker.spy(time, "sleep")
                with pytest.raises(LockAcquisitionTimeout):
                    with storage.acquire_algorithm_lock(
                        experiment, timeout=0.05, retry_interval=retry_interval
                    ):
                        pass

                # TODO: Add back when PR #791 is merged
                # assert sleep_mock.call_count == 1
                # sleep_mock.assert_called_with(retry_interval)

    def test_acquire_algorithm_lock_handle_fail(self, storage):
        with OrionState(experiments=[base_experiment], storage=storage) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)
            with storage.acquire_algorithm_lock(experiment) as locked_algo_state:
                assert locked_algo_state.state is None
                locked_algo_state.set_state("new original state")

            with pytest.raises(RuntimeError):
                with storage.acquire_algorithm_lock(experiment) as locked_algo_state:
                    assert locked_algo_state.state == "new original state"
                    locked_algo_state.set_state("should not be set")
                    raise RuntimeError

            with storage.acquire_algorithm_lock(experiment) as locked_algo_state:
                assert locked_algo_state.state == "new original state"

    def test_acquire_algorithm_lock_not_initialised(self, storage):
        with OrionState(experiments=[base_experiment], storage=storage) as cfg:
            storage = cfg.storage
            experiment = cfg.get_experiment("default_name", version=None)
            experiment._id = "bad id"
            with pytest.raises(LockAcquisitionTimeout):
                with storage.acquire_algorithm_lock(experiment, timeout=0.1) as what:
                    pass
