Switching to test_contexts

Author: giadarol

tests/test_aperture_turn_ele_and_monitor.py

@@ -1,19 +1,14 @@
 import numpy as np
 
 import xobjects as xo
-from xobjects.context import available
 import xtrack as xt
 import xline as xl
 
 
 def test_aperture_turn_ele_and_monitor():
 
-    for CTX in xo.ContextCpu, xo.ContextPyopencl, xo.ContextCupy:
-        if CTX not in available:
-            continue
-
-        print(f"Test {CTX}")
-        context = CTX()
+    for context in xo.context.get_test_contexts():
+        print(f"Test {context.__class__}")
 
         x_aper_min = -0.1
         x_aper_max = 0.2
@@ -117,12 +112,8 @@ def test_aperture_turn_ele_and_monitor():
 
 def test_custom_monitor():
 
-    for CTX in xo.ContextCpu, xo.ContextPyopencl, xo.ContextCupy:
-        if CTX not in available:
-            continue
-
-        print(f"Test {CTX}")
-        context = CTX()
+    for context in xo.context.get_test_contexts():
+        print(f"Test {context.__class__}")
 
         x_aper_min = -0.1
         x_aper_max = 0.2

tests/test_collective_tracker.py

@@ -3,19 +3,14 @@
 import numpy as np
 
 import xobjects as xo
-from xobjects.context import available
 import xline as xl
 import xtrack as xt
 import xfields as xf
 
 def test_collective_tracker():
 
-    for CTX in xo.ContextCpu, xo.ContextPyopencl, xo.ContextCupy:
-        if CTX not in available:
-            continue
-
-        print(f"Test {CTX}")
-        context = CTX()
+    for context in xo.context.get_test_contexts():
+        print(f"Test {context.__class__}")
 
         test_data_folder = pathlib.Path(
             __file__).parent.joinpath('../test_data').absolute()

tests/test_dress.py

@@ -10,24 +10,28 @@ class ElementData(xo.Struct):
 
     class Element(dress(ElementData)):
 
-        def __init__(self, vv):
-            self.xoinitialize(n=len(vv), b=np.sum(vv), vv=vv)
-
-    ele = Element([1,2,3])
-    assert ele.n == ele._xobject.n == 3
-    assert ele.b == ele._xobject.b == 6
-    assert ele.vv[1] == ele._xobject.vv[1] == 2
-
-    ele.vv = [7,8,9]
-    assert ele.n == ele._xobject.n == 3
-    assert ele.b == ele._xobject.b == 6
-    assert ele.vv[1] == ele._xobject.vv[1] == 8
-
-    ele.n = 5.
-    assert ele.n == ele._xobject.n == 5
-
-    ele.b = 50
-    assert ele.b == ele._xobject.b == 50.
+        def __init__(self, vv, **kwargs):
+            self.xoinitialize(n=len(vv), b=np.sum(vv), vv=vv,
+                              **kwargs)
+    for context in xo.context.get_test_contexts():
+        print(f"Test {context.__class__}")
+
+        ele = Element([1,2,3], _context=context)
+        assert ele.n == ele._xobject.n == 3
+        assert ele.b == ele._xobject.b == 6
+        assert ele.vv[1] == ele._xobject.vv[1] == 2
+
+        new_vv = context.nparray_to_context_array(np.array([7,8,9]))
+        ele.vv = new_vv
+        assert ele.n == ele._xobject.n == 3
+        assert ele.b == ele._xobject.b == 6
+        assert ele.vv[1] == ele._xobject.vv[1] == 8
+
+        ele.n = 5.
+        assert ele.n == ele._xobject.n == 5
+
+        ele.b = 50
+        assert ele.b == ele._xobject.b == 50.
 
 
 def test_explicit_buffer():