test: increase test coverage for prepare_beam_unpol

src/matvis/core/beams.py

@@ -94,22 +94,6 @@ def _wrangle_beams(
     if polarized:
         if any(b.beam_type != "efield" for b in beam_list):
             raise ValueError("beam type must be efield if using polarized=True")
-    else:
-        # The following applies if we're not polarized
-        for b in beam_list:
-            if b.beam_type != "power":
-                raise ValueError(
-                    f"beam type must be power if polarized=False. Have beam_type={b.beam_type}"
-                )
-            if b.Npols > 1:
-                raise ValueError(
-                    f"beam type must be power and have only one pol if polarized=False. Have {b.Npols}"
-                )
-
-            if b.polarization_array[0] not in [-5, -6]:
-                raise ValueError(
-                    f"beam type must be power and have only one pol (either xx or yy) if polarized=False. Have {b.polarization_array[0]}"
-                )
 
     return beam_list, nbeam, beam_idx
 

tests/test_beams.py

@@ -21,6 +21,12 @@ def efield_beam(uvbeam):
     return BeamInterface(uvbeam)
 
 
+@pytest.fixture(scope="module")
+def efield_beam_single_feed(efield_beam):
+    """An e-field uvbeam."""
+    return efield_beam.with_feeds(["x"])
+
+
 @pytest.fixture(scope="function")
 def efield_single_freq(uvbeam):
     """Single frequency beam."""
@@ -92,6 +98,7 @@ class TestPrepareBeamUnpolarized:
             lf("power_beam"),
             lf("efield_analytic_beam"),
             lf("power_analytic_beam"),
+            lf("efield_beam_single_feed"),
         ],
     )
     def test_different_input_beams(self, beam):
@@ -103,6 +110,13 @@ def test_different_input_beams(self, beam):
         assert len(new_beam.beam.polarization_array) == 1
         assert polnum2str(new_beam.beam.polarization_array[0]).lower() == "xx"
 
+    def test_noop(self, power_beam):
+        """Test that passing in a power beam with a single pol is a no-op."""
+        new_beam = power_beam.with_feeds(["x"])
+
+        same_beam = prepare_beam_unpolarized(new_beam)
+        assert same_beam is new_beam
+
 
 class TestUVBeamInterpolator:
     """Test the UVBeamInterpolator."""