protocols: Give rust wallet types an alias, don't use `__all__` (#15186)

* protocols: Give rust wallet types an alias, don't use `__all__`

* Refactor `types_in_module` in `test_network_protocol_test.py`

chia/protocols/wallet_protocol.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import List, Optional, Tuple
 
-from chia_rs import CoinState, RespondToPhUpdates
+import chia_rs
 
 from chia.full_node.fee_estimate import FeeEstimateGroup
 from chia.types.blockchain_format.coin import Coin
@@ -20,7 +20,8 @@ Note: When changing this file, also change protocol_message_types.py, and the pr
 """
 
 
-__all__ = ["CoinState", "RespondToPhUpdates"]
+CoinState = chia_rs.CoinState
+RespondToPhUpdates = chia_rs.RespondToPhUpdates
 
 
 @streamable

tests/util/test_network_protocol_test.py

@@ -1,7 +1,9 @@
 # flake8: noqa
 from __future__ import annotations
 
-from typing import Any, List, Set
+import ast
+import inspect
+from typing import Any, Set, cast
 
 from chia.protocols import (
     farmer_protocol,
@@ -19,16 +21,16 @@ from chia.protocols import (
 
 
 def types_in_module(mod: Any) -> Set[str]:
-    ret: List[str] = []
-    mod_name = mod.__name__
-    for sym in dir(mod):
-        obj = getattr(mod, sym)
-        if hasattr(obj, "__module__") and obj.__module__ == mod_name:
-            ret.append(sym)
-
-    if hasattr(mod, "__all__"):
-        ret += getattr(mod, "__all__")
-    return set(ret)
+    parsed = ast.parse(inspect.getsource(mod))
+    types = set()
+    for line in parsed.body:
+        if isinstance(line, ast.Assign):
+            name = cast(ast.Name, line.targets[0])
+            if inspect.isclass(getattr(mod, name.id)):
+                types.add(name.id)
+        elif isinstance(line, ast.ClassDef):
+            types.add(line.name)
+    return types
 
 
 def test_missing_messages_state_machine() -> None:
@@ -155,8 +157,6 @@ def test_missing_messages() -> None:
         "PutFarmerPayload",
         "PutFarmerRequest",
         "PutFarmerResponse",
-        "get_current_authentication_token",
-        "validate_authentication_token",
     }
 
     timelord_msgs = {