From 447887645ed7dbf2682b3130e9b8cc3656a35dc4 Mon Sep 17 00:00:00 2001
From: Amelia Trainer <ameliajo@mit.edu>
Date: Thu, 25 Mar 2021 22:42:04 +0000
Subject: [PATCH] added tests for collision filter

---
 openmc/lib/filter.py                   |  2 +-
 src/tallies/filter_collision.cpp       | 10 ++++------
 tests/regression_tests/tallies/test.py |  9 ++++++++-
 tests/unit_tests/test_filters.py       | 16 ++++++++++++++++
 4 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/openmc/lib/filter.py b/openmc/lib/filter.py
index 9c1eea4da5..59897b50fa 100644
--- a/openmc/lib/filter.py
+++ b/openmc/lib/filter.py
@@ -176,7 +176,7 @@ class EnergyFilter(Filter):
             self._index, len(energies), energies_p)
 
 
-class CollisionFilterFilter(CollisionFilter):
+class CollisionFilter(Filter):
     filter_type = 'collision'
 
 
diff --git a/src/tallies/filter_collision.cpp b/src/tallies/filter_collision.cpp
index 60229941a6..17350acc26 100644
--- a/src/tallies/filter_collision.cpp
+++ b/src/tallies/filter_collision.cpp
@@ -42,12 +42,10 @@ void CollisionFilter::get_all_bins(
   auto n = p.n_collision_;
 
   // Bin the collision number. Must fit exactly the desired collision number .
-  if (n >= bins_.front() && n <= bins_.back()) {
-    auto search = map_.find(n);
-    if (search != map_.end()){
-      match.bins_.push_back(search->second);
-      match.weights_.push_back(1.0);
-    }
+  auto search = map_.find(n);
+  if (search != map_.end()){
+    match.bins_.push_back(search->second);
+    match.weights_.push_back(1.0);
   }
 }
 
diff --git a/tests/regression_tests/tallies/test.py b/tests/regression_tests/tallies/test.py
index d4857b3ec7..8e1e8b60b9 100644
--- a/tests/regression_tests/tallies/test.py
+++ b/tests/regression_tests/tallies/test.py
@@ -121,6 +121,12 @@ def test_tallies():
     harmonics_tally3.scores = ['flux', 'total']
     harmonics_tally3.estimatir = 'tracklength'
 
+    n_collision = (1,2,5,3,6)
+    collision_filter = CollisionFilter(n_collision)
+    collision_tally = Tally()
+    collision_tally .filters = [collision_filter]
+    collision_tally .scores = ['scatter']
+
     universe_tally = Tally()
     universe_tally.filters = [
         UniverseFilter((model.geometry.get_all_universes()[1],
@@ -174,7 +180,8 @@ def test_tallies():
         cellborn_tally, dg_tally, energy_tally, energyout_tally,
         transfer_tally, material_tally, mu_tally1, mu_tally2,
         polar_tally1, polar_tally2, polar_tally3, legendre_tally,
-        harmonics_tally, harmonics_tally2, harmonics_tally3, universe_tally]
+        harmonics_tally, harmonics_tally2, harmonics_tally3, 
+        universe_tally,collision_tally]
     model.tallies += score_tallies
     model.tallies += flux_tallies
     model.tallies += all_nuclide_tallies
diff --git a/tests/unit_tests/test_filters.py b/tests/unit_tests/test_filters.py
index bf62abb65d..d52775abd2 100644
--- a/tests/unit_tests/test_filters.py
+++ b/tests/unit_tests/test_filters.py
@@ -46,6 +46,22 @@ def test_cell_instance():
     assert instances.apply(lambda x: x in (0, 1, 2)).all()
 
 
+def test_collision():
+    f = openmc.CollisionFilter([1,5,3,2,8])
+    assert f.bins[0] == 1
+    assert f.bins[1] == 5
+    assert f.bins[-1] == 8
+    assert len(f.bins) == 5
+
+    # Make sure __repr__ works
+    repr(f)
+
+    # to_xml_element()
+    elem = f.to_xml_element()
+    assert elem.tag == 'filter'
+    assert elem.attrib['type'] == 'collision'
+
+
 def test_legendre():
     n = 5
     f = openmc.LegendreFilter(n)