From 2f4ea8641ec5c121aba429d4697a37ec6b736f09 Mon Sep 17 00:00:00 2001
From: Evan <evanev7@gmail.com>
Date: Mon, 22 Dec 2025 03:46:00 +0000
Subject: [PATCH] fix test

---
 src/exo/master/tests/conftest.py       | 12 ++----------
 src/exo/master/tests/test_placement.py | 12 ++++++------
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/src/exo/master/tests/conftest.py b/src/exo/master/tests/conftest.py
index e931d1a7..5830ce6b 100644
--- a/src/exo/master/tests/conftest.py
+++ b/src/exo/master/tests/conftest.py
@@ -1,7 +1,6 @@
 from ipaddress import ip_address
 from itertools import count
 
-from exo.routing.connection_message import SocketAddress
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
     MemoryPerformanceProfile,
@@ -11,7 +10,6 @@ from exo.shared.types.profiling import (
 from exo.shared.types.topology import Connection, ConnectionProfile, NodeInfo
 
 ip_octet_iter = count()
-port_iter = count(5000)
 
 
 def create_node(memory: int, node_id: NodeId | None = None) -> NodeInfo:
@@ -39,21 +37,15 @@ def create_connection(
     source_node_id: NodeId,
     sink_node_id: NodeId,
     *,
-    port: int | None = None,
     ip_octet: int | None = None,
 ) -> Connection:
     global ip_octet_iter
-    global port_iter
 
     return Connection(
         source_id=source_node_id,
         sink_id=sink_node_id,
-        sink_addr=SocketAddress(
-            ip=ip_address(
-                f"169.254.0.{ip_octet if ip_octet is not None else next(ip_octet_iter)}"
-            ),
-            port=port if port is not None else next(port_iter),
-            zone_id=None,
+        sink_addr=ip_address(
+            f"169.254.0.{ip_octet if ip_octet is not None else next(ip_octet_iter)}"
         ),
         connection_profile=ConnectionProfile(
             throughput=1000, latency=1000, jitter=1000
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 53b28aad..370850a0 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -351,11 +351,11 @@ def test_tensor_rdma_backend_connectivity_matrix(
         network_interfaces=[
             NetworkInterfaceInfo(
                 name="en3",
-                ip_address=conn_c_a.sink_addr.ip,
+                ip_address=conn_c_a.sink_addr,
             ),
             NetworkInterfaceInfo(
                 name="en4",
-                ip_address=conn_b_a.sink_addr.ip,
+                ip_address=conn_b_a.sink_addr,
             ),
             ethernet_interface,
         ],
@@ -369,11 +369,11 @@ def test_tensor_rdma_backend_connectivity_matrix(
         network_interfaces=[
             NetworkInterfaceInfo(
                 name="en3",
-                ip_address=conn_c_b.sink_addr.ip,
+                ip_address=conn_c_b.sink_addr,
             ),
             NetworkInterfaceInfo(
                 name="en4",
-                ip_address=conn_a_b.sink_addr.ip,
+                ip_address=conn_a_b.sink_addr,
             ),
             ethernet_interface,
         ],
@@ -387,11 +387,11 @@ def test_tensor_rdma_backend_connectivity_matrix(
         network_interfaces=[
             NetworkInterfaceInfo(
                 name="en3",
-                ip_address=conn_a_c.sink_addr.ip,
+                ip_address=conn_a_c.sink_addr,
             ),
             NetworkInterfaceInfo(
                 name="en4",
-                ip_address=conn_b_c.sink_addr.ip,
+                ip_address=conn_b_c.sink_addr,
             ),
             ethernet_interface,
         ],