better mass access functions which improved testing code

4 files changed, 60 insertions, 54 deletions

diff --git a/src/mass.rs b/src/mass.rs
index 8d9f9fc..7cc7d70 100644
--- a/src/mass.rs
+++ b/src/mass.rs
@@ -272,4 +272,40 @@ impl Mass {
             }
         }
     }
+
+    pub fn take_item(&mut self, item_type: ItemType) -> Option<Item> {
+        match self.mass_type {
+            MassType::Ship {
+                ref mut storage, ..
+            } => storage.take_item(item_type),
+            MassType::Astroid {
+                ref mut resources, ..
+            } => resources.take_item(item_type),
+            _ => None,
+        }
+    }
+
+    pub fn give_item(&mut self, item: Item) -> bool {
+        match self.mass_type {
+            MassType::Ship {
+                ref mut storage, ..
+            } => storage.give_item(item),
+            MassType::Astroid {
+                ref mut resources, ..
+            } => resources.give_item(item),
+            _ => false,
+        }
+    }
+
+    pub fn item_count(&self, item_type: ItemType) -> usize {
+        match &self.mass_type {
+            MassType::Ship {
+                storage, ..
+            } => storage.item_count(item_type),
+            MassType::Astroid {
+                resources, ..
+            } => resources.item_count(item_type),
+            _ => 0,
+        }
+    }
 }
diff --git a/src/modules/construction.rs b/src/modules/construction.rs
index 0eec155..091ebc6 100644
--- a/src/modules/construction.rs
+++ b/src/modules/construction.rs
@@ -33,7 +33,7 @@ impl Construction {
         masses: &mut HashMap<String, Mass>,
         storage: &mut Storage,
     ) {
-        if !self.has_enough(storage) {
+        if storage.item_count(ItemType::Iron) < constants::SHIP_CONSTRUCTION_IRON_COST {
             self.off();
         }
         if let Some(timer) = self.start {
@@ -43,9 +43,10 @@ impl Construction {
             }
         }
         if self.status == Status::Constructed {
-            storage
-                .take_items(ItemType::Iron, constants::SHIP_CONSTRUCTION_IRON_COST)
-                .unwrap();
+            for _ in 0..constants::SHIP_CONSTRUCTION_IRON_COST {
+                storage.take_item(ItemType::Iron).unwrap();
+            }
+
             masses.insert(
                 "Station".to_string(),
                 Mass::new_station(
@@ -60,7 +61,8 @@ impl Construction {
 
     pub fn get_client_data(&self, storage: &Storage) -> String {
         let client_data = ClientData {
-            has_enough: self.has_enough(storage),
+            has_enough: storage.item_count(ItemType::Iron)
+                >= constants::SHIP_CONSTRUCTION_IRON_COST,
             status: self.status.clone(),
         };
         serde_json::to_string(&client_data).unwrap() + "\n"
@@ -72,15 +74,6 @@ impl Construction {
         }
     }
 
-    fn has_enough(&self, storage: &Storage) -> bool {
-        storage
-            .items
-            .iter()
-            .filter(|item| item.item_type == ItemType::Iron)
-            .count()
-            >= constants::SHIP_CONSTRUCTION_IRON_COST
-    }
-
     fn toggle(&mut self) {
         match self.status {
             Status::None => self.on(),
diff --git a/src/modules/mining.rs b/src/modules/mining.rs
index 816313e..01d7663 100644
--- a/src/modules/mining.rs
+++ b/src/modules/mining.rs
@@ -42,23 +42,18 @@ impl Mining {
             }
         }
         if self.status == Status::Mined {
-            if let MassType::Astroid {
-                ref mut resources, ..
-            } = target.mass_type
-            {
-                match resources.take_item(ItemType::CrudeMinerals) {
-                    Some(item) => {
-                        if !storage.give_item(item.clone()) {
-                            let mass = Mass::new_item(
-                                item.clone(),
-                                target.position.clone(),
-                                target.velocity.clone(),
-                            );
-                            masses.insert(item.name.clone(), mass);
-                        }
+            match target.take_item(ItemType::CrudeMinerals) {
+                Some(item) => {
+                    if !storage.give_item(item.clone()) {
+                        let mass = Mass::new_item(
+                            item.clone(),
+                            target.position.clone(),
+                            target.velocity.clone(),
+                        );
+                        masses.insert(item.name.clone(), mass);
                     }
-                    None => self.off(),
                 }
+                None => self.off(),
             }
             self.mined();
         }
diff --git a/src/storage.rs b/src/storage.rs
index 5cbd55e..fbb99b1 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -35,31 +35,6 @@ impl Storage {
         }
     }
 
-    pub fn take_items(&mut self, item_type: ItemType, count: usize) -> Option<Vec<Item>> {
-        if self
-            .items
-            .iter()
-            .filter(|item| item.item_type == item_type)
-            .count()
-            >= count
-        {
-            let mut items = Vec::new();
-            for _ in 0..count {
-                let index = self
-                    .items
-                    .iter()
-                    .position(|item| item.item_type == item_type)
-                    .unwrap();
-                let item = self.items.remove(index);
-                self.carrying -= item.size;
-                items.push(item);
-            }
-            Some(items)
-        } else {
-            None
-        }
-    }
-
     pub fn give_item(&mut self, item: Item) -> bool {
         if self.capacity >= self.carrying + item.size {
             self.carrying += item.size;
@@ -69,4 +44,11 @@ impl Storage {
             false
         }
     }
+
+    pub fn item_count(&self, item_type: ItemType) -> usize {
+        self.items
+            .iter()
+            .filter(|item| item.item_type == item_type)
+            .count()
+    }
 }