diff --git a/t/test_decay.cc b/t/test_decay.cc
index 38cd4b7..bc4cae3 100644
--- a/t/test_decay.cc
+++ b/t/test_decay.cc
@@ -1,244 +1,264 @@
 /**
  *  \authors   The HEJ collaboration (see AUTHORS for details)
  *  \date      2019-2020
  *  \copyright GPLv2 or later
  *
  *  \brief Test classification for (invalid) W decays
  */
 #include "hej_test.hh"
 
 #include <array>
 #include <cstdlib>
 #include <initializer_list>
 #include <iostream>
 #include <memory>
 #include <string>
 #include <utility>
 
 #include "fastjet/JetDefinition.hh"
 #include "fastjet/PseudoJet.hh"
 
 #include "HEJ/Event.hh"
 #include "HEJ/event_types.hh"
 #include "HEJ/Particle.hh"
 #include "HEJ/PDG_codes.hh"
 
 namespace {
   const fastjet::JetDefinition jet_def{fastjet::JetAlgorithm::antikt_algorithm, 0.4};
   const double min_jet_pt{30.};
 
   HEJ::Event::EventData new_event() {
     HEJ::Event::EventData ev;
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -11,  -96,  -76,  123}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -15,  -70,  -22,   75}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, {  68,   93,  -20,  117}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -12,   95,   56,  111}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -30,  -22,   25,   45}, {}});
     ev.incoming[0] = {HEJ::ParticleID::gluon, {   0,    0, -254,  254}, {}};
     ev.incoming[1] = {HEJ::ParticleID::gluon, {   0,    0,  217,  217}, {}};
     return ev;
   }
 
   bool test_event(HEJ::Event::EventData data, bool const valid
   ){
     using namespace HEJ::event_type;
     EventType const expected{ valid?FKL:bad_final_state };
     shuffle_particles(data);
     auto const ev = std::move(data).cluster(jet_def, min_jet_pt);
     if(ev.type() != expected){
       std::cerr << "Event does not match expectation, expected "
         << name(expected) << "\n" << ev << std::endl;
       return false;
     }
     return true;
   }
-} // namespace anonymous
 
-int main() {
-  using namespace HEJ::pid;
-  auto ev = new_event();
-  // basic FKL
-  test_event(ev, true);
-  // W position shouldn't matter
-  for(auto const W_type: {Wp, Wm}){
-    for(std::size_t w_pos = 1; w_pos<ev.outgoing.size()-1; ++w_pos){
-      ev = new_event();
-      ev.outgoing[w_pos].type = W_type;
-      ev.outgoing.back().type = (W_type==Wp)?d:u;
-      ev.incoming.back().type = (W_type==Wp)?u:d;
+  // Check basic FKL event
+  bool check_fkl() {
+    auto ev = new_event();
+    return test_event(ev, true);
+  }
+
+  // Check W decays
+  bool check_W_decay() {
+    using namespace HEJ::pid;
+
+    auto ev = new_event();
+    // W position shouldn't matter
+    for(auto const W_type: {Wp, Wm}){
+      for(std::size_t w_pos = 1; w_pos<ev.outgoing.size()-1; ++w_pos){
+        ev = new_event();
+        ev.outgoing[w_pos].type = W_type;
+        ev.outgoing.back().type = (W_type==Wp)?d:u;
+        ev.incoming.back().type = (W_type==Wp)?u:d;
+
+        // no decay
+        if(!test_event(ev, false))
+          return false;
+
+        // working decay
+        ev.decays[w_pos] = decay_W(ev.outgoing[w_pos]);
+        if(!test_event(ev, true))
+          return false;
+
+        // swapped W+ <-> W-
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        -ev.decays[w_pos].at(0).type );
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        -ev.decays[w_pos].at(1).type );
+        if(!test_event(ev, false))
+          return false;
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        -ev.decays[w_pos].at(0).type );
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        -ev.decays[w_pos].at(1).type );
+
+        // replace e -> mu (normal)
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type+2 );
+        if(!test_event(ev, false))
+          return false;
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type-2 );
+
+        // replace e -> mu (anti)
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(1).type-2 );
+        if(!test_event(ev, false))
+          return false;
+
+        // all mu
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type+2 );
+        if(!test_event(ev, true))
+          return false;
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type-2 );
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(1).type+2 );
+
+        // partonic
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type-10 );
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(1).type+10 );
+        if(!test_event(ev, false))
+          return false;
+        ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(0).type+10 );
+        ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
+                                        ev.decays[w_pos].at(1).type-10 );
+
+        // double check that we undid all changes
+        if(!test_event(ev, true))
+          return false;
+
+        // 1->3 decay
+        ev.decays[w_pos].emplace_back(
+            HEJ::Particle{photon, fastjet::PseudoJet(0,0,0,0), {}}
+        );
+        if(!test_event(ev, false))
+          return false;
+        ev.decays[w_pos].pop_back();
+
+        // unexpected decay
+        ev.decays[0] = decay_W(ev.outgoing[0]);
+        ev.decays[0].at(0).type = ev.outgoing[0].type;
+        ev.decays[0].at(1).type = gluon;
+        if(!test_event(ev, false))
+          return false;
+      }
+    }
+    return true;
+  }
+
+  // Check Z decays
+  bool check_Z_decay() {
+    using namespace HEJ::pid;
+    auto ev = new_event();
+    for(size_t z_pos = 1; z_pos<ev.outgoing.size()-1; ++z_pos){
+      auto ev = new_event();
+      ev.outgoing[z_pos].type = Z_photon_mix;
+      ev.outgoing.back().type = u;
+      ev.incoming.back().type = u;
 
       // no decay
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
+        return false;
 
       // working decay
-      ev.decays[w_pos] = decay_W(ev.outgoing[w_pos]);
+      ev.decays[z_pos] = decay_Z(ev.outgoing[z_pos]);
       if(!test_event(ev, true))
-        return EXIT_FAILURE;
+        return false;
 
-      // swapped W+ <-> W-
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      -ev.decays[w_pos].at(0).type );
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      -ev.decays[w_pos].at(1).type );
-      if(!test_event(ev, false))
-        return EXIT_FAILURE;
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      -ev.decays[w_pos].at(0).type );
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      -ev.decays[w_pos].at(1).type );
-
-      // replace e -> mu (normal)
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type+2 );
+      // replace e- -> mu-
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type+2 );
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type-2 );
+        return false;
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type-2 );
 
-      // replace e -> mu (anti)
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(1).type-2 );
+      // replace e+ -> mu+
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type-2 );
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
+        return false;
 
       // all mu
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type+2 );
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type+2 );
       if(!test_event(ev, true))
-        return EXIT_FAILURE;
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type-2 );
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(1).type+2 );
+        return false;
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type-2 );
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type+2 );
+
+      // replace e- -> nu_e
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type+1 );
+      if(!test_event(ev, false))
+        return false;
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type-1 );
+
+      // replace e+ -> nu_e_bar
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type-1 );
+      if(!test_event(ev, false))
+        return false;
+
+      // neutrino-antineutrino
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type+1 );
+      if(!test_event(ev, false))
+        return false;
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type-1 );
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type+1 );
 
       // partonic
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type-10 );
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(1).type+10 );
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type-10 );
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type+10 );
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
-      ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(0).type+10 );
-      ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
-                                      ev.decays[w_pos].at(1).type-10 );
+        return false;
+      ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(0).type+10 );
+      ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
+                                      ev.decays[z_pos].at(1).type-10 );
 
       // double check that we undid all changes
       if(!test_event(ev, true))
-        return EXIT_FAILURE;
+        return false;
 
       // 1->3 decay
-      ev.decays[w_pos].emplace_back(
+      ev.decays[z_pos].emplace_back(
           HEJ::Particle{photon, fastjet::PseudoJet(0,0,0,0), {}}
       );
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
-      ev.decays[w_pos].pop_back();
+        return false;
+      ev.decays[z_pos].pop_back();
 
-      // unexpected decay
-      ev.decays[0] = decay_W(ev.outgoing[0]);
+      // second decay
+      ev.decays[0] = decay_Z(ev.outgoing[0]);
       ev.decays[0].at(0).type = ev.outgoing[0].type;
       ev.decays[0].at(1).type = gluon;
       if(!test_event(ev, false))
-        return EXIT_FAILURE;
+        return false;
     }
+    return true;
   }
+} // namespace anonymous
 
-  for(size_t z_pos = 1; z_pos<ev.outgoing.size()-1; ++z_pos){
-    ev = new_event();
-    ev.outgoing[z_pos].type = Z_photon_mix;
-    ev.outgoing.back().type = u;
-    ev.incoming.back().type = u;
-
-    // no decay
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-
-    // working decay
-    ev.decays[z_pos] = decay_Z(ev.outgoing[z_pos]);
-    if(!test_event(ev, true))
-      return EXIT_FAILURE;
-
-    // replace e- -> mu-
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type+2 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type-2 );
-
-    // replace e+ -> mu+
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type-2 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-
-    // all mu
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type+2 );
-    if(!test_event(ev, true))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type-2 );
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type+2 );
-
-    // replace e- -> nu_e
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type+1 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type-1 );
-
-    // replace e+ -> nu_e_bar
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type-1 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-
-    // neutrino-antineutrino
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type+1 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type-1 );
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type+1 );
-
-    // partonic
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type-10 );
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type+10 );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(0).type+10 );
-    ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
-                                    ev.decays[z_pos].at(1).type-10 );
-
-    // double check that we undid all changes
-    if(!test_event(ev, true))
-      return EXIT_FAILURE;
-
-    // 1->3 decay
-    ev.decays[z_pos].emplace_back(
-        HEJ::Particle{photon, fastjet::PseudoJet(0,0,0,0), {}}
-    );
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-    ev.decays[z_pos].pop_back();
-
-    // second decay
-    ev.decays[0] = decay_Z(ev.outgoing[0]);
-    ev.decays[0].at(0).type = ev.outgoing[0].type;
-    ev.decays[0].at(1).type = gluon;
-    if(!test_event(ev, false))
-      return EXIT_FAILURE;
-  }
+
+int main() {
+  if(!check_fkl()) return EXIT_FAILURE;
+  if(!check_W_decay()) return EXIT_FAILURE;
+  if(!check_Z_decay()) return EXIT_FAILURE;
 
   return EXIT_SUCCESS;
 }