How to use the faculty.clients.experiment._client.ExperimentClient function in faculty

def test_delete_runs(mocker):
    delete_runs_response = mocker.Mock()
    mocker.patch.object(
        ExperimentClient, "_post", return_value=delete_runs_response
    )
    response_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.DeleteExperimentRunsResponseSchema"
    )
    filter_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.FilterSchema"
    )
    filter_dump_mock = filter_schema_mock.return_value.dump

    run_ids = [uuid4(), uuid4()]

    client = ExperimentClient(mocker.Mock())
    response = client.delete_runs(PROJECT_ID, run_ids)

    assert response == delete_runs_response

    expected_filter = CompoundFilter(
        LogicalOperator.OR,
        [
            RunIdFilter(ComparisonOperator.EQUAL_TO, run_ids[0]),
            RunIdFilter(ComparisonOperator.EQUAL_TO, run_ids[1]),
        ],
    )
    filter_dump_mock.assert_called_once_with(expected_filter)
    ExperimentClient._post.assert_called_once_with(
        "/project/{}/run/delete/query".format(PROJECT_ID),
        response_schema_mock.return_value,
        json={"filter": filter_dump_mock.return_value},

def test_experiment_client_list_runs(mocker):
    mocker.patch.object(ExperimentClient, "query_runs")

    client = ExperimentClient(mocker.Mock())
    response = client.list_runs(
        PROJECT_ID,
        experiment_ids=[123, 456],
        lifecycle_stage=LifecycleStage.DELETED,
        start=20,
        limit=10,
    )

    assert response == ExperimentClient.query_runs.return_value
    expected_filter = CompoundFilter(
        LogicalOperator.AND,
        [
            CompoundFilter(
                LogicalOperator.OR,
                [
                    ExperimentIdFilter(ComparisonOperator.EQUAL_TO, 123),

def test_experiment_client_update(mocker, name, description):
    mocker.patch.object(ExperimentClient, "_patch_raw")

    client = ExperimentClient(mocker.Mock())
    client.update(
        PROJECT_ID, EXPERIMENT_ID, name=name, description=description
    )

    ExperimentClient._patch_raw.assert_called_once_with(
        "/project/{}/experiment/{}".format(PROJECT_ID, EXPERIMENT_ID),
        json={"name": name, "description": description},
    )

def test_log_run_data(mocker):
    mocker.patch.object(ExperimentClient, "_patch_raw")
    run_data_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.ExperimentRunDataSchema"
    )
    run_data_dump_mock = run_data_schema_mock.return_value.dump

    metric = mocker.Mock()
    param = mocker.Mock()
    tag = mocker.Mock()

    client = ExperimentClient(mocker.Mock())
    client.log_run_data(
        PROJECT_ID,
        EXPERIMENT_RUN_ID,
        metrics=[metric],
        params=[param],
        tags=[tag],
    )

    run_data_schema_mock.assert_called_once_with()
    run_data_dump_mock.assert_called_once_with(
        {"metrics": [metric], "params": [param], "tags": [tag]}
    )
    ExperimentClient._patch_raw.assert_called_once_with(
        "/project/{}/run/{}/data".format(PROJECT_ID, EXPERIMENT_RUN_ID),
        json=run_data_dump_mock.return_value,
    )

def test_log_run_data_empty(mocker):
    mocker.patch.object(ExperimentClient, "_patch_raw")

    client = ExperimentClient(mocker.Mock())

    client.log_run_data(PROJECT_ID, EXPERIMENT_RUN_ID)
    ExperimentClient._patch_raw.assert_not_called()

def test_experiment_client_query_runs(mocker):
    list_response = mocker.Mock()
    mocker.patch.object(ExperimentClient, "_post", return_value=list_response)
    response_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.ListExperimentRunsResponseSchema"
    )
    request_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.RunQuerySchema"
    )
    request_dump_mock = request_schema_mock.return_value.dump

    filter = mocker.Mock()
    sort = mocker.Mock()

    client = ExperimentClient(mocker.Mock())
    list_result = client.query_runs(
        PROJECT_ID, filter, sort, start=20, limit=10
    )

    assert list_result == list_response

    request_dump_mock.assert_called_once_with(
        RunQuery(filter, sort, Page(20, 10))
    )
    response_schema_mock.assert_called_once_with()
    ExperimentClient._post.assert_called_once_with(
        "/project/{}/run/query".format(PROJECT_ID),
        response_schema_mock.return_value,
        json=request_dump_mock.return_value,
    )

def test_log_run_data_other_conflict(mocker):
    response_mock = mocker.Mock()
    exception = Conflict(response_mock, "", "")

    mocker.patch.object(ExperimentClient, "_patch_raw", side_effect=exception)
    client = ExperimentClient(mocker.Mock())

    with pytest.raises(Conflict):
        client.log_run_data(
            PROJECT_ID, EXPERIMENT_RUN_ID, params=[mocker.Mock()]
        )

def test_experiment_create_run(mocker):
    run = mocker.Mock()
    mocker.patch.object(ExperimentClient, "_post", return_value=run)
    request_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.CreateRunSchema"
    )
    dump_mock = request_schema_mock.return_value.dump
    response_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.ExperimentRunSchema"
    )
    run_name = mocker.Mock()
    started_at = mocker.Mock()
    artifact_location = mocker.Mock()

    client = ExperimentClient(mocker.Mock())
    returned_run = client.create_run(
        PROJECT_ID,
        EXPERIMENT_ID,
        run_name,
        started_at,
        PARENT_RUN_ID,
        artifact_location=artifact_location,
    )
    assert returned_run == run

    request_schema_mock.assert_called_once_with()
    dump_mock.assert_called_once_with(
        {
            "name": run_name,
            "parent_run_id": PARENT_RUN_ID,
            "started_at": started_at,

restore_runs_response = mocker.Mock()
    mocker.patch.object(
        ExperimentClient, "_post", return_value=restore_runs_response
    )
    response_schema_mock = mocker.patch(
        "faculty.clients.experiment._client."
        "RestoreExperimentRunsResponseSchema"
    )
    filter_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.FilterSchema"
    )
    filter_dump_mock = filter_schema_mock.return_value.dump

    run_ids = [uuid4(), uuid4()]

    client = ExperimentClient(mocker.Mock())
    response = client.restore_runs(PROJECT_ID, run_ids)

    assert response == restore_runs_response

    expected_filter = CompoundFilter(
        LogicalOperator.OR,
        [
            RunIdFilter(ComparisonOperator.EQUAL_TO, run_ids[0]),
            RunIdFilter(ComparisonOperator.EQUAL_TO, run_ids[1]),
        ],
    )
    filter_dump_mock.assert_called_once_with(expected_filter)
    ExperimentClient._post.assert_called_once_with(
        "/project/{}/run/restore/query".format(PROJECT_ID),
        response_schema_mock.return_value,
        json={"filter": filter_dump_mock.return_value},

def test_update_run_info(mocker):
    run = mocker.Mock()
    mocker.patch.object(ExperimentClient, "_patch", return_value=run)
    run_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.ExperimentRunSchema"
    )
    run_info_schema_mock = mocker.patch(
        "faculty.clients.experiment._client.ExperimentRunInfoSchema"
    )
    run_info_dump_mock = run_info_schema_mock.return_value.dump

    status = mocker.Mock()
    ended_at = mocker.Mock()

    client = ExperimentClient(mocker.Mock())
    returned_run = client.update_run_info(
        PROJECT_ID, EXPERIMENT_RUN_ID, status, ended_at
    )
    assert returned_run == run

    run_schema_mock.assert_called_once_with()
    run_info_schema_mock.assert_called_once_with()
    run_info_dump_mock.assert_called_once_with(
        {"status": status, "ended_at": ended_at}
    )
    ExperimentClient._patch.assert_called_once_with(
        "/project/{}/run/{}/info".format(PROJECT_ID, EXPERIMENT_RUN_ID),
        run_schema_mock.return_value,
        json=run_info_dump_mock.return_value,
    )