Corrected tests

Author: softwaremagico

think-machine-random/src/main/java/com/softwaremagico/tm/random/RandomSelector.java

@@ -163,7 +163,7 @@ private void assignMandatory() throws InvalidXmlElementException {
 
     private TreeMap<Integer, Element> assignElementsWeight() throws InvalidXmlElementException {
         final TreeMap<Integer, Element> weightedElements = new TreeMap<>();
-        int count = 0;
+        int count = 1;
         for (final Element element : getAllElements()) {
             try {
                 validateElement(element);
@@ -178,9 +178,6 @@ private TreeMap<Integer, Element> assignElementsWeight() throws InvalidXmlElemen
                 count += weight;
             }
         }
-        if (!weightedElements.isEmpty()) {
-            weightedElements.put(count, null);
-        }
         return weightedElements;
     }
 
@@ -355,7 +352,7 @@ protected Element selectElementByWeight() throws InvalidRandomElementSelectedExc
             throw new InvalidRandomElementSelectedException("No elements to select");
         }
         final int value = RANDOM.nextInt(totalWeight) + 1;
-        Element selectedElement = weightedElements.values().iterator().next();
+        Element selectedElement;
         final SortedMap<Integer, Element> view = weightedElements.headMap(value, true);
         try {
             selectedElement = view.get(view.lastKey());
@@ -400,22 +397,20 @@ protected void updateWeight(Element element, int newWeight) {
         weightedElements.put(newWeight, element);
     }
 
-    public SortedMap<Integer, Element> getWeightedElements() {
+    public TreeMap<Integer, Element> getWeightedElements() {
         return weightedElements;
     }
 
     public Integer getAssignedWeight(Element element) {
         if (element == null) {
             return null;
         }
-        Integer elementWeight = null;
-        for (final Map.Entry<Integer, Element> entry : weightedElements.entrySet()) {
-            if (Objects.equals(entry.getValue(), element)) {
-                elementWeight = entry.getKey();
-                continue;
+        int previousWeight = 0;
+        for (final Entry<Integer, Element> entry : weightedElements.entrySet()) {
+            if (entry.getValue().equals(element)) {
+                return entry.getKey() - previousWeight;
             }
-            if (elementWeight != null)
-                return entry.getKey() - elementWeight;
+            previousWeight = entry.getKey();
         }
         return null;
     }