feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'

tests/test_search_agent.py

@@ -0,0 +1,40 @@
+"""
+Unit tests for SearchAgent.
+"""
+
+import torch
+import pytest
+from src.search_agent import SearchAgent, PolicyValueNet
+from src.utils import encode_state, get_actions, step
+
+
+def test_policy_value_net_forward():
+    net = PolicyValueNet(input_dim=1, action_space=2)
+    x = torch.tensor([[3.0]])
+    policy, value = net(x)
+    assert policy.shape == (1, 2)
+    assert value.shape == (1, 1)
+
+
+def test_search_agent_action_selection():
+    net = PolicyValueNet(input_dim=1, action_space=2)
+    agent = SearchAgent(policy_value_net=net, max_depth=2)
+    # Start from state 0; actions are 1 and 2
+    action = agent.act(0)
+    assert action in [1, 2]
+
+
+def test_search_agent_value_estimation():
+    net = PolicyValueNet(input_dim=1, action_space=2)
+    agent = SearchAgent(policy_value_net=net, max_depth=3)
+    # For state 8, the optimal action is 10 (reward 1)
+    action = agent.act(8)
+    assert action == 10 or action == 9  # depending on policy, 10 is better
+
+
+def test_search_agent_terminal_state():
+    net = PolicyValueNet(input_dim=1, action_space=2)
+    agent = SearchAgent(policy_value_net=net, max_depth=1)
+    # State 10 is terminal; agent should return 10
+    action = agent.act(10)
+    assert action == 10