Repository: qdrant/qdrant
File: tests/consensus_tests/utils.py (Python)

Code before:
        except requests.exceptions.ConnectionError:
            print(f"Can't retrieve cluster info for offline peer {uri}")
    return max_peer_url


def get_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}/cluster", headers=headers)
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_shard_transfer_count(peer_api_uri: str, collection_name: str) -> int:
    info = get_collection_cluster_info(peer_api_uri, collection_name)
    return len(info["shard_transfers"])


def get_collection_info(peer_api_uri: str, collection_name: str) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}")
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_collection_point_count(peer_api_uri: str, collection_name: str, exact: bool = False,
                               shard_key: Optional[str] = None, filter: Optional[str] = None) -> int:
    r = requests.post(f"{peer_api_uri}/collections/{collection_name}/points/count", json={"exact": exact, "shard_key": shard_key, "filter": filter})
    assert_http_ok(r)
    res = r.json()["result"]["count"]
    return res


def print_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}):
    print(json.dumps(get_collection_cluster_info(peer_api_uri, collection_name, headers=headers), indent=4))


def get_leader(peer_api_uri: str, headers={}) -> str:
    r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
    assert_http_ok(r)
    return r.json()["result"]["raft_info"]["leader"]


def check_leader(peer_api_uri: str, expected_leader: str, headers={}) -> bool:
    try:
        r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
        assert_http_ok(r)
        leader = r.json()["result"]["raft_info"]["leader"]
        correct_leader = leader == expected_leader
        if not correct_leader:
            print(f"Cluster leader invalid for peer {peer_api_uri} {leader}/{expected_leader}")
        return correct_leader

Review comment (suggestion):
```suggestion
                               shard_key: Optional[str] = None, filter: Optional[dict] = None) -> int:
```

Code after:
        try:
            # do not crash if the peer is not online
            peer_id = get_cluster_info(uri)['peer_id']
            if peer_id > max_peer_id:
                max_peer_id = peer_id
                max_peer_url = uri
        except requests.exceptions.ConnectionError:
            print(f"Can't retrieve cluster info for offline peer {uri}")
    return max_peer_url


def get_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}/cluster", headers=headers)
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_shard_transfer_count(peer_api_uri: str, collection_name: str) -> int:
    info = get_collection_cluster_info(peer_api_uri, collection_name)
    return len(info["shard_transfers"])


def get_collection_info(peer_api_uri: str, collection_name: str) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}")
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_collection_point_count(peer_api_uri: str, collection_name: str, exact: bool = False,
                               shard_key: Optional[str] = None, filter: Optional[dict] = None) -> int:
    r = requests.post(f"{peer_api_uri}/collections/{collection_name}/points/count", json={"exact": exact, "shard_key": shard_key, "filter": filter})
    assert_http_ok(r)
    res = r.json()["result"]["count"]
    return res


def print_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}):
    print(json.dumps(get_collection_cluster_info(peer_api_uri, collection_name, headers=headers), indent=4))


def get_leader(peer_api_uri: str, headers={}) -> str:
    r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
    assert_http_ok(r)
    return r.json()["result"]["raft_info"]["leader"]


def check_leader(peer_api_uri: str, expected_leader: str, headers={}) -> bool:
    try:
        r = requests.get(f"{peer_api_uri}/cluster", headers=headers)