Fix SignalsAPI.stake_get (KeyError on removed totalStake field)

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'])

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)