diff --git a/numerapi/signalsapi.py b/numerapi/signalsapi.py index c698422..c491926 100644 --- a/numerapi/signalsapi.py +++ b/numerapi/signalsapi.py @@ -284,18 +284,29 @@ def ticker_universe(self) -> List[str]: path = self.download_dataset("signals/v1.0/live.parquet") return pd.read_parquet(path).numerai_ticker.tolist() - def stake_get(self, username) -> decimal.Decimal: - """get current stake for a given users + def stake_get(self, username) -> decimal.Decimal | None: + """get current stake for a given user Args: username (str) Returns: - decimal.Decimal: current stake + decimal.Decimal or None: current stake, or None if the model has + no stake Example: >>> SignalsAPI().stake_get("uuazed") Decimal('14.63') """ - data = self.public_user_profile(username) - return data['totalStake'] + query = """ + query($model_name: String! + $tournament: Int) { + v3UserProfile(model_name: $model_name + tournament: $tournament) { + stakeValue + } + } + """ + arguments = {'model_name': username, 'tournament': self.tournament_id} + data = self.raw_query(query, arguments)['data']['v3UserProfile'] + return utils.parse_float_string(data['stakeValue']) diff --git a/tests/test_signalsapi.py b/tests/test_signalsapi.py index 0ccaf0f..dd1db70 100644 --- a/tests/test_signalsapi.py +++ b/tests/test_signalsapi.py @@ -1,3 +1,6 @@ +import decimal +from unittest.mock import patch + import pytest import responses @@ -13,6 +16,27 @@ def api_fixture(): return api +@patch("numerapi.signalsapi.SignalsAPI.raw_query") +def test_stake_get(mocked, api): + mocked.return_value = {"data": {"v3UserProfile": {"stakeValue": "14.63"}}} + + stake = api.stake_get("uuazed") + + assert stake == decimal.Decimal("14.63") + args, _ = mocked.call_args + # current stake lives in `stakeValue`; `totalStake` no longer exists + assert "stakeValue" in args[0] + assert "totalStake" not in args[0] + assert args[1]["tournament"] == api.tournament_id == 11 + + +@patch("numerapi.signalsapi.SignalsAPI.raw_query") +def test_stake_get_no_stake(mocked, api): + # a model with no stake returns null -> None, not a KeyError + mocked.return_value = {"data": {"v3UserProfile": {"stakeValue": None}}} + assert api.stake_get("uuazed") is None + + @pytest.mark.live_api def test_get_leaderboard(api): lb = api.get_leaderboard(1)