From 7f0e28905c6f364f1e7a39649e117c54a7d0bbd7 Mon Sep 17 00:00:00 2001
From: David Waterman <dgwaterman@gmail.com>
Date: Wed, 2 Aug 2023 17:43:17 +0100
Subject: [PATCH] Fixes for tests

---
 src/dxtbx/model/beam.py  | 6 ++++--
 tests/model/test_beam.py | 7 +++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/dxtbx/model/beam.py b/src/dxtbx/model/beam.py
index 10741b433..5ae08b3d5 100644
--- a/src/dxtbx/model/beam.py
+++ b/src/dxtbx/model/beam.py
@@ -8,9 +8,9 @@
 import libtbx.phil
 
 try:
-    from ..dxtbx_model_ext import Beam, Probe, PolychromaticBeam
+    from ..dxtbx_model_ext import Beam, PolychromaticBeam, Probe
 except ModuleNotFoundError:
-    from dxtbx_model_ext import Beam, Probe, PolychromaticBeam  # type: ignore
+    from dxtbx_model_ext import Beam, PolychromaticBeam, Probe  # type: ignore
 
 Vec3Float = Tuple[float, float, float]
 
@@ -181,6 +181,7 @@ def make_polychromatic_beam(
         polarization_fraction: float = 0.5,
         flux: float = 0.0,
         transmission: float = 1.0,
+        probe=Probe.xray,
         deg: bool = True,
     ) -> PolychromaticBeam:
         return PolychromaticBeam(
@@ -191,6 +192,7 @@ def make_polychromatic_beam(
             float(polarization_fraction),
             float(flux),
             float(transmission),
+            probe,
             bool(deg),
         )
 
diff --git a/tests/model/test_beam.py b/tests/model/test_beam.py
index 2f83a790c..48a6eb5a1 100644
--- a/tests/model/test_beam.py
+++ b/tests/model/test_beam.py
@@ -6,7 +6,7 @@
 from scitbx import matrix
 
 from dxtbx.model import Beam, PolychromaticBeam
-from dxtbx.model.beam import BeamFactory, beam_phil_scope
+from dxtbx.model.beam import BeamFactory, Probe, beam_phil_scope
 
 
 def test_setting_direction_and_wavelength():
@@ -222,6 +222,7 @@ def test_make_polychromatic_beam():
     polarization_fraction = 0.65
     transmission = 0.5
     flux = 0.75
+    probe = Probe.neutron
 
     beam = BeamFactory.make_polychromatic_beam(
         direction=direction,
@@ -231,6 +232,7 @@ def test_make_polychromatic_beam():
         polarization_fraction=polarization_fraction,
         transmission=transmission,
         flux=flux,
+        probe=probe,
     )
 
     assert beam.get_sample_to_source_direction() == pytest.approx((0.0, 0.0, 1.0))
@@ -240,6 +242,7 @@ def test_make_polychromatic_beam():
     assert beam.get_polarization_fraction() == pytest.approx(0.65)
     assert beam.get_transmission() == pytest.approx(0.5)
     assert beam.get_flux() == pytest.approx(0.75)
+    assert beam.get_probe() == Probe.neutron
 
 
 def test_polychromatic_beam_wavelength_guards():
@@ -266,5 +269,5 @@ def test_polychromatic_beam_str():
     beam = PolychromaticBeam()
     assert (
         beam.__str__()
-        == "Beam:\n    sample to source direction : {0,0,1}\n    divergence: 0\n    sigma divergence: 0\n    polarization normal: {0,1,0}\n    polarization fraction: 0.5\n    flux: 0\n    transmission: 1\n"
+        == "Beam:\n    probe: x-ray\n    sample to source direction : {0,0,1}\n    divergence: 0\n    sigma divergence: 0\n    polarization normal: {0,1,0}\n    polarization fraction: 0.5\n    flux: 0\n    transmission: 1\n"
     )