bug correction -- fit logic

AGENTS.md

@@ -59,7 +59,7 @@ Each `process()` call passes through a state machine controlled by config fields
 
 Phases in order: **CONFIGURE → TRAIN → DETECT** (phases are skipped when the corresponding field is `None`).
 
-State control enums allow overriding automatic transitions: `TrainState.KEEP_TRAINING` / `STOP_TRAINING` and `ConfigState.KEEP_CONFIGURE` / `STOP_CONFIGURE` on the component's `fit_logic`.
+State control enums allow overriding automatic transitions: `TrainState.KEEP_TRAINING` / `STOP_TRAINING` and `EnumState.KEEP` / `STOP_CONFIGURE` on the component's `fit_logic`.
 
 #### CoreDetectorConfig: EventsConfig
 

src/detectmatelibrary/common/_core_op/_fit_logic.py

@@ -3,75 +3,60 @@
 from enum import Enum
 import warnings
 
+# Individual state of config or train ################################################
 
-class TrainState(Enum):
+class EnumState(Enum):
     DEFAULT = 0
-    STOP_TRAINING = 1
-    KEEP_TRAINING = 2
+    STOP = 1
+    KEEP = 2
 
     def describe(self) -> str:
         descriptions = [
-            "Follow default training behavior.",
-            "Force stop training.",
-            "Keep training regardless of default behavior."
+            "Follow default behavior.",
+            "Force stop",
+            "Keep doing it regardless of default behavior."
         ]
-
         return descriptions[self.value]
 
 
-class ConfigState(Enum):
-    DEFAULT = 0
-    STOP_CONFIGURE = 1
-    KEEP_CONFIGURE = 2
+class State:
+    def __init__(self, total_need_data: int | None) -> None:
+        self.total_need_data = total_need_data
 
-    def describe(self) -> str:
-        descriptions = [
-            "Follow default configuration behavior.",
-            "Force stop configuration.",
-            "Keep configuring regardless of default behavior."
-        ]
-        return descriptions[self.value]
+        self.ready_to_finish = False
+        self.finished = False
+        self.data_used = 0
+        self.current = EnumState.DEFAULT
 
+    def force_keep_doing(self) -> None:
+        self.current = EnumState.KEEP
 
-StatesL = Literal["keep_training", "stop_training", "keep_configuring", "stop_configuring"]
+    def keep_doing(self) -> bool:
+        if self.current == EnumState.STOP:
+            return False
+        if self.current == EnumState.KEEP:
+            return True
 
-def update_state(
-    state: StatesL, train_state: TrainState, config_state: ConfigState
-) -> tuple[TrainState, ConfigState]:
-    if state == "keep_training":
-        train_state = TrainState.KEEP_TRAINING
-    elif state == "stop_training":
-        train_state = TrainState.STOP_TRAINING
-    elif state == "keep_configuring":
-        config_state = ConfigState.KEEP_CONFIGURE
-    elif state == "stop_configuring":
-        config_state = ConfigState.STOP_CONFIGURE
-    else:
-        warnings.warn(f"State {state} unknown, use: {get_args(StatesL)}")
-
-    return train_state, config_state
-
-
-def do_training(
-    data_use_training: int | None, index: int, train_state: TrainState
-) -> bool:
-    if train_state == TrainState.STOP_TRAINING:
-        return False
-    elif train_state == TrainState.KEEP_TRAINING:
-        return True
+        return self.total_need_data is not None and self.total_need_data > self.data_used
 
-    return data_use_training is not None and data_use_training > index
+    def force_finish(self) -> None:
+        self.current = EnumState.STOP
+        self.finished, self.ready_to_finish = False, True
 
+    def check_if_ready_finish(self) -> None:
+        if self.data_used > 0 and not self.ready_to_finish:
+                self.ready_to_finish = True
 
-def do_configure(
-    data_use_configure: int | None, index: int, configure_state: ConfigState
-) -> bool:
-    if configure_state == ConfigState.STOP_CONFIGURE:
+    def is_finish(self) -> bool:
+        if self.ready_to_finish and not self.finished:
+            self.finished = True
+            return True
         return False
-    elif configure_state == ConfigState.KEEP_CONFIGURE:
-        return True
 
-    return data_use_configure is not None and data_use_configure > index
+
+# Overall state of the core component ##################################################
+
+StatesL = Literal["keep_training", "stop_training", "keep_configuring", "stop_configuring"]
 
 
 class FitLogicState(Enum):
@@ -93,58 +78,43 @@ def __init__(
         self, data_use_configure: int | None, data_use_training: int | None
     ) -> None:
 
-        self.train_state = TrainState.DEFAULT
-        self.configure_state = ConfigState.DEFAULT
         self.last_state = FitLogicState.NOTHING
 
-        self.data_used_train, self.data_used_configure = 0, 0
-        self._configuration_done, self.config_finished = False, False
-        self._training_done, self.training_finished = False, False
-
-        self.data_use_configure = data_use_configure
-        self.data_use_training = data_use_training
+        self.config_state = State(data_use_configure)
+        self.train_state = State(data_use_training)
 
     def get_last_state(self) -> str:
         return self.last_state.describe()
 
     def update_state(self, state: StatesL) -> None:
-        self.train_state, self.configure_state = update_state(
-            state=state, train_state=self.train_state, config_state=self.configure_state
-        )
+        if state == "keep_training":
+            self.train_state.force_keep_doing()
+        elif state == "stop_training":
+            self.train_state.force_finish()
+        elif state == "keep_configuring":
+            self.config_state.force_keep_doing()
+        elif state == "stop_configuring":
+            self.config_state.force_finish()
+        else:
+            warnings.warn(f"State {state} unknown, use: {get_args(StatesL)}")
 
     def finish_config(self) -> bool:
-        if self._configuration_done and not self.config_finished:
-            self.config_finished = True
-            return True
-        return False
+        return self.config_state.is_finish()
 
     def finish_training(self) -> bool:
-        if self._training_done and not self.training_finished:
-            self.training_finished = True
-            return True
-        return False
+        return self.train_state.is_finish()
 
     def __check_state(self) -> FitLogicState:
-        if do_configure(
-            data_use_configure=self.data_use_configure,
-            index=self.data_used_configure,
-            configure_state=self.configure_state
-        ):
-            self.data_used_configure += 1
+        if self.config_state.keep_doing():
+            self.config_state.data_used += 1
             return FitLogicState.DO_CONFIG
         else:
-            if self.data_used_configure > 0 and not self._configuration_done:
-                self._configuration_done = True
-
-            if do_training(
-                data_use_training=self.data_use_training,
-                index=self.data_used_train,
-                train_state=self.train_state
-            ):
-                self.data_used_train += 1
+            self.config_state.check_if_ready_finish()
+
+            if self.train_state.keep_doing():
+                self.train_state.data_used += 1
                 return FitLogicState.DO_TRAIN
-            elif self.data_used_train > 0 and not self._training_done:
-                self._training_done = True
+            self.train_state.check_if_ready_finish()
 
         return FitLogicState.NOTHING
 

tests/test_common/test_core.py

@@ -241,7 +241,7 @@ def test_training(self) -> None:
                 })
             )
 
-        assert len(component.train_data) == component.fitlogic.data_used_train
+        assert len(component.train_data) == component.fitlogic.train_state.data_used
         for i, log in enumerate(component.train_data):
             expected = schemas.LogSchema({
                 "__version__": "1.0.0",
@@ -266,7 +266,8 @@ def test_training_use_config_data(self) -> None:
                     "hostname": "test_hostname"
                 })
             )
-        total = component.fitlogic.data_use_training + component.fitlogic.data_use_configure
+        total = component.fitlogic.train_state.total_need_data
+        total += component.fitlogic.config_state.total_need_data
         assert len(component.train_data) == total
         for i, log in enumerate(component.train_data):
             expected = schemas.LogSchema({
@@ -315,7 +316,7 @@ def test_configuration(self) -> None:
 
         results = [component.process(_make_log(i)) for i in range(10)]
 
-        assert component.fitlogic.data_used_configure == 3
+        assert component.fitlogic.config_state.data_used == 3
         assert len(component.configure_data) == 3
         assert all(r is None for r in results[:3])
         assert component.set_configuration_called == 1
@@ -340,7 +341,7 @@ def test_configuration_force_stop(self) -> None:
             component.process(_make_log(i))
 
         assert len(component.configure_data) == 0
-        assert component.set_configuration_called == 0
+        assert component.set_configuration_called == 1
 
     def test_configuration_keep_configure(self) -> None:
         component = MockComponentWithConfigure(name="DummyCfg4")

tests/test_common/test_fit_logic.py

@@ -1,7 +1,7 @@
 """Tests for FitLogic training lifecycle hooks."""
 
 from detectmatelibrary.common._core_op._fit_logic import (
-    FitLogic, FitLogicState, TrainState
+    FitLogic, FitLogicState, EnumState
 )
 
 
@@ -46,7 +46,7 @@ def test_finish_training_not_called_during_training(self) -> None:
 
     def test_finish_training_not_called_with_keep_training(self) -> None:
         logic = FitLogic(data_use_configure=None, data_use_training=None)
-        logic.train_state = TrainState.KEEP_TRAINING
+        logic.train_state.current = EnumState.KEEP
         for _ in range(10):
             state = logic.run()
             assert state == FitLogicState.DO_TRAIN
@@ -60,3 +60,33 @@ def test_finish_training_after_configure_and_training(self) -> None:
             logic.run()
             finish_calls.append(logic.finish_training())
         assert finish_calls.count(True) == 1
+
+    def test_force_stop_config(self) -> None:
+        logic = FitLogic(data_use_configure=20, data_use_training=None)
+
+        # Stop in the middle
+        assert not logic.finish_config()
+        logic.update_state("stop_configuring")
+        assert logic.finish_config()
+        assert not logic.finish_config()
+
+        # Stop after a force start 
+        logic.update_state("keep_configuring")
+        assert not logic.finish_config()
+        logic.update_state("stop_configuring")
+        assert logic.finish_config()
+
+    def test_force_stop_train(self) -> None:
+        logic = FitLogic(data_use_configure=None, data_use_training=20)
+
+        # Stop in the middle
+        assert not logic.finish_training()
+        logic.update_state("stop_training")
+        assert logic.finish_training()
+        assert not logic.finish_training()
+
+        # Stop after a force start 
+        logic.update_state("keep_training")
+        assert not logic.finish_training()
+        logic.update_state("stop_training")
+        assert logic.finish_training()

tests/test_detectors/test_bigram_frequency_detector.py

@@ -8,11 +8,10 @@
 - Input/output schema validation
 """
 
-from detectmatelibrary.common._core_op._fit_logic import TrainState
 from detectmatelibrary.detectors.bigram_frequency_detector import (
     BigramFrequencyDetector, BigramFrequencyDetectorConfig, BufferMode
 )
-from detectmatelibrary.common._core_op._fit_logic import ConfigState
+from detectmatelibrary.common._core_op._fit_logic import EnumState
 from detectmatelibrary.constants import GLOBAL_EVENT_ID
 from detectmatelibrary.parsers.template_matcher import MatcherParser
 from detectmatelibrary.helper.from_to import From
@@ -273,20 +272,20 @@ def test_audit_log_anomalies_via_process(self):
         logs = list(From.log(parser, in_path=AUDIT_LOG, do_process=True))
 
         # Phase 1: configure — keep configuring for logs[:TRAIN_UNTIL]
-        detector.fitlogic.configure_state = ConfigState.KEEP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Transition: stop configure so next process() call triggers set_configuration()
-        detector.fitlogic.configure_state = ConfigState.STOP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.STOP
 
         # Phase 2: train — keep training for logs[:TRAIN_UNTIL]
-        detector.fitlogic.train_state = TrainState.KEEP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Phase 3: detect — stop training so process() only calls detect()
-        detector.fitlogic.train_state = TrainState.STOP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.STOP
         detected_ids: set[str] = set()
         for log in logs[TRAIN_UNTIL:]:
             if detector.process(log) is not None:

tests/test_detectors/test_charset_detector.py

@@ -8,9 +8,8 @@
 - Input/output schema validation
 """
 
-from detectmatelibrary.common._core_op._fit_logic import TrainState
 from detectmatelibrary.detectors.charset_detector import CharsetDetector, CharsetDetectorConfig, BufferMode
-from detectmatelibrary.common._core_op._fit_logic import ConfigState
+from detectmatelibrary.common._core_op._fit_logic import EnumState
 from detectmatelibrary.constants import GLOBAL_EVENT_ID
 from detectmatelibrary.parsers.template_matcher import MatcherParser
 from detectmatelibrary.helper.from_to import From
@@ -322,20 +321,20 @@ def test_audit_log_anomalies_via_process(self):
         logs = list(From.log(parser, in_path=AUDIT_LOG, do_process=True))
 
         # Phase 1: configure — keep configuring for logs[:TRAIN_UNTIL]
-        detector.fitlogic.configure_state = ConfigState.KEEP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Transition: stop configure so next process() call triggers set_configuration()
-        detector.fitlogic.configure_state = ConfigState.STOP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.STOP
 
         # Phase 2: train — keep training for logs[:TRAIN_UNTIL]
-        detector.fitlogic.train_state = TrainState.KEEP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Phase 3: detect — stop training so process() only calls detect()
-        detector.fitlogic.train_state = TrainState.STOP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.STOP
         detected_ids: set[str] = set()
         for log in logs[TRAIN_UNTIL:]:
             if detector.process(log) is not None:

tests/test_detectors/test_new_event_detector.py

@@ -14,7 +14,7 @@
 from detectmatelibrary.helper.from_to import From
 import detectmatelibrary.schemas as schemas
 from detectmatelibrary.utils.aux import time_test_mode
-from detectmatelibrary.common._core_op._fit_logic import ConfigState, TrainState
+from detectmatelibrary.common._core_op._fit_logic import EnumState
 from detectmatelibrary.constants import GLOBAL_EVENT_ID
 from tests.test_data import AUDIT_LOG, AUDIT_TEMPLATES, TRAIN_UNTIL
 
@@ -187,20 +187,20 @@ def test_audit_log_anomalies_via_process(self):
         logs = list(From.log(pars, in_path=AUDIT_LOG, do_process=True))
 
         # Phase 1: configure — keep configuring for logs[:TRAIN_UNTIL]
-        detector.fitlogic.configure_state = ConfigState.KEEP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Transition: stop configure so next process() call triggers set_configuration()
-        detector.fitlogic.configure_state = ConfigState.STOP_CONFIGURE
+        detector.fitlogic.config_state.current = EnumState.STOP
 
         # Phase 2: train — keep training for logs[:TRAIN_UNTIL]
-        detector.fitlogic.train_state = TrainState.KEEP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.KEEP
         for log in logs[:TRAIN_UNTIL]:
             detector.process(log)
 
         # Phase 3: detect — stop training so process() only calls detect()
-        detector.fitlogic.train_state = TrainState.STOP_TRAINING
+        detector.fitlogic.train_state.current = EnumState.STOP
         detected_ids: set[str] = set()
         for log in logs[TRAIN_UNTIL:]:
             if detector.process(log) is not None: