1 files changed, 6 insertions, 7 deletions
diff --git a/Examples/Tests/pass_mpi_communicator/PICMI_inputs_2d.py b/Examples/Tests/pass_mpi_communicator/PICMI_inputs_2d.py
index 0d87699c5..540766152 100755
--- a/Examples/Tests/pass_mpi_communicator/PICMI_inputs_2d.py
+++ b/Examples/Tests/pass_mpi_communicator/PICMI_inputs_2d.py
@@ -6,6 +6,7 @@
 # --- if the correct amount of processors are initialized in AMReX.
 
 from mpi4py import MPI
+import pywarpx
 from pywarpx import picmi
 constants = picmi.constants
 
@@ -124,18 +125,16 @@ sim.step(max_steps, mpi_comm=new_comm)
 comm_world_size = comm_world.size
 new_comm_size = new_comm.size
 
-from pywarpx import wx
-
 if color == 0:
     # verify that communicator contains correct number of procs (1)
-    assert wx.libwarpx.warpx_getNProcs() == comm_world_size - 1
-    assert wx.libwarpx.warpx_getNProcs() == new_comm_size
+    assert pywarpx.getNProcs() == comm_world_size - 1
+    assert pywarpx.getNProcs() == new_comm_size
 
 else:
     # verify that amrex initialized with 1 fewer proc than comm world
-    assert wx.libwarpx.warpx_getNProcs() == comm_world_size - 1
-    assert wx.libwarpx.warpx_getNProcs() == new_comm_size
+    assert pywarpx.getNProcs() == comm_world_size - 1
+    assert pywarpx.getNProcs() == new_comm_size
 
     # verify that amrex proc ranks are offset by -1 from
     # world comm proc ranks
-    assert wx.libwarpx.warpx_getMyProc() == rank - 1
+    assert pywarpx.getMyProc() == rank - 1