Fixed bug when buying armor at base

gen/defenses/armor_group_generator.py

@@ -3,22 +3,38 @@ import random
 from dcs.vehicles import Armor
 
 from game import db
-from gen.defenses.armored_group_generator import ArmoredGroupGenerator
+from gen.defenses.armored_group_generator import ArmoredGroupGenerator, FixedSizeArmorGroupGenerator
 
 
 def generate_armor_group(faction:str, game, ground_object):
     """
     This generate a group of ground units
-    :param parentCp: The parent control point
-    :param ground_object: The ground object which will own the group
-    :param country: Owner country
     :return: Generated group
     """
 
     possible_unit = [u for u in db.FACTIONS[faction]["units"] if u in Armor.__dict__.values()]
     if len(possible_unit) > 0:
         unit_type = random.choice(possible_unit)
-        generator = ArmoredGroupGenerator(game, ground_object, unit_type)
-        generator.generate()
-        return generator.get_generated_group()
+        return generate_armor_group_of_type(game, ground_object, unit_type)
     return None
+
+
+def generate_armor_group_of_type(game, ground_object, unit_type):
+    """
+    This generate a group of ground units of given type
+    :return: Generated group
+    """
+    generator = ArmoredGroupGenerator(game, ground_object, unit_type)
+    generator.generate()
+    return generator.get_generated_group()
+
+
+def generate_armor_group_of_type_and_size(game, ground_object, unit_type, size: int):
+    """
+    This generate a group of ground units of given type and size
+    :return: Generated group
+    """
+    generator = FixedSizeArmorGroupGenerator(game, ground_object, unit_type, size)
+    generator.generate()
+    return generator.get_generated_group()
+

gen/defenses/armored_group_generator.py

@@ -25,3 +25,20 @@ class ArmoredGroupGenerator(GroupGenerator):
                               self.position.y + spacing * j, self.heading)
 
 
+class FixedSizeArmorGroupGenerator(GroupGenerator):
+
+    def __init__(self, game, ground_object, unit_type, size):
+        super(FixedSizeArmorGroupGenerator, self).__init__(game, ground_object)
+        self.unit_type = unit_type
+        self.size = size
+
+    def generate(self):
+        spacing = random.randint(20, 70)
+
+        index = 0
+        for i in range(self.size):
+            index = index + 1
+            self.add_unit(self.unit_type, "Armor#" + str(index),
+                            self.position.x + spacing * i,
+                            self.position.y, self.heading)
+