diff --git a/src/node.py b/src/node.py index db01496..aece4ad 100644 --- a/src/node.py +++ b/src/node.py @@ -51,8 +51,7 @@ def send_ip(): # 用环境变量获取本机ip -def get_local_ip(): - global ip +def get_local_ip() -> str | None: ip = os.environ.get("HOST_IP") if not ip: # 如果环境变量中没有IP try: @@ -61,12 +60,16 @@ def get_local_ip(): s.connect(("8.8.8.8", 80)) # 通过连接Google DNS获取IP ip = str(s.getsockname()[0]) s.close() + return ip except IndexError: raise ValueError("Unable to get IP") + else: + return ip def init(): - get_local_ip() + global ip + ip = get_local_ip() send_ip() asyncio.create_task(send_heartbeat_internal()) print("Finish init") diff --git a/tests/node_test5.py b/tests/node_test5.py index 12aae1b..06e81b5 100644 --- a/tests/node_test5.py +++ b/tests/node_test5.py @@ -21,7 +21,6 @@ from node import ( init, clear, send_user_des_message, - ip, id, ) @@ -34,26 +33,23 @@ server_address = "http://60.204.236.38:8000/server" class TestGetLocalIP(unittest.TestCase): - os.environ["HOST_IP"] = "60.204.193.58" # 模拟设置 HOST_IP 环境变量 - def test_get_ip_from_env(self): - global ip - # 调用被测函数 - get_local_ip() + os.environ["HOST_IP"] = "60.204.193.58" # 模拟设置 HOST_IP 环境变量 + ip = get_local_ip() # 检查函数是否正确获取到 HOST_IP self.assertEqual(ip, "60.204.193.58") @patch("socket.socket") # Mock socket 连接行为 - @patch.dict("os.environ", {}) # 模拟没有 HOST_IP 环境变量 def test_get_ip_from_socket(self, mock_socket): - global ip + os.environ.pop("HOST_IP", None) + # 模拟 socket 返回的 IP 地址 mock_socket_instance = MagicMock() mock_socket.return_value = mock_socket_instance mock_socket_instance.getsockname.return_value = ("110.41.155.96", 0) # 调用被测函数 - get_local_ip() + ip = get_local_ip() # 确认 socket 被调用过 mock_socket_instance.connect.assert_called_with(("8.8.8.8", 80))