1 files changed, 7 insertions, 3 deletions
diff --git a/src/modules/record_sheet.js b/src/modules/record_sheet.js
index 5f14852..0e26155 100644
--- a/src/modules/record_sheet.js
+++ b/src/modules/record_sheet.js
@@ -63,15 +63,19 @@ function createWeaponIcon(type) {
 
 function createRecord(unit) {
   const { dataset: { allegiance, number, squad }} = unit,
-    primaryWeapon = unit.querySelector('.primary-weapon'),
-    pw = primaryWeapon.getAttributeNS(null, 'href').split('#').pop() || 'rifle',
+
+    pw = unit.classList && unit.classList.contains('counter')
+      ? unit.querySelector('.primary-weapon').getAttributeNS(null, 'href').split('#').pop()
+      : unit.dataset.weapon || 'rifle',
+
     div = document.createElement('soldier-record-block'),
     spans = Array(6).fill('span').map(t => document.createElement(t)),
     [tn, sn, pwt, pwd, pwrs, pwrl] = spans;
 
   div.classList.add('soldier-record');
-  if (unit.classList.contains('selected')) div.classList.add('selected');
+  if (unit.classList && unit.classList.contains('selected')) div.classList.add('selected');
   div.dataset.number = number;
+  div.dataset.squad = squad;
   div.dataset.allegiance = allegiance;
 
   tn.setAttribute('slot', 'troop-number');