diff --git a/tests/test_inverse_kinematics.py b/tests/test_inverse_kinematics.py
index 95c006a..2659f16 100644
--- a/tests/test_inverse_kinematics.py
+++ b/tests/test_inverse_kinematics.py
@@ -24,16 +24,27 @@ def make_transparent(link):
     for link in visual_data:
         make_transparent(link)
 
-
-def test_jacobian_follower():
+def create_test_chain(robot="kuka_iiwa", device="cpu"):
+    if robot == "kuka_iiwa":
+        urdf = "kuka_iiwa/model.urdf"
+        search_path = pybullet_data.getDataPath()
+        full_urdf = os.path.join(search_path, urdf)
+        chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
+        chain = chain.to(device=device)
+    elif robot == "widowx":
+        urdf = "widowx/wx250s.urdf"
+        full_urdf = urdf
+        chain = pk.build_serial_chain_from_urdf(open(full_urdf, "rb").read(), "ee_gripper_link")
+        chain = chain.to(device=device)
+    else:
+        raise NotImplementedError(f"Robot {robot} not implemented")
+    return chain, urdf
+
+def test_jacobian_follower(robot="kuka_iiwa"):
     pytorch_seed.seed(2)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    # device = "cpu"
-    urdf = "kuka_iiwa/model.urdf"
     search_path = pybullet_data.getDataPath()
-    full_urdf = os.path.join(search_path, urdf)
-    chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
-    chain = chain.to(device=device)
+    chain, urdf = create_test_chain(robot=robot, device=device)
 
     # robot frame
     pos = torch.tensor([0.0, 0.0, 0.0], device=device)
@@ -45,7 +56,7 @@ def test_jacobian_follower():
     # generate random goal joint angles (so these are all achievable)
     # use the joint limits to generate random joint angles
     lim = torch.tensor(chain.get_joint_limits(), device=device)
-    goal_q = torch.rand(M, 7, device=device) * (lim[1] - lim[0]) + lim[0]
+    goal_q = torch.rand(M, lim.shape[1], device=device) * (lim[1] - lim[0]) + lim[0]
 
     # get ee pose (in robot frame)
     goal_in_rob_frame_tf = chain.forward_kinematics(goal_q)
@@ -159,16 +170,11 @@ def test_jacobian_follower():
             p.stepSimulation()
 
 
-def test_ik_in_place_no_err():
+def test_ik_in_place_no_err(robot="kuka_iiwa"):
     pytorch_seed.seed(2)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     # device = "cpu"
-    urdf = "kuka_iiwa/model.urdf"
-    search_path = pybullet_data.getDataPath()
-    full_urdf = os.path.join(search_path, urdf)
-    chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
-    chain = chain.to(device=device)
-
+    chain, urdf = create_test_chain(robot=robot, device=device)
     # robot frame
     pos = torch.tensor([0.0, 0.0, 0.0], device=device)
     rot = torch.tensor([0.0, 0.0, 0.0], device=device)
@@ -176,7 +182,7 @@ def test_ik_in_place_no_err():
 
     # goal equal to current configuration
     lim = torch.tensor(chain.get_joint_limits(), device=device)
-    cur_q = torch.rand(7, device=device) * (lim[1] - lim[0]) + lim[0]
+    cur_q = torch.rand(lim.shape[1], device=device) * (lim[1] - lim[0]) + lim[0]
     M = 1
     goal_q = cur_q.unsqueeze(0).repeat(M, 1)
 
@@ -209,5 +215,9 @@ def test_ik_in_place_no_err():
 
 
 if __name__ == "__main__":
-    test_jacobian_follower()
-    test_ik_in_place_no_err()
\ No newline at end of file
+    print("Testing kuka_iiwa IK")
+    test_jacobian_follower(robot="kuka_iiwa")
+    test_ik_in_place_no_err(robot="kuka_iiwa")
+    print("Testing widowx IK")
+    test_jacobian_follower(robot="widowx")
+    test_ik_in_place_no_err(robot="widowx")
\ No newline at end of file