wip: add load_balance test

Author: triceo

tests/test_collectors.py

@@ -563,24 +563,39 @@ def define_constraints(constraint_factory: ConstraintFactory):
     assert score_manager.explain(problem).score == SimpleScore.of(4)
 
 
-def test_flatten_last():
+def test_load_balance():
     @constraint_provider
     def define_constraints(constraint_factory: ConstraintFactory):
         return [
             constraint_factory.for_each(Entity)
-            .map(lambda entity: (1, 2, 3))
-            .flatten_last(lambda the_tuple: the_tuple)
-            .reward(SimpleScore.ONE)
-            .as_constraint('Count')
+            .group_by(ConstraintCollectors.load_balance(
+                lambda entity: entity.value
+            ))
+            .reward(SimpleScore.ONE,
+                    lambda balance: balance.unfairness().multiply(BigDecimal.valueOf(1000)).intValue())
+            .as_constraint('Balanced value')
         ]
 
     score_manager = create_score_manager(define_constraints)
 
     entity_a: Entity = Entity('A')
+    entity_b: Entity = Entity('B')
+    entity_c: Entity = Entity('C')
 
     value_1 = Value(1)
+    value_2 = Value(2)
 
-    problem = Solution([entity_a], [value_1])
+    problem = Solution([entity_a, entity_b, entity_c], [value_1, value_2])
     entity_a.value = value_1
+    entity_b.value = value_1
+    entity_c.value = value_1
 
-    assert score_manager.explain(problem).score == SimpleScore.of(3)
+    assert score_manager.explain(problem).score == SimpleScore.of(0)
+
+    entity_c.value = value_2
+
+    assert score_manager.explain(problem).score == SimpleScore.of(2) // FIXME
+
+    entity_b.value = value_2
+
+    assert score_manager.explain(problem).score == SimpleScore.of(4) // FIXME

tests/test_constraint_streams.py

@@ -223,6 +223,29 @@ def define_constraints(constraint_factory: ConstraintFactory):
     assert score_manager.explain(problem).score.score == 1
 
 
+def test_flatten_last():
+    @constraint_provider
+    def define_constraints(constraint_factory: ConstraintFactory):
+        return [
+            constraint_factory.for_each(Entity)
+            .map(lambda entity: (1, 2, 3))
+            .flatten_last(lambda the_tuple: the_tuple)
+            .reward(SimpleScore.ONE)
+            .as_constraint('Count')
+        ]
+
+    score_manager = create_score_manager(define_constraints)
+
+    entity_a: Entity = Entity('A')
+
+    value_1 = Value(1)
+
+    problem = Solution([entity_a], [value_1])
+    entity_a.value = value_1
+
+    assert score_manager.explain(problem).score == SimpleScore.of(3)
+
+
 def test_join_uni():
     @constraint_provider
     def define_constraints(constraint_factory: ConstraintFactory):

timefold-solver-python-core/src/main/python/score/_group_by.py

@@ -122,18 +122,18 @@ def perform_group_by(constraint_stream, package, group_by_args, *type_arguments)
             created_collector = extract_collector(collector_info, *type_arguments)
             actual_group_by_args.append(created_collector)
 
-    if len(group_by_args) == 1:
+    if len(group_by_args) is 1:
         return UniConstraintStream(constraint_stream.groupBy(*actual_group_by_args), package,
                                    JClass('java.lang.Object'))
-    elif len(group_by_args) == 2:
+    elif len(group_by_args) is 2:
         return BiConstraintStream(constraint_stream.groupBy(*actual_group_by_args), package,
                                   JClass('java.lang.Object'),
                                   JClass('java.lang.Object'))
-    elif len(group_by_args) == 3:
+    elif len(group_by_args) is 3:
         return TriConstraintStream(constraint_stream.groupBy(*actual_group_by_args), package,
                                    JClass('java.lang.Object'),
                                    JClass('java.lang.Object'), JClass('java.lang.Object'))
-    elif len(group_by_args) == 4:
+    elif len(group_by_args) is 4:
         return QuadConstraintStream(constraint_stream.groupBy(*actual_group_by_args), package,
                                     JClass('java.lang.Object'),
                                     JClass('java.lang.Object'), JClass('java.lang.Object'),
@@ -1107,27 +1107,28 @@ def load_balance(balanced_item_function, load_function=None, initial_load_functi
 
         Parameters
         ----------
-        balanced_item_function:
+        balanced_item_function: Callable[[ParameterTypes, ...], Balanced_]
              The function that returns the item which should be load-balanced.
-        load_function:
+        load_function: Callable[[ParameterTypes, ...], int]
             How much the item should count for in the formula.
-        initial_load_function:
+        initial_load_function: Callable[[ParameterTypes, ...], int]
             The initial value of the metric, allowing to provide initial state
             without requiring the entire previous planning windows in the working memory.
             If this function is provided, load_function must be provided as well.
         """
-        if None == load_function == initial_load_function:
+        if load_function is None and initial_load_function is None:
             return LoadBalanceCollector(ConstraintCollectors._delegate().loadBalance, balanced_item_function, None,
                                         None)
-        elif None == initial_load_function:
+        elif initial_load_function is None:
             return LoadBalanceCollector(ConstraintCollectors._delegate().loadBalance, balanced_item_function,
                                         load_function, None)
-        elif None == load_function:
+        elif load_function is None:
             raise ValueError("load_function cannot be None if initial_load_function is not None")
         else:
             return LoadBalanceCollector(ConstraintCollectors._delegate().loadBalance, balanced_item_function,
                                         load_function, initial_load_function)
 
+
 # Must be at the bottom, constraint_stream depends on this module
 from ._constraint_stream import *
 from ._function_translator import *