diff --git a/examples/ditto_example/client_dynamic.py b/examples/ditto_example/client_dynamic.py
index 217bef249..1c2805f0f 100644
--- a/examples/ditto_example/client_dynamic.py
+++ b/examples/ditto_example/client_dynamic.py
@@ -12,7 +12,7 @@
 from torch.utils.data import DataLoader
 
 from examples.models.cnn_model import MnistNet
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.mixins.personalized import PersonalizedMode, make_it_personal
 from fl4health.reporting import JsonReporter
 from fl4health.utils.config import narrow_dict_type
@@ -22,8 +22,8 @@
 from fl4health.utils.sampler import DirichletLabelBasedSampler
 
 
-class MnistClient(BasicClient):
-    """A simple `BasicClient` type that we dynamically personalize via Ditto."""
+class MnistClient(FlexibleClient):
+    """A simple `FlexibleClient` type that we dynamically personalize via Ditto."""
 
     def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
         sample_percentage = narrow_dict_type(config, "downsampling_ratio", float)
diff --git a/fl4health/clients/flexible_client.py b/fl4health/clients/flexible_client.py
new file mode 100644
index 000000000..bc4dd0f69
--- /dev/null
+++ b/fl4health/clients/flexible_client.py
@@ -0,0 +1,326 @@
+import warnings
+from collections.abc import Sequence
+from logging import WARN
+from pathlib import Path
+from typing import Any
+
+import torch
+from flwr.common.logger import log
+from torch import nn
+from torch.optim import Optimizer
+
+from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule
+from fl4health.clients.basic_client import BasicClient
+from fl4health.metrics.base_metrics import Metric
+from fl4health.reporting.base_reporter import BaseReporter
+from fl4health.utils.losses import EvaluationLosses, LossMeterType, TrainingLosses
+from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType
+
+
+class FlexibleClient(BasicClient):
+    def __init__(
+        self,
+        data_path: Path,
+        metrics: Sequence[Metric],
+        device: torch.device,
+        loss_meter_type: LossMeterType = LossMeterType.AVERAGE,
+        checkpoint_and_state_module: ClientCheckpointAndStateModule | None = None,
+        reporters: Sequence[BaseReporter] | None = None,
+        progress_bar: bool = False,
+        client_name: str | None = None,
+    ) -> None:
+        """
+        Flexible FL Client with functionality to train, evaluate, log, report and checkpoint.
+
+        `FlexibleClient` is similar to `BasicClient` but provides added flexibility through the
+        ability to inject models and optimizers in the methods responsible for making predictions
+        and performing both train and validation steps.
+
+        This added flexibility allows for `FlexibleClient` to be automatically adapted with our
+        personalized methods: ~fl4health.mixins.personalized.
+
+        As with `BasicClient`, users are responsible for implementing methods:
+
+            - ``get_model``
+            - ``get_optimizer``
+            - ``get_data_loaders``,
+            - ``get_criterion``
+
+        However, unlike `BasicClient`, users looking to specialize logic for making predictions,
+        and performing train and validation steps, should instead override:
+
+            - ``predict_with_model``
+            - ``_train_step_with_model_and_optimizer`` (and its delegated helpers)
+            - ``_val_step_with_model``
+
+        Other methods can be overridden to achieve custom functionality.
+
+        Args:
+            data_path (Path): path to the data to be used to load the data for client-side training
+            metrics (Sequence[Metric]): Metrics to be computed based on the labels and predictions of the client model
+            device (torch.device): Device indicator for where to send the model, batches, labels etc. Often "cpu" or
+                "cuda"
+            loss_meter_type (LossMeterType, optional): Type of meter used to track and compute the losses over
+                each batch. Defaults to ``LossMeterType.AVERAGE``.
+            checkpoint_and_state_module (ClientCheckpointAndStateModule | None, optional): A module meant to handle
+                both checkpointing and state saving. The module, and its underlying model and state checkpointing
+                components will determine when and how to do checkpointing during client-side training.
+                No checkpointing (state or model) is done if not provided. Defaults to None.
+            reporters (Sequence[BaseReporter] | None, optional): A sequence of FL4Health reporters which the client
+                should send data to. Defaults to None.
+            progress_bar (bool, optional): Whether or not to display a progress bar during client training and
+                validation. Uses ``tqdm``. Defaults to False
+            client_name (str | None, optional): An optional client name that uniquely identifies a client.
+                If not passed, a hash is randomly generated. Client state will use this as part of its state file
+                name. Defaults to None.
+        """
+        super().__init__(
+            data_path,
+            metrics,
+            device,
+            loss_meter_type,
+            checkpoint_and_state_module,
+            reporters,
+            progress_bar,
+            client_name,
+        )
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Perform some validations on subclasses of FlexibleClient."""
+        super().__init_subclass__(**kwargs)
+
+        # check that specific methods are not overridden, otherwise throw warning
+        methods_should_not_be_overridden = [
+            (
+                "predict",
+                (
+                    f"`{cls.__name__}` overrides `predict()`, but this method should no longer be overridden. "
+                    "Please use `predict_with_model()` instead."
+                ),
+            ),
+            (
+                "val_step",
+                (
+                    f"`{cls.__name__}` overrides `val_step()`, but this method should no longer be overridden. "
+                    "Please use `_val_step_with_model()` instead."
+                ),
+            ),
+            (
+                "train_step",
+                (
+                    f"`{cls.__name__}` overrides `train_step()`, but this method should no longer be overridden. "
+                    "Please use `_train_step_with_model_and_optimizer()` and its helper methods instead "
+                    "for proper customization."
+                ),
+            ),
+        ]
+
+        for method_name, msg in methods_should_not_be_overridden:
+            if method_name in cls.__dict__:  # method was overridden by subclass
+                log(WARN, msg)
+                warnings.warn(msg, RuntimeWarning, stacklevel=2)
+
+    def _compute_preds_and_losses(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        """
+        Helper method within the train step for computing preds and losses.
+
+        NOTE: Subclasses should implement this helper method if there is a need
+        to specialize this part of the overall train step.
+
+        Args:
+            model (nn.Module): the model used to make predictions
+            optimizer (Optimizer): the associated optimizer
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
+            a dictionary of any predictions produced by the model prior to the
+            application of the backwards phase
+        """
+        # Clear gradients from optimizer if they exist
+        optimizer.zero_grad()
+
+        # Call user defined methods to get predictions and compute loss
+        preds, features = self.predict_with_model(model, input)
+        target = self.transform_target(target)
+        losses = self.compute_training_loss(preds, features, target)
+
+        return losses, preds
+
+    def _apply_backwards_on_losses_and_take_step(
+        self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
+    ) -> TrainingLosses:
+        """
+        Helper method within the train step for applying backwards on losses and taking step with optimizer.
+
+        NOTE: Subclasses should implement this helper method if there is a need
+        to specialize this part of the overall train step.
+
+        Args:
+            model (nn.Module): the model used for making predictions. Passed here in case subclasses need it.
+            optimizer (Optimizer): the optimizer with which we take the step
+            losses (TrainingLosses): the losses to apply backwards on
+
+        Returns:
+            TrainingLosses: The losses object post backwards application
+        """
+        # Compute backward pass and update parameters with optimizer
+        losses.backward["backward"].backward()
+        self.transform_gradients(losses)
+        optimizer.step()
+
+        return losses
+
+    def _train_step_with_model_and_optimizer(
+        self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        """
+        Helper train step method that allows for injection of model and optimizer.
+
+        NOTE: Subclasses should implement this method if there is a need to specialize
+        the train_step logic.
+
+        Args:
+            model (nn.Module): the model used for making predictions. Passed here in case subclasses need it.
+            optimizer (Optimizer): the optimizer with which we take the step
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
+            a dictionary of any predictions produced by the model.
+        """
+        losses, preds = self._compute_preds_and_losses(model, optimizer, input, target)
+        losses = self._apply_backwards_on_losses_and_take_step(model, optimizer, losses)
+
+        return losses, preds
+
+    def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]:
+        """
+        Given a single batch of input and target data, generate predictions, compute loss, update parameters and
+        optionally update metrics if they exist. (i.e. backprop on a single batch of data).
+        Assumes ``self.model`` is in train mode already.
+
+        Args:
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
+            a dictionary of any predictions produced by the model.
+        """
+        return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target)
+
+    def _val_step_with_model(
+        self, model: nn.Module, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        """
+        Helper method for val_step that allows for injection of model.
+
+        NOTE: Subclasses should implement this method if there is a need to
+        specialize the val_step logic.
+
+        Args:
+            model (nn.Module): the model used for making predictions. Passed here in case subclasses need it.
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the
+            predictions produced by the model.
+        """
+        # Get preds and compute loss
+        with torch.no_grad():
+            preds, features = self.predict_with_model(model, input)
+            target = self.transform_target(target)
+            losses = self.compute_evaluation_loss(preds, features, target)
+
+        return losses, preds
+
+    def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]:
+        """
+        Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already.
+
+        Args:
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the
+            predictions produced by the model.
+        """
+        return self._val_step_with_model(self.model, input, target)
+
+    def predict_with_model(
+        self, model: torch.nn.Module, input: TorchInputType
+    ) -> tuple[TorchPredType, TorchFeatureType]:
+        """
+        Helper predict method that allows for injection of model.
+
+        NOTE: Subclasses should implement this method if there is need to specialize
+        the predict logic of the client.
+
+        Args:
+            model (torch.nn.Module): the model with which to make predictions
+            input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``,
+                it is assumed that the keys of input match the names of the keyword arguments of
+                ``self.model.forward().`
+
+        Returns:
+            tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of
+            predictions indexed by name and the second element contains intermediate activations indexed by name. By
+            passing features, we can compute losses such as the contrastive loss in MOON. All predictions included in
+            dictionary will by default be used to compute metrics separately.
+
+        Raises:
+            TypeError: Occurs when something other than a tensor or dict of tensors is passed in to the model's
+                forward method.
+            ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
+                forward.
+        """
+        if isinstance(input, torch.Tensor):
+            output = model(input)
+        elif isinstance(input, dict):
+            # If input is a dictionary, then we unpack it before computing the forward pass.
+            # Note that this assumes the keys of the input match (exactly) the keyword args
+            # of self.model.forward().
+            output = model(**input)
+        else:
+            raise TypeError("'input' must be of type torch.Tensor or dict[str, torch.Tensor].")
+
+        if isinstance(output, dict):
+            return output, {}
+        if isinstance(output, torch.Tensor):
+            return {"prediction": output}, {}
+        if isinstance(output, tuple):
+            if len(output) != 2:
+                raise ValueError(f"Output tuple should have length 2 but has length {len(output)}")
+            preds, features = output
+            return preds, features
+        raise ValueError("Model forward did not return a tensor, dictionary of tensors, or tuple of tensors")
+
+    def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
+        """
+        Helper transform gradients method that allows for injection of model.
+
+        NOTE: Subclasses should implement this helper should there be a need to specialize the logic
+        for transforming gradients.
+
+        Args:
+            model (torch.nn.Module): the model used to generate predictions to compute losses
+            losses (TrainingLosses): The losses object from the train step
+        """
+        pass
+
+    def transform_gradients(self, losses: TrainingLosses) -> None:
+        """
+        Hook function for model training only called after backwards pass but before optimizer step. Useful for
+        transforming the gradients (such as with gradient clipping) before they are applied to the model weights.
+
+        Args:
+            losses (TrainingLosses): The losses object from the train step
+        """
+        return self._transform_gradients_with_model(self.model, losses)
diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py
index f649b120a..4bef6fe0b 100644
--- a/fl4health/mixins/adaptive_drift_constrained.py
+++ b/fl4health/mixins/adaptive_drift_constrained.py
@@ -8,19 +8,19 @@
 from flwr.common.logger import log
 from flwr.common.typing import Config, NDArrays
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.losses.weight_drift_loss import WeightDriftLoss
-from fl4health.mixins.core_protocols import BasicClientProtocol, BasicClientProtocolPreSetup
+from fl4health.mixins.core_protocols import FlexibleClientProtocol, FlexibleClientProtocolPreSetup
 from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
 from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger
 from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
 from fl4health.utils.losses import TrainingLosses
-from fl4health.utils.typing import TorchFeatureType, TorchPredType, TorchTargetType
+from fl4health.utils.typing import TorchInputType, TorchPredType, TorchTargetType
 
 
 @runtime_checkable
-class AdaptiveDriftConstrainedProtocol(BasicClientProtocol, Protocol):
+class AdaptiveDriftConstrainedProtocol(FlexibleClientProtocol, Protocol):
     loss_for_adaptation: float
     drift_penalty_tensors: list[torch.Tensor] | None
     drift_penalty_weight: float | None
@@ -38,8 +38,12 @@ def __init__(self, *args: Any, **kwargs: Any):
         To be used with `~fl4health.BaseClient` in order to add the ability to compute
         losses via a constrained adaptive drift.
 
+        NOTE: Rather than using `AdaptiveDriftConstraintClient`, if a client subclasses
+        `FlexibleClient`, than this mixin could be used on that subclass to implement the
+        adaptive drift constraint.
+
         Raises:
-            RuntimeError: when the inheriting class does not satisfy `BasicClientProtocolPreSetup`.
+            RuntimeError: when the inheriting class does not satisfy `FlexibleClientProtocolPreSetup`.
         """
         # Initialize mixin-specific attributes with default values
         self.loss_for_adaptation = 0.1
@@ -54,8 +58,8 @@ def __init__(self, *args: Any, **kwargs: Any):
             super().__init__()
 
         # set penalty_loss_function
-        if not isinstance(self, BasicClientProtocolPreSetup):
-            raise RuntimeError("This object needs to satisfy `BasicClientProtocolPreSetup`.")
+        if not isinstance(self, FlexibleClientProtocolPreSetup):
+            raise RuntimeError("This object needs to satisfy `FlexibleClientProtocolPreSetup`.")
         self.penalty_loss_function = WeightDriftLoss(self.device)
 
     def __init_subclass__(cls, **kwargs: Any):
@@ -70,15 +74,15 @@ def __init_subclass__(cls, **kwargs: Any):
         if hasattr(cls, "_dynamically_created"):
             return
 
-        # Check at class definition time if the parent class satisfies BasicClientProtocol
+        # Check at class definition time if the parent class satisfies FlexibleClientProtocol
         for base in cls.__bases__:
-            if base is not AdaptiveDriftConstrainedMixin and issubclass(base, BasicClient):
+            if base is not AdaptiveDriftConstrainedMixin and issubclass(base, FlexibleClient):
                 return
 
         # If we get here, no compatible base was found
         msg = (
             f"Class {cls.__name__} inherits from AdaptiveDriftConstrainedMixin but none of its other "
-            f"base classes is a BasicClient. This may cause runtime errors."
+            f"base classes is a FlexibleClient. This may cause runtime errors."
         )
         log(WARN, msg)
         warnings.warn(msg, RuntimeWarning, stacklevel=2)
@@ -140,39 +144,26 @@ def set_parameters(
 
         super().set_parameters(server_model_state, config, fitting_round)  # type: ignore[safe-super]
 
-    def compute_training_loss(
-        self: AdaptiveDriftConstrainedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> TrainingLosses:
-        """
-        Computes training loss given predictions of the model and ground truth data. Adds to objective by including
-        penalty loss.
+    def train_step(
+        self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        losses, preds = self._compute_preds_and_losses(self.model, self.optimizers["global"], input, target)
+        loss_clone = losses.backward["backward"].clone()
 
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in
-                dictionary will be used to compute metrics.
-            features: (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target: (TorchTargetType): Ground truth data to evaluate predictions against.
-
-        Returns:
-            TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed
-            by name. Additional losses includes penalty loss.
-        """
-        loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target)
-        if additional_losses is None:
-            additional_losses = {}
-
-        additional_losses["loss"] = loss.clone()
-        # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation
-        additional_losses["loss_for_adaptation"] = loss.clone()
-
-        # Compute the drift penalty loss and store it in the additional losses dictionary.
+        # apply penalty
         penalty_loss = self.compute_penalty_loss()
-        additional_losses["penalty_loss"] = penalty_loss.clone()
+        losses.backward["backward"] = losses.backward["backward"] + penalty_loss
+        losses = self._apply_backwards_on_losses_and_take_step(self.model, self.optimizers["global"], losses)
+
+        # prepare return values
+        additional_losses = {
+            "penalty_loss": penalty_loss.clone(),
+            "local_loss": loss_clone,
+            "loss_for_adaptation": loss_clone.clone(),
+        }
+        losses.additional_losses = additional_losses
 
-        return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses)
+        return losses, preds
 
     def get_parameter_exchanger(self: AdaptiveDriftConstrainedProtocol, config: Config) -> ParameterExchanger:
         """
@@ -218,15 +209,15 @@ def compute_penalty_loss(self: AdaptiveDriftConstrainedProtocol) -> torch.Tensor
         return self.penalty_loss_function(self.model, self.drift_penalty_tensors, self.drift_penalty_weight)
 
 
-def apply_adaptive_drift_to_client(client_base_type: type[BasicClient]) -> type[BasicClient]:
+def apply_adaptive_drift_to_client(client_base_type: type[FlexibleClient]) -> type[FlexibleClient]:
     """
     Dynamically create an adapted client class.
 
     Args:
-        client_base_type (type[BasicClient]): The class to be mixed.
+        client_base_type (type[FlexibleClient]): The class to be mixed.
 
     Returns:
-        type[BasicClient]: A basic client that has been mixed with `AdaptiveDriftConstrainedMixin`.
+        type[FlexibleClient]: A basic client that has been mixed with `AdaptiveDriftConstrainedMixin`.
     """
     return type(
         f"AdaptiveDrift{client_base_type.__name__}",
diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py
index 701026690..9d91acfcb 100644
--- a/fl4health/mixins/core_protocols.py
+++ b/fl4health/mixins/core_protocols.py
@@ -32,7 +32,7 @@ def update_after_train(self, local_steps: int, loss_dict: dict[str, float], conf
 
 
 @runtime_checkable
-class BasicClientProtocolPreSetup(NumPyClientMinimalProtocol, Protocol):
+class FlexibleClientProtocolPreSetup(NumPyClientMinimalProtocol, Protocol):
     """A minimal protocol for BasicClient focused on methods."""
 
     device: torch.device
@@ -61,7 +61,7 @@ def compute_loss_and_additional_losses(
 
 
 @runtime_checkable
-class BasicClientProtocol(BasicClientProtocolPreSetup, Protocol):
+class FlexibleClientProtocol(FlexibleClientProtocolPreSetup, Protocol):
     """A minimal protocol for BasicClient focused on methods."""
 
     model: nn.Module
@@ -77,12 +77,38 @@ def initialize_all_model_weights(self, parameters: NDArrays, config: Config) ->
     def update_before_train(self, current_server_round: int) -> None:
         pass  # pragma: no cover
 
-    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+    def _compute_preds_and_losses(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def _apply_backwards_on_losses_and_take_step(
+        self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
+    ) -> TrainingLosses:
+        pass  # pragma: no cover
+
+    def _train_step_with_model_and_optimizer(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def _val_step_with_model(
+        self, model: nn.Module, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def predict_with_model(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
         pass  # pragma: no cover
 
     def transform_target(self, target: TorchTargetType) -> TorchTargetType:
         pass  # pragma: no cover
 
+    def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
+        pass  # pragma: no cover
+
+    def transform_gradients(self, losses: TrainingLosses) -> None:
+        pass  # pragma: no cover
+
     def compute_training_loss(
         self,
         preds: TorchPredType,
diff --git a/fl4health/mixins/personalized/__init__.py b/fl4health/mixins/personalized/__init__.py
index 08e25d302..daa3f08fe 100644
--- a/fl4health/mixins/personalized/__init__.py
+++ b/fl4health/mixins/personalized/__init__.py
@@ -1,6 +1,6 @@
 from enum import Enum
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.mixins.personalized.ditto import DittoPersonalizedMixin, DittoPersonalizedProtocol
 
 
@@ -11,7 +11,7 @@ class PersonalizedMode(Enum):
 PersonalizedMixinRegistry = {PersonalizedMode.DITTO: DittoPersonalizedMixin}
 
 
-def make_it_personal(client_base_type: type[BasicClient], mode: PersonalizedMode) -> type[BasicClient]:
+def make_it_personal(client_base_type: type[FlexibleClient], mode: PersonalizedMode) -> type[FlexibleClient]:
     """A mixed class factory for converting basic clients to personalized versions."""
     if mode == PersonalizedMode.DITTO:
         return type(
diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py
index 26e618534..8be311774 100644
--- a/fl4health/mixins/personalized/ditto.py
+++ b/fl4health/mixins/personalized/ditto.py
@@ -1,7 +1,8 @@
 """Ditto Personalized Mixin."""
 
+import copy
 import warnings
-from logging import DEBUG, ERROR, INFO, WARN
+from logging import ERROR, INFO, WARN
 from typing import Any, Protocol, cast, runtime_checkable
 
 import torch
@@ -10,9 +11,9 @@
 from torch import nn
 from torch.optim import Optimizer
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.mixins.adaptive_drift_constrained import AdaptiveDriftConstrainedMixin, AdaptiveDriftConstrainedProtocol
-from fl4health.mixins.core_protocols import BasicClientProtocolPreSetup
+from fl4health.mixins.core_protocols import FlexibleClientProtocolPreSetup
 from fl4health.mixins.personalized.utils import ensure_protocol_compliance
 from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 from fl4health.utils.config import narrow_dict_type
@@ -34,9 +35,6 @@ def _copy_optimizer_with_new_params(self, original_optimizer: Optimizer) -> Opti
     def set_initial_global_tensors(self) -> None:
         pass  # pragma: no cover
 
-    def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        pass  # pragma: no cover
-
     def safe_global_model(self) -> nn.Module:
         pass  # pragma: no cover
 
@@ -46,7 +44,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         """
         This mixin implements the Ditto algorithm from Ditto: Fair and Robust Federated Learning Through
         Personalization. This mixin inherits from the `AdaptiveDriftConstrainedMixin`, and like that mixin,
-        this should be mixed with a `BasicClient` type in order to apply the Ditto personalization method
+        this should be mixed with a `FlexibleClient` type in order to apply the Ditto personalization method
         to that client.
 
         Background Context:
@@ -56,7 +54,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
 
         Raises:
-            RuntimeError: If the object does not satisfy the `BasicClientProtocolPreSetup`
+            RuntimeError: If the object does not satisfy the `FlexibleClientProtocolPreSetup`
             then it will raise an error. This is additional validation to ensure that the mixin was
             applied to an appropriate base class.
         """
@@ -70,8 +68,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             # if a parent class doesn't take args/kwargs
             super().__init__()
 
-        if not isinstance(self, BasicClientProtocolPreSetup):
-            raise RuntimeError("This object needs to satisfy `BasicClientProtocolPreSetup`.")
+        if not isinstance(self, FlexibleClientProtocolPreSetup):
+            raise RuntimeError("This object needs to satisfy `FlexibleClientProtocolPreSetup`.")  # pragma: no cover
 
     def __init_subclass__(cls, **kwargs: Any) -> None:
         """This method is called when a class inherits from AdaptiveMixin."""
@@ -85,15 +83,15 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
         if hasattr(cls, "_dynamically_created"):
             return
 
-        # Check at class definition time if the parent class satisfies BasicClientProtocol
+        # Check at class definition time if the parent class satisfies FlexibleClientProtocol
         for base in cls.__bases__:
-            if base is not DittoPersonalizedMixin and issubclass(base, BasicClient):
+            if base is not DittoPersonalizedMixin and issubclass(base, FlexibleClient):
                 return
 
         # If we get here, no compatible base was found
         msg = (
             f"Class {cls.__name__} inherits from DittoPersonalizedMixin but none of its other "
-            f"base classes implement BasicClient. This may cause runtime errors."
+            f"base classes implement FlexibleClient. This may cause runtime errors."
         )
         log(WARN, msg)
         warnings.warn(msg, RuntimeWarning, stacklevel=2)
@@ -128,7 +126,7 @@ def _copy_optimizer_with_new_params(self: DittoPersonalizedProtocol, original_op
         Helper method to make a copy of the original optimizer for the global model.
 
         Args:
-            original_optimizer (Optimizer): original optimizer of the underlying `BasicClient`.
+            original_optimizer (Optimizer): original optimizer of the underyling `FlexibleClient`.
 
         Returns:
             Optimizer: a copy of the original optimizer to be used by the global model.
@@ -173,7 +171,8 @@ def get_global_model(self: DittoPersonalizedProtocol, config: Config) -> nn.Modu
         Returns:
             nn.Module: The PyTorch model serving as the global model for Ditto
         """
-        return self.get_model(config).to(self.device)
+        model_copy = copy.deepcopy(self.get_model(config))
+        return model_copy.to(self.device)
 
     @ensure_protocol_compliance
     def get_optimizer(self: DittoPersonalizedProtocol, config: Config) -> dict[str, Optimizer]:
@@ -366,179 +365,60 @@ def train_step(
             model optimization steps. The prediction dictionary contains predictions indexed a "global" and "local"
             corresponding to predictions from the global and local Ditto models for metric evaluations.
         """
-        # Clear gradients from optimizers if they exist
-        self.optimizers["global"].zero_grad()
-        self.optimizers["local"].zero_grad()
-
-        # Forward pass on both the global and local models
-        preds, features = self.predict(input)
-        target = self.transform_target(target)  # Apply transformation (Defaults to identity)
-
-        # Compute all relevant losses
-        losses = self.compute_training_loss(preds, features, target)
-
-        # Take a step with the global model vanilla loss
-        losses.additional_losses["global_loss"].backward()
-        self.optimizers["global"].step()
-
-        # Take a step with the local model using the local loss and Ditto constraint
-        losses.backward["backward"].backward()
-        self.optimizers["local"].step()
-
-        # Return dictionary of predictions where key is used to name respective MetricMeters
-        return losses, preds
-
-    def predict(
-        self: DittoPersonalizedProtocol,
-        input: TorchInputType,
-    ) -> tuple[TorchPredType, TorchFeatureType]:
-        """
-        Computes the predictions for both the **GLOBAL** and **LOCAL** models and pack them into the prediction
-        dictionary.
+        # global
+        global_losses, global_preds = self._compute_preds_and_losses(
+            self.safe_global_model(), self.optimizers["global"], input, target
+        )
+        # local
+        local_losses, local_preds = self._compute_preds_and_losses(self.model, self.optimizers["local"], input, target)
+        local_loss_clone = local_losses.backward["backward"].clone()  # need a clone for later
 
-        Args:
-            input (TorchInputType): Inputs to be fed into both models.
+        # take step global
+        global_losses = self._apply_backwards_on_losses_and_take_step(
+            self.safe_global_model(), self.optimizers["global"], global_losses
+        )
+        # take step local
+        penalty_loss = self.compute_penalty_loss()
+        local_losses.backward["backward"] = local_losses.backward["backward"] + penalty_loss
+        local_losses = self._apply_backwards_on_losses_and_take_step(
+            self.model, self.optimizers["local"], local_losses
+        )
 
-        Returns:
-            tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains predictions indexed by
-            name and the second element contains intermediate activations index by name. For Ditto, we only need the
-            predictions, so the second dictionary is simply empty.
-
-        Raises:
-            ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
-                forward.
-        """
-        if hasattr(self, "_predict"):
-            log(INFO, "Using '_predict' to make predictions")
-            global_preds, _ = self._predict(self.safe_global_model(), input)
-            local_preds, _ = self._predict(self.model, input)
-            log(INFO, "Successfully predicted for global and local models")
-        elif isinstance(input, torch.Tensor):
-            global_preds = self.safe_global_model()(input)
-            local_preds = self.model(input)
-        elif isinstance(input, dict):
-            # If input is a dictionary, then we unpack it before computing the forward pass.
-            # Note that this assumes the keys of the input match (exactly) the keyword args
-            # of the forward method.
-            global_preds = self.safe_global_model()(**input)
-            local_preds = self.model(**input)
+        # prepare return values
+        additional_losses = {
+            "penalty_loss": penalty_loss.clone(),
+            "local_loss": local_loss_clone,
+            "global_loss": global_losses.backward["backward"],
+            "loss_for_adaptation": local_loss_clone.clone(),
+        }
+        local_losses.additional_losses = additional_losses
 
+        # combined preds
         if isinstance(global_preds, torch.Tensor) and isinstance(local_preds, torch.Tensor):
-            return {"global": global_preds, "local": local_preds}, {}
-        if isinstance(global_preds, dict) and isinstance(local_preds, dict):
-            retval = {f"global-{k}": v for k, v in global_preds.items()}
-            retval.update(**{f"local-{k}": v for k, v in local_preds.items()})
-            return retval, {}
-        raise ValueError(f"Unsupported pred types. Global: {type(global_preds)}, local: {type(local_preds)}")
-
-    def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        """
-        Helper method to extract predictions from global and local models.
-
-        Args:
-            kind (str): the kind of predictions to be extracted i.e., for 'local' or for 'global'
-            preds (dict[str, torch.Tensor]): the full preds result from `self.predict()`
+            combined_preds = {"global": global_preds, "local": local_preds}
+        elif isinstance(global_preds, dict) and isinstance(local_preds, dict):
+            combined_preds = {f"global-{k}": v for k, v in global_preds.items()}
+            combined_preds.update(**{f"local-{k}": v for k, v in local_preds.items()})
 
-        Raises:
-            ValueError: If supplied `kind` is not of 'global' or 'local'
-
-        Returns:
-            dict[str, torch.Tensor]: the predictions associated with the specified kind
-        """
-        if kind not in ["global", "local"]:
-            raise ValueError("Unsupported kind of prediction. Must be 'global' or 'local'.")
+        return local_losses, combined_preds
 
-        # filter
-        retval = {k: v for k, v in preds.items() if kind in k}
-        # remove prefix
-        return {k.replace(f"{kind}-", ""): v for k, v in retval.items()}
-
-    def compute_loss_and_additional_losses(
-        self: DittoPersonalizedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
-        """
-        Computes the local model loss and the global Ditto model loss (stored in additional losses) for reporting and
-        training of the global model.
-
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name.
-            features (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target (TorchTargetType): Ground truth data to evaluate predictions against.
-
-        Returns:
-            tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple with:
-
-            - The tensor for the model loss
-            - A dictionary with ``local_loss``, ``global_loss`` as additionally reported loss values.
-        """
-        global_preds = self._extract_pred(kind="global", preds=preds)
-        local_preds = self._extract_pred(kind="local", preds=preds)
-
-        log(DEBUG, f"global_preds has keys {global_preds.keys()}")
-        log(DEBUG, f"local_preds has keys {local_preds.keys()}")
-
-        # Compute global model vanilla loss
-
-        if hasattr(self, "_special_compute_loss_and_additional_losses"):
-            log(INFO, "Using '_special_compute_loss_and_additional_losses' to compute loss")
-            global_loss, _ = self._special_compute_loss_and_additional_losses(global_preds, features, target)
-
-            # Compute local model loss
-            local_loss, _ = self._special_compute_loss_and_additional_losses(local_preds, features, target)
-
-        else:
-            global_loss = self.criterion(global_preds["global"], target)
-
-            # Compute local model loss
-            local_loss = self.criterion(local_preds["local"], target)
-
-        additional_losses = {"local_loss": local_loss.clone(), "global_loss": global_loss}
-
-        return local_loss, additional_losses
-
-    def compute_training_loss(
-        self: DittoPersonalizedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> TrainingLosses:
-        """
-        Computes training losses given predictions of the global and local models and ground truth data.
-        For the local model we add to the vanilla loss function by including Ditto penalty loss which is the l2 inner
-        product between the initial global model weights and weights of the local model. This is stored in backward
-        The loss to optimize the global model is stored in the additional losses dictionary under "global_loss".
-
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in
-                dictionary will be used to compute metrics.
-            features: (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target: (TorchTargetType): Ground truth data to evaluate predictions against.
-
-        Returns:
-            TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed by
-            name. Additional losses includes each loss component and the global model
-            loss tensor.
-        """
-        # Check that both models are in training mode
-        assert self.safe_global_model().training and self.model.training
-
-        # local loss is stored in loss, global model loss is stored in additional losses.
-        loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target)
-        additional_losses = additional_losses or {}  # make mypy happy
-
-        # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or
-        # less weight is used to constrain it to the global model (as in FedProx)
-        additional_losses["loss_for_adaptation"] = additional_losses["local_loss"].clone()
-
-        # This is the Ditto penalty loss of the local model compared with the original Global model weights, scaled
-        # by drift_penalty_weight (or lambda in the original paper)
-        penalty_loss = self.compute_penalty_loss()
-        additional_losses["penalty_loss"] = penalty_loss.clone()
-
-        return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses)
+    def val_step(
+        self: DittoPersonalizedProtocol, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        # global
+        global_losses, global_preds = self._val_step_with_model(self.safe_global_model(), input, target)
+        # local
+        local_losses, local_preds = self._val_step_with_model(self.model, input, target)
+
+        # combine
+        losses = EvaluationLosses(
+            local_losses.checkpoint,
+            additional_losses={"global_loss": global_losses.checkpoint, "local_loss": local_losses.checkpoint},
+        )
+        preds: TorchPredType = {}
+        preds.update(**{f"global-{k}": v for k, v in global_preds.items()})
+        preds.update(**{f"local-{k}": v for k, v in local_preds.items()})
+        return losses, preds
 
     @ensure_protocol_compliance
     def validate(
diff --git a/fl4health/mixins/personalized/utils.py b/fl4health/mixins/personalized/utils.py
index dcd052daa..5449a0241 100644
--- a/fl4health/mixins/personalized/utils.py
+++ b/fl4health/mixins/personalized/utils.py
@@ -3,15 +3,15 @@
 
 import wrapt
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 
 
 @wrapt.decorator
 def ensure_protocol_compliance(func: Callable, instance: Any | None, args: Any, kwargs: Any) -> Any:
     """
-    Wrapper to ensure that a the instance is of `BasicClient` type.
+    Wrapper to ensure that a the instance is of `FlexibleClient` type.
 
-    NOTE: this should only be used within a `BasicClient`.
+    NOTE: this should only be used within a `FlexibleClient`.
 
     Args:
         # are params specified and supplied by the `wrapt` decorator
@@ -21,13 +21,13 @@ def ensure_protocol_compliance(func: Callable, instance: Any | None, args: Any,
         kwargs (Any): kwargs passed to func
 
     Raises:
-        TypeError: we raise error if the instance is not a `BasicClient`.
+        TypeError: we raise error if the instance is not a `FlexibleClient`.
 
     Returns:
         _type_: _description_
     """
-    # validate self is a BasicClient
-    if not isinstance(instance, BasicClient):
+    # validate self is a FlexibleClient
+    if not isinstance(instance, FlexibleClient):
         raise TypeError("Protocol requirements not met.")
 
     return func(*args, **kwargs)
diff --git a/tests/clients/test_flexible_client.py b/tests/clients/test_flexible_client.py
new file mode 100644
index 000000000..d65c51c49
--- /dev/null
+++ b/tests/clients/test_flexible_client.py
@@ -0,0 +1,286 @@
+import datetime
+import re
+from collections.abc import Sequence
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import freezegun
+import pytest
+import torch
+from flwr.common import Scalar
+from flwr.common.typing import Config
+from freezegun import freeze_time
+from torch.utils.data import DataLoader
+
+from fl4health.clients.flexible_client import FlexibleClient
+from fl4health.reporting import JsonReporter
+from fl4health.reporting.base_reporter import BaseReporter
+from fl4health.utils.client import fold_loss_dict_into_metrics
+from fl4health.utils.dataset import TensorDataset
+from fl4health.utils.logging import LoggingMode
+from fl4health.utils.losses import EvaluationLosses, TrainingLosses
+from fl4health.utils.random import set_all_random_seeds, unset_all_random_seeds
+from fl4health.utils.typing import TorchInputType, TorchPredType, TorchTargetType
+from tests.test_utils.assert_metrics_dict import assert_metrics_dict
+from tests.test_utils.models_for_test import LinearModel
+
+
+freezegun.configure(extend_ignore_list=["transformers"])  # type: ignore
+
+
+def get_dummy_dataset() -> TensorDataset:
+    data = torch.randn(100, 10, 8)
+    targets = torch.randint(5, (100,))
+    return TensorDataset(data=data, targets=targets)
+
+
+DUMMY_DATASET = get_dummy_dataset()
+
+
+@freeze_time("2012-12-12 12:12:12")
+def test_json_reporter_setup_client() -> None:
+    reporter = JsonReporter()
+    fl_client = MockFlexibleClient(reporters=[reporter])
+    fl_client.setup_client({})
+
+    metric_dict = {
+        "host_type": "client",
+        "initialized": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+    }
+    errors = assert_metrics_dict(metric_dict, reporter.metrics)
+    assert len(errors) == 0, f"Metrics check failed. Errors: {errors}"
+
+
+@freeze_time("2012-12-12 12:12:12")
+def test_json_reporter_shutdown() -> None:
+    reporter = JsonReporter()
+    fl_client = MockFlexibleClient(reporters=[reporter])
+    fl_client.shutdown()
+
+    metric_dict = {
+        "shutdown": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+    }
+    errors = assert_metrics_dict(metric_dict, reporter.metrics)
+    assert len(errors) == 0, f"Metrics check failed. Errors: {errors}"
+
+
+@freeze_time("2012-12-12 12:12:12")
+def test_metrics_reporter_fit() -> None:
+    test_current_server_round = 2
+    test_loss_dict = {"test_loss": 123.123}
+    test_metrics: dict[str, Scalar] = {"test_metric": 1234}
+    reporter = JsonReporter()
+
+    fl_client = MockFlexibleClient(loss_dict=test_loss_dict, metrics=test_metrics, reporters=[reporter])
+    fl_client.fit([], {"current_server_round": test_current_server_round, "local_epochs": 0})
+    metric_dict = {
+        "host_type": "client",
+        "initialized": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+        "rounds": {
+            test_current_server_round: {
+                "round_start": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+                "fit_round_losses": test_loss_dict,
+                "fit_round_metrics": test_metrics,
+                "round": test_current_server_round,
+            },
+        },
+    }
+
+    errors = assert_metrics_dict(metric_dict, reporter.metrics)
+    assert len(errors) == 0, f"Metrics check failed. Errors: {errors}"
+
+
+@freeze_time("2012-12-12 12:12:12")
+def test_metrics_reporter_evaluate() -> None:
+    test_current_server_round = 2
+    test_loss = 123.123
+    test_metrics: dict[str, Scalar] = {"test_metric": 1234}
+    test_metrics_testing: dict[str, Scalar] = {"testing_metric": 1234}
+    test_metrics_final = {
+        "test_metric": 1234,
+        "testing_metric": 1234,
+        "val - checkpoint": 123.123,
+        "test - checkpoint": 123.123,
+        "test - num_examples": 32,
+    }
+    reporter = JsonReporter()
+    fl_client = MockFlexibleClient(
+        loss=test_loss,
+        loss_dict={"checkpoint": test_loss},
+        metrics=test_metrics,
+        test_set_metrics=test_metrics_testing,
+        reporters=[reporter],
+    )
+    fl_client.evaluate(
+        [],
+        {"current_server_round": test_current_server_round, "local_epochs": 0, "pack_losses_with_val_metrics": True},
+    )
+
+    metric_dict = {
+        "host_type": "client",
+        "initialized": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+        "rounds": {
+            test_current_server_round: {
+                "eval_round_start": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+                "eval_round_loss": test_loss,
+                "eval_round_metrics": test_metrics_final,
+                "eval_round_end": str(datetime.datetime(2012, 12, 12, 12, 12, 12)),
+            },
+        },
+    }
+
+    errors = assert_metrics_dict(metric_dict, reporter.metrics)
+    assert len(errors) == 0, f"Metrics check failed. Errors: {errors}"
+
+
+def test_evaluate_after_fit_enabled() -> None:
+    fl_client = MockFlexibleClient()
+    fl_client.validate = MagicMock()  # type: ignore
+    fl_client.validate.return_value = fl_client.mock_loss, fl_client.mock_metrics
+
+    fl_client.fit([], {"current_server_round": 2, "local_epochs": 0, "evaluate_after_fit": True})
+
+    fl_client.validate.assert_called_once()  # type: ignore
+
+
+def test_evaluate_after_fit_disabled() -> None:
+    fl_client = MockFlexibleClient()
+    fl_client.validate = MagicMock()  # type: ignore
+    fl_client.validate.return_value = fl_client.mock_loss, fl_client.mock_metrics
+
+    fl_client.fit([], {"current_server_round": 2, "local_epochs": 0, "evaluate_after_fit": False})
+    fl_client.validate.assert_not_called()  # type: ignore
+
+    fl_client.fit([], {"current_server_round": 2, "local_epochs": 0})
+    fl_client.validate.assert_not_called()  # type: ignore
+
+
+def test_validate_by_steps() -> None:
+    # Set the random seeds
+    set_all_random_seeds(2023)
+
+    fl_client = MockFlexibleClient()
+    fl_client.num_validation_steps = 2
+    fl_client.model = LinearModel()
+    fl_client.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    fl_client.val_loader = DataLoader(DUMMY_DATASET, batch_size=15, shuffle=False)
+    loss, metrics = fl_client._validate_by_steps(fl_client.val_loss_meter, fl_client.val_metric_manager)
+    assert loss == (1.0 + 2.0) / 2.0
+    assert fl_client.val_iterator is not None
+
+    unset_all_random_seeds()
+
+
+def test_num_val_samples_correct() -> None:
+    fl_client_no_max = MockFlexibleClient()
+    fl_client_no_max.setup_client({})
+    assert fl_client_no_max.num_validation_steps is None
+    assert fl_client_no_max.num_val_samples == 32
+
+    fl_client_max = MockFlexibleClient()
+    config: Config = {"num_validation_steps": 2}
+    fl_client_max.setup_client(config)
+    assert fl_client_max.num_validation_steps == 2
+    assert fl_client_max.num_val_samples == 8
+
+
+class MockFlexibleClient(FlexibleClient):
+    def __init__(
+        self,
+        loss_dict: dict[str, float] | None = None,
+        metrics: dict[str, Scalar] | None = None,
+        test_set_metrics: dict[str, Scalar] | None = None,
+        loss: float | None = 0,
+        reporters: Sequence[BaseReporter] | None = None,
+    ):
+        super().__init__(Path(""), [], torch.device(0), reporters=reporters)
+
+        self.mock_loss_dict = loss_dict
+        if self.mock_loss_dict is None:
+            self.mock_loss_dict = {}
+
+        self.mock_metrics = metrics
+        if self.mock_metrics is None:
+            self.mock_metrics = {}
+
+        self.mock_metrics_test = test_set_metrics
+
+        self.mock_loss = loss
+
+        # Mocking attributes
+        self.train_loader = MagicMock()
+        self.test_loader = MagicMock()
+        self.num_train_samples = 0
+        self.num_val_samples = 0
+        self.num_validation_steps = None
+
+        # Mocking methods
+        self.set_parameters = MagicMock()  # type: ignore
+        self.get_parameters = MagicMock()  # type: ignore
+        self.train_by_epochs = MagicMock()  # type: ignore
+        self.train_by_epochs.return_value = self.mock_loss_dict, self.mock_metrics
+        self.train_by_steps = MagicMock()  # type: ignore
+        self.train_by_steps.return_value = self.mock_loss_dict, self.mock_metrics
+        self.get_model = MagicMock()  # type: ignore
+        self.get_data_loaders = MagicMock()  # type: ignore
+        mock_data_loader = MagicMock()  # type: ignore
+        mock_data_loader.batch_size = 4
+        mock_data_loader.dataset = [None] * 32
+        mock_data_loader.__len__ = lambda _: len(mock_data_loader.dataset) // mock_data_loader.batch_size
+        self.get_data_loaders.return_value = mock_data_loader, mock_data_loader
+        self.get_test_data_loader = MagicMock()  # type: ignore
+        self.get_test_data_loader.return_value = mock_data_loader
+        self.get_optimizer = MagicMock()  # type: ignore
+        self.get_criterion = MagicMock()  # type: ignore
+        self.val_step = MagicMock()  # type: ignore
+        self.val_step.side_effect = self.mock_val_step
+
+        self._fully_validate_or_test = MagicMock()  # type: ignore
+        self._fully_validate_or_test.side_effect = self.mock_validate_or_test
+
+    def mock_val_step(self, input: TorchInputType, target: TorchTargetType):  # type: ignore
+        loss_tensor = torch.randint(1, 4, (1,))
+        return EvaluationLosses(loss_tensor), self.mock_metrics
+
+    def mock_validate_or_test(  # type: ignore
+        self, loader, loss_meter, metric_manager, logging_mode=LoggingMode.VALIDATION, include_losses_in_metrics=False
+    ):
+        if include_losses_in_metrics:
+            assert self.mock_loss_dict is not None and self.mock_metrics is not None
+            fold_loss_dict_into_metrics(self.mock_metrics, self.mock_loss_dict, logging_mode)
+        if logging_mode == LoggingMode.VALIDATION:
+            return self.mock_loss, self.mock_metrics
+        return self.mock_loss, self.mock_metrics_test
+
+
+def test_subclass_raises_warning_if_override_val_step() -> None:
+    msg = (
+        "`_TestSubclass` overrides `val_step()`, but this method should no longer be overridden. "
+        "Please use `_val_step_with_model()` instead."
+    )
+    with pytest.warns(RuntimeWarning, match=re.escape(msg)):
+
+        class _TestSubclass(FlexibleClient):
+            """A subclass that overrides val_step."""
+
+            def val_step(
+                self, input: TorchInputType, target: TorchTargetType
+            ) -> tuple[EvaluationLosses, TorchPredType]:
+                return super().val_step(input, target)
+
+
+def test_subclass_raises_warning_if_override_train_step() -> None:
+    msg = (
+        "`_TestSubclass` overrides `train_step()`, but this method should no longer be overridden. "
+        "Please use `_train_step_with_model_and_optimizer()` and its helper methods instead "
+        "for proper customization."
+    )
+    with pytest.warns(RuntimeWarning, match=re.escape(msg)):
+
+        class _TestSubclass(FlexibleClient):
+            """A subclass that overrides train_step."""
+
+            def train_step(
+                self, input: TorchInputType, target: TorchTargetType
+            ) -> tuple[TrainingLosses, TorchPredType]:
+                return super().train_step(input, target)
diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py
index 39a47677a..bb2feef36 100644
--- a/tests/mixins/personalized/test_ditto.py
+++ b/tests/mixins/personalized/test_ditto.py
@@ -1,3 +1,4 @@
+import re
 import warnings
 from logging import INFO
 from pathlib import Path
@@ -11,12 +12,13 @@
 from torch import nn
 from torch.nn.modules.loss import _Loss
 from torch.optim import Optimizer
-from torch.testing import assert_close
+
+# from torch.testing import assert_close
 from torch.utils.data import DataLoader, TensorDataset
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.metrics import Accuracy
-from fl4health.mixins.core_protocols import BasicClientProtocol
+from fl4health.mixins.core_protocols import FlexibleClientProtocol
 from fl4health.mixins.personalized import (
     DittoPersonalizedMixin,
     DittoPersonalizedProtocol,
@@ -25,11 +27,10 @@
 )
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
 from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
-from fl4health.utils.losses import TrainingLosses
-from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType
+from fl4health.utils.losses import EvaluationLosses, TrainingLosses
 
 
-class _TestBasicClient(BasicClient):
+class _TestFlexibleClient(FlexibleClient):
     def get_model(self, config: Config) -> nn.Module:
         return self.model
 
@@ -43,28 +44,16 @@ def get_criterion(self, config: Config) -> _Loss:
         return torch.nn.CrossEntropyLoss()
 
 
-class _TestBasicClientV2(BasicClient):
-    def get_model(self, config: Config) -> nn.Module:
-        return self.model
-
-    def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
-        return self.train_loader, self.val_loader
-
-    def get_optimizer(self, config: Config) -> Optimizer | dict[str, Optimizer]:
-        return self.optimizers["local"]
-
-    def get_criterion(self, config: Config) -> _Loss:
-        return torch.nn.CrossEntropyLoss()
-
-    def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
-        return {}, {}
+class _TestDittoedClient(DittoPersonalizedMixin, _TestFlexibleClient):
+    pass
 
 
-class _TestDittoedClient(DittoPersonalizedMixin, _TestBasicClient):
-    pass
+class _DummyParent:
+    def __init__(self) -> None:
+        pass
 
 
-class _TestDittoedClientV2(DittoPersonalizedMixin, _TestBasicClientV2):
+class _TestInvalidDittoedClient(DittoPersonalizedMixin, _DummyParent):
     pass
 
 
@@ -79,7 +68,7 @@ def test_init() -> None:
     client.initialized = True
     client.setup_client({})
 
-    assert isinstance(client, BasicClientProtocol)
+    assert isinstance(client, FlexibleClientProtocol)
     assert isinstance(client, DittoPersonalizedProtocol)
 
 
@@ -90,20 +79,20 @@ def test_init_raises_value_error_when_basic_client_protocol_not_satisfied() -> N
     class _InvalidTestDittoClient(DittoPersonalizedMixin):
         pass
 
-    with pytest.raises(RuntimeError, match="This object needs to satisfy `BasicClientProtocolPreSetup`."):
+    with pytest.raises(RuntimeError, match="This object needs to satisfy `FlexibleClientProtocolPreSetup`."):
         _InvalidTestDittoClient(data_path=Path(""), metrics=[Accuracy()])
 
 
 def test_subclass_checks_raise_no_warning() -> None:
     with warnings.catch_warnings(record=True) as recorded_warnings:
 
-        class _TestInheritanceMixin(DittoPersonalizedMixin, _TestBasicClient):
-            """subclass should skip validation if is itself a Mixin that inherits DittoPersonalizedMixin."""
+        class _TestInheritanceMixin(DittoPersonalizedMixin, _TestFlexibleClient):
+            """Subclass should skip validation if is itself a Mixin that inherits DittoPersonalizedMixin."""
 
             pass
 
         # attaches _dynamically_created attr
-        _ = make_it_personal(_TestBasicClient, PersonalizedMode.DITTO)
+        _ = make_it_personal(_TestFlexibleClient, PersonalizedMode.DITTO)
 
     assert len(recorded_warnings) == 0
 
@@ -113,7 +102,7 @@ def test_subclass_checks_raise_warning() -> None:
     with pytest.warns((RuntimeWarning, RuntimeWarning)):
 
         class _InvalidSubclass(DittoPersonalizedMixin):
-            """Invalid subclass that warns the user that it expects this class to be mixed with a BasicClient."""
+            """Invalid subclass that warns the user that it expects this class to be mixed with a FlexibleClient."""
 
             pass
 
@@ -147,6 +136,74 @@ def test_get_parameters() -> None:
     assert_array_equal(packed_params, pack_params_return_val)
 
 
+@patch.object(_TestDittoedClient, "compute_penalty_loss")
+@patch.object(_TestDittoedClient, "_apply_backwards_on_losses_and_take_step")
+@patch.object(_TestDittoedClient, "_compute_preds_and_losses")
+def test_train_step(
+    mock_private_compute_preds_and_losses: MagicMock,
+    mock_private_apply_backwards_on_losses_and_take_step: MagicMock,
+    mock_compute_penalty_loss: MagicMock,
+) -> None:
+    # setup client
+    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.model = torch.nn.Linear(5, 5)
+    client.global_model = torch.nn.Linear(5, 5)
+    client.optimizers = {
+        "global": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+        "local": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+    }
+    # setup mocks
+    mock_param_exchanger = MagicMock()
+    push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights
+    mock_param_exchanger.pack_parameters.return_value = pack_params_return_val
+    client.parameter_exchanger = mock_param_exchanger
+    client.initialized = True
+
+    mock_backward_loss = MagicMock()
+    mock_additional_global_loss = MagicMock()
+    dummy_training_losses = TrainingLosses(
+        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    )
+    dummy_training_losses_for_local = TrainingLosses(
+        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    )
+    mock_private_compute_preds_and_losses.side_effect = [
+        (
+            dummy_training_losses,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+        (
+            dummy_training_losses_for_local,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+    ]
+    mock_private_apply_backwards_on_losses_and_take_step.side_effect = [
+        dummy_training_losses,
+        dummy_training_losses_for_local,
+    ]
+
+    # act
+    input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3)
+    # TODO: fix the mixin/protocol typing that leads to mypy complaint
+    _ = client.train_step(input, target)  # type: ignore
+
+    mock_private_compute_preds_and_losses.assert_has_calls(
+        [
+            _Call(((client.global_model, client.optimizers["global"], input, target), {})),
+            _Call(((client.model, client.optimizers["local"], input, target), {})),
+        ]
+    )
+    mock_private_apply_backwards_on_losses_and_take_step.assert_has_calls(
+        [
+            _Call(((client.global_model, client.optimizers["global"], dummy_training_losses), {})),
+            _Call(((client.model, client.optimizers["local"], dummy_training_losses_for_local), {})),
+        ]
+    )
+    mock_compute_penalty_loss.assert_called_once()
+
+
 @patch.object(_TestDittoedClient, "setup_client")
 @patch("fl4health.mixins.personalized.ditto.FullParameterExchanger")
 def test_get_parameters_uninitialized(mock_param_exchanger: MagicMock, mock_setup_client: MagicMock) -> None:
@@ -233,94 +290,8 @@ def test_get_optimizer(mock_copy_optimizer: MagicMock) -> None:
     mock_copy_optimizer.assert_called_once_with(client.optimizers["local"])
 
 
-def test_predict() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    mock_model = MagicMock()
-    mock_global_model = MagicMock()
-
-    mock_model.return_value = torch.ones(5)
-    mock_global_model.return_value = torch.zeros(5)
-
-    client.model = mock_model
-    client.global_model = mock_global_model
-
-    client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
-    }
-
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
-    client.initialized = True
-
-    # act
-    # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    res, _ = client.predict(input=torch.zeros(5))  # type: ignore
-    print(f"res: {res}")
-    print(torch.zeros(5))
-
-    # assert
-    assert_close(res["global"], torch.zeros(5))
-    assert_close(res["local"], torch.ones(5))
-
-
-@patch.object(_TestDittoedClientV2, "_predict")
-def test_predict_delagation(private_predict: MagicMock) -> None:
-    # setup client
-    client = _TestDittoedClientV2(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-    client.model = torch.nn.Linear(5, 5)
-    client.global_model = torch.nn.Linear(5, 5)
-
-    private_predict.side_effect = [(torch.zeros(5), {}), (torch.ones(5), {})]
-
-    client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
-    }
-
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
-    client.initialized = True
-
-    # act
-    # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    res, _ = client.predict(input=torch.zeros(5))  # type: ignore
-    print(f"res: {res}")
-    print(torch.zeros(5))
-
-    # assert
-    assert_close(res["global"], torch.zeros(5))
-    assert_close(res["local"], torch.ones(5))
-
-
-def test_extract_pred() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    res = client._extract_pred(
-        kind="global", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)}
-    )
-
-    assert_close(res["xyz"], torch.ones(5))
-    assert_close(res["abc"], torch.zeros(5))
-
-
-def test_extract_pred_raises_error() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    with pytest.raises(ValueError):
-        client._extract_pred(
-            kind="oops", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)}
-        )
-
-
 @patch.object(_TestDittoedClient, "set_initial_global_tensors")
-@patch.object(_TestBasicClient, "update_before_train")
+@patch.object(_TestFlexibleClient, "update_before_train")
 def test_update_before_train(
     mock_super_update_before_train: MagicMock, mock_set_initial_global_tensors: MagicMock
 ) -> None:
@@ -344,43 +315,60 @@ def test_safe_model_raises_error() -> None:
         client.safe_global_model()  # type: ignore
 
 
-@patch("torch.optim.Optimizer")
-@patch.object(_TestDittoedClient, "predict")
-@patch.object(_TestDittoedClient, "compute_training_loss")
-def test_train_step(
-    mock_compute_training_loss: MagicMock, mock_predict: MagicMock, mock_optimizer_class: MagicMock
+@patch.object(_TestDittoedClient, "_val_step_with_model")
+def test_val_step(
+    mock_private_val_step_with_model: MagicMock,
 ) -> None:
     # setup client
     client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.model = torch.nn.Linear(5, 5)
+    client.global_model = torch.nn.Linear(5, 5)
     client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
+        "global": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+        "local": torch.optim.SGD(client.model.parameters(), lr=0.0001),
     }
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
+    # setup mocks
+    mock_param_exchanger = MagicMock()
+    push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights
+    mock_param_exchanger.pack_parameters.return_value = pack_params_return_val
+    client.parameter_exchanger = mock_param_exchanger
     client.initialized = True
 
-    # arrange mocks
-    pred_return: tuple[TorchPredType, TorchFeatureType] = {"local": torch.ones(5)}, {}
-    preds, features = pred_return
-    mock_predict.return_value = preds, features
-    mock_backward_loss = MagicMock()
-    mock_additional_global_loss = MagicMock()
-    losses = TrainingLosses(
-        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    dummy_training_losses = EvaluationLosses(
+        checkpoint=torch.ones(5),
+    )
+    dummy_training_losses_for_local = EvaluationLosses(
+        checkpoint=torch.ones(5),
     )
-    mock_compute_training_loss.return_value = losses
+    mock_private_val_step_with_model.side_effect = [
+        (
+            dummy_training_losses,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+        (
+            dummy_training_losses_for_local,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+    ]
 
     # act
     input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3)
     # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    retval = client.train_step(input, target)  # type: ignore
-
-    mock_predict.assert_called_once()
-    client.optimizers["global"].zero_grad.assert_called_once()
-    client.optimizers["local"].zero_grad.assert_called_once()
-    mock_compute_training_loss.assert_called_once_with(preds, features, target)
-    mock_backward_loss.backward.assert_called_once()
-    mock_additional_global_loss.backward.assert_called_once()
-    assert retval == (losses, preds)
+    _ = client.val_step(input, target)  # type: ignore
+
+    mock_private_val_step_with_model.assert_has_calls(
+        [
+            _Call(((client.global_model, input, target), {})),
+            _Call(((client.model, input, target), {})),
+        ]
+    )
+
+
+def test_raise_runtime_error_not_flexible_client() -> None:
+    """Test that an invalid parent raises RuntimeError."""
+    with pytest.raises(
+        RuntimeError, match=re.escape("This object needs to satisfy `FlexibleClientProtocolPreSetup`.")
+    ):
+        _TestInvalidDittoedClient()
diff --git a/tests/mixins/personalized/test_factory.py b/tests/mixins/personalized/test_factory.py
index 1e9128bcf..2bc9c8f43 100644
--- a/tests/mixins/personalized/test_factory.py
+++ b/tests/mixins/personalized/test_factory.py
@@ -6,11 +6,11 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.mixins.personalized import DittoPersonalizedMixin, PersonalizedMode, make_it_personal
 
 
-class MyClient(BasicClient):
+class MyClient(FlexibleClient):
     def get_model(self, config: Config) -> nn.Module:
         return self.model
 
@@ -27,7 +27,7 @@ def get_criterion(self, config: Config) -> _Loss:
 def test_make_it_personal_factory_method() -> None:
     ditto_my_client_cls = make_it_personal(MyClient, mode=PersonalizedMode.DITTO)
 
-    assert issubclass(ditto_my_client_cls, (BasicClient, DittoPersonalizedMixin))
+    assert issubclass(ditto_my_client_cls, (FlexibleClient, DittoPersonalizedMixin))
 
 
 def test_make_it_personal_raises_value_error() -> None:
diff --git a/tests/mixins/personalized/test_utils.py b/tests/mixins/personalized/test_utils.py
index 3942a4071..16217711d 100644
--- a/tests/mixins/personalized/test_utils.py
+++ b/tests/mixins/personalized/test_utils.py
@@ -9,7 +9,7 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, TensorDataset
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.metrics import Accuracy
 from fl4health.mixins.personalized.utils import ensure_protocol_compliance
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
@@ -18,7 +18,7 @@
 
 def test_ensure_protocol_compliance_does_not_raise() -> None:
     # arrange
-    class MyClient(BasicClient):
+    class MyClient(FlexibleClient):
         def get_model(self, config: Config) -> nn.Module:
             return self.model
 
@@ -53,7 +53,7 @@ def some_method(self, x: int) -> int:
 def test_ensure_protocol_compliance_does_raise_type_error() -> None:
     # arrange
     class MyClient:
-        """My Client DOES not satisfy the protocol of BasicClient."""
+        """My Client DOES not satisfy the protocol of FlexibleClient."""
 
         @ensure_protocol_compliance
         def some_method(self, x: int) -> int:
diff --git a/tests/mixins/test_adaptive_drift_constrained.py b/tests/mixins/test_adaptive_drift_constrained.py
index 17608c884..f11a919d4 100644
--- a/tests/mixins/test_adaptive_drift_constrained.py
+++ b/tests/mixins/test_adaptive_drift_constrained.py
@@ -13,19 +13,20 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, TensorDataset
 
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.flexible_client import FlexibleClient
 from fl4health.metrics import Accuracy
 from fl4health.mixins.adaptive_drift_constrained import (
     AdaptiveDriftConstrainedMixin,
     AdaptiveDriftConstrainedProtocol,
     apply_adaptive_drift_to_client,
 )
-from fl4health.mixins.core_protocols import BasicClientProtocol
+from fl4health.mixins.core_protocols import FlexibleClientProtocol
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
 from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
+from fl4health.utils.losses import TrainingLosses
 
 
-class _TestBasicClient(BasicClient):
+class _TestFlexibleClient(FlexibleClient):
     def get_model(self, config: Config) -> nn.Module:
         return self.model
 
@@ -39,7 +40,7 @@ def get_criterion(self, config: Config) -> _Loss:
         return torch.nn.CrossEntropyLoss()
 
 
-class _TestAdaptedClient(AdaptiveDriftConstrainedMixin, _TestBasicClient):
+class _TestAdaptedClient(AdaptiveDriftConstrainedMixin, _TestFlexibleClient):
     pass
 
 
@@ -54,7 +55,7 @@ def test_init() -> None:
     client.initialized = True
     client.setup_client({})
 
-    assert isinstance(client, BasicClientProtocol)
+    assert isinstance(client, FlexibleClientProtocol)
     assert isinstance(client, AdaptiveDriftConstrainedProtocol)
 
 
@@ -65,19 +66,19 @@ def test_init_raises_value_error_when_basic_client_protocol_not_satisfied() -> N
     class _InvalidTestAdaptedClient(AdaptiveDriftConstrainedMixin):
         pass
 
-    with pytest.raises(RuntimeError, match="This object needs to satisfy `BasicClientProtocolPreSetup`."):
+    with pytest.raises(RuntimeError, match="This object needs to satisfy `FlexibleClientProtocolPreSetup`."):
         _InvalidTestAdaptedClient(data_path=Path(""), metrics=[Accuracy()])
 
 
 def test_subclass_checks_raise_no_warning() -> None:
     with warnings.catch_warnings(record=True) as recorded_warnings:
 
-        class _TestInheritanceMixin(AdaptiveDriftConstrainedMixin, _TestBasicClient):
-            """subclass should skip validation if is itself a Mixin that inherits AdaptiveDriftConstrainedMixin."""
+        class _TestInheritanceMixin(AdaptiveDriftConstrainedMixin, _TestFlexibleClient):
+            """Subclass should skip validation if is itself a Mixin that inherits AdaptiveDriftConstrainedMixin."""
 
             pass
 
-        class _DynamicallyCreatedClass(AdaptiveDriftConstrainedMixin, _TestBasicClient):
+        class _DynamicallyCreatedClass(AdaptiveDriftConstrainedMixin, _TestFlexibleClient):
             """subclass used for dynamic creation of clients with mixins."""
 
             _dynamically_created = True
@@ -88,12 +89,12 @@ class _DynamicallyCreatedClass(AdaptiveDriftConstrainedMixin, _TestBasicClient):
 def test_subclass_checks_raise_warning() -> None:
     msg = (
         "Class _InvalidSubclass inherits from AdaptiveDriftConstrainedMixin but none of its other "
-        "base classes is a BasicClient. This may cause runtime errors."
+        "base classes is a FlexibleClient. This may cause runtime errors."
     )
     with pytest.warns(RuntimeWarning, match=msg):
 
         class _InvalidSubclass(AdaptiveDriftConstrainedMixin):
-            """Invalid subclass that warns the user that it expects this class to be mixed with a BasicClient."""
+            """Invalid subclass that warns the user that it expects this class to be mixed with a FlexibleClient."""
 
             pass
 
@@ -152,7 +153,7 @@ def test_get_parameters_uninitialized(mock_param_exchanger: MagicMock, mock_setu
 
 
 @patch("fl4health.mixins.adaptive_drift_constrained.log")
-@patch.object(_TestBasicClient, "set_parameters")
+@patch.object(_TestFlexibleClient, "set_parameters")
 def test_set_parameters(mock_super_set_parameters: MagicMock, mock_logger: MagicMock) -> None:
     # setup client
     client = _TestAdaptedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
@@ -184,7 +185,7 @@ def test_set_parameters(mock_super_set_parameters: MagicMock, mock_logger: Magic
 
 
 def test_dynamically_created_class() -> None:
-    adapted_class = apply_adaptive_drift_to_client(_TestBasicClient)
+    adapted_class = apply_adaptive_drift_to_client(_TestFlexibleClient)
 
     client = adapted_class(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
     client.model = torch.nn.Linear(5, 5)
@@ -195,5 +196,50 @@ def test_dynamically_created_class() -> None:
     client.initialized = True
     client.setup_client({})
 
-    assert isinstance(client, BasicClientProtocol)
+    assert isinstance(client, FlexibleClientProtocol)
     assert isinstance(client, AdaptiveDriftConstrainedProtocol)
+
+
+@patch.object(_TestAdaptedClient, "_compute_preds_and_losses")
+@patch.object(_TestAdaptedClient, "compute_penalty_loss")
+@patch.object(_TestAdaptedClient, "_apply_backwards_on_losses_and_take_step")
+def test_train_step(
+    mock_apply_backwards_on_losses_and_take_step: MagicMock,
+    mock_compute_penalty_loss: MagicMock,
+    mock_compute_preds_and_losses: MagicMock,
+) -> None:
+    # setup client
+    client = _TestAdaptedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.model = torch.nn.Linear(5, 5)
+    client.optimizers = {"global": torch.optim.SGD(client.model.parameters(), lr=0.0001)}
+    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
+    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
+    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
+    client.initialized = True
+    client.setup_client({})
+
+    # setup mocks
+    dummy_training_losses = TrainingLosses(
+        backward=torch.ones(5),
+    )
+    dummy_pred_type = {"prediction": torch.ones(5)}
+    dummy_penalty_loss = torch.zeros(5)
+    mock_compute_preds_and_losses.return_value = (dummy_training_losses, dummy_pred_type)
+    mock_compute_penalty_loss.return_value = dummy_penalty_loss
+    mock_apply_backwards_on_losses_and_take_step.side_effect = lambda x, y, z: z
+
+    # act
+    dummy_input = torch.ones(5)
+    dummy_target = torch.zeros(5)
+    # TODO: fix the mixin/protocol typing that leads to mypy complaint
+    result = client.train_step(dummy_input, dummy_target)  # type: ignore
+
+    # assert
+    mock_compute_preds_and_losses.assert_called_once_with(
+        client.model, client.optimizers["global"], dummy_input, dummy_target
+    )
+    mock_compute_penalty_loss.assert_called_once()
+    mock_apply_backwards_on_losses_and_take_step.assert_called_once_with(
+        client.model, client.optimizers["global"], dummy_training_losses
+    )
+    assert result[1] == dummy_pred_type
diff --git a/tests/smoke_tests/ditto_config_epochs.yaml b/tests/smoke_tests/ditto_config_epochs.yaml
new file mode 100644
index 000000000..a1408afe5
--- /dev/null
+++ b/tests/smoke_tests/ditto_config_epochs.yaml
@@ -0,0 +1,10 @@
+# Parameters that describe server
+n_server_rounds: 1 # The number of rounds to run FL
+
+# Parameters that describe clients
+n_clients: 2 # The number of clients in the FL experiment
+local_epochs: 1 # The number of epochs to complete for client
+batch_size: 32 # The batch size for client training
+
+# Downsampling settings per client
+downsampling_ratio: 0.1 # percentage of original MNIST data to keep for minority numbers
diff --git a/tests/smoke_tests/test_smoke_tests.py b/tests/smoke_tests/test_smoke_tests.py
index c061f12d5..20eaa351d 100644
--- a/tests/smoke_tests/test_smoke_tests.py
+++ b/tests/smoke_tests/test_smoke_tests.py
@@ -297,6 +297,32 @@ async def test_ditto_mnist() -> None:
     assert_on_done_task(task)
 
 
+@pytest.mark.smoketest
+async def test_ditto_mnist_dynamic() -> None:
+    coroutine = run_smoke_test(
+        server_python_path="examples.ditto_example.server",
+        client_python_path="examples.ditto_example.client_dynamic",
+        config_path="tests/smoke_tests/ditto_config.yaml",
+        dataset_path="examples/datasets/mnist_data/",
+    )
+    task = asyncio.create_task(coroutine)
+    await try_running_test_task(task)
+    assert_on_done_task(task)
+
+
+@pytest.mark.smoketest
+async def test_ditto_mnist_dynamic_train_by_epochs() -> None:
+    coroutine = run_smoke_test(
+        server_python_path="examples.ditto_example.server",
+        client_python_path="examples.ditto_example.client_dynamic",
+        config_path="tests/smoke_tests/ditto_config_epochs.yaml",
+        dataset_path="examples/datasets/mnist_data/",
+    )
+    task = asyncio.create_task(coroutine)
+    await try_running_test_task(task)
+    assert_on_done_task(task)
+
+
 @pytest.mark.smoketest
 async def test_mr_mtl_mnist(tolerance: float) -> None:
     coroutine = run_smoke_test(