code: def test_xgb_classification_learner():
      df_train_binary = pd.DataFrame(
          {
              "id": ["id1", "id2", "id3", "id4"],
              "x1": [10.0, 13.0, 10.0, 13.0],
              "x2": [0, 1, 1, 0],
              "w": [2, 1, 2, 0.5],
              "y": [0, 1, 0, 1],
          }
      )
  
      df_test_binary = pd.DataFrame(
          {
              "id": ["id4", "id4", "id5", "id6"],
              "x1": [12.0, 1000.0, -4.0, 0.0],
              "x2": [1, 1, 0, 1],
              "w": [1, 2, 0, 0.5],
              "y": [1, 0, 0, 1],
          }
      )
  
      df_train_multinomial = pd.DataFrame(
          {
              "id": ["id1", "id2", "id3", "id4", "id3", "id4"],
              "x1": [10.0, 13.0, 10.0, 13.0, 10.0, 13.0],
              "x2": [0, 1, 1, 0, 1, 0],
              "w": [2, 1, 2, 0.5, 2, 0.5],
              "y": [0, 1, 2, 1, 2, 0],
          }
      )
  
      df_test_multinomial = pd.DataFrame(
          {
              "id": ["id4", "id4", "id5", "id6", "id5", "id6"],
              "x1": [12.0, 1000.0, -4.0, 0.0, -4.0, 0.0],
              "x2": [1, 1, 0, 1, 0, 1],
              "w": [1, 2, 0, 0.5, 0, 0.5],
              "y": [1, 2, 0, 1, 2, 0],
          }
      )
  
      features = ["x1", "x2"]
  
      learner_binary = xgb_classification_learner(
          features=features,
          target="y",
          learning_rate=0.1,
          num_estimators=20,
          extra_params={"max_depth": 4, "seed": 42},
          prediction_column="prediction",
          weight_column="w",
      )
  
      predict_fn_binary, pred_train_binary, log = learner_binary(
          df_train_binary
      )
  
      pred_test_binary = predict_fn_binary(df_test_binary)
  
      expected_col_train = df_train_binary.columns.tolist() + [
          "prediction"
      ]
      expected_col_test = df_test_binary.columns.tolist() + [
          "prediction"
      ]
  
      assert Counter(expected_col_train) == Counter(
          pred_train_binary.columns.tolist()
      )
      assert Counter(expected_col_test) == Counter(
          pred_test_binary.columns.tolist()
      )
      assert pred_test_binary.prediction.max() < 1
      assert pred_test_binary.prediction.min() > 0
      assert (
          pred_test_binary.columns == pred_train_binary.columns
      ).all()
  
      # SHAP test
      pred_shap = predict_fn_binary(df_test_binary, apply_shap=True)
      assert "shap_values" in pred_shap.columns
      assert "shap_expected_value" in pred_shap.columns
      assert np.vstack(pred_shap["shap_values"]).shape == (4, 2)
  
      # test multinomial case
      learner_multinomial = xgb_classification_learner(
          features=features,
          target="y",
          learning_rate=0.1,
          num_estimators=20,
          extra_params={
              "max_depth": 2,
              "seed": 42,
              "objective": "multi:softprob",
              "num_class": 3,
          },
          prediction_column="prediction",
      )
  
      (
          predict_fn_multinomial,
          pred_train_multinomial,
          log,
      ) = learner_multinomial(df_train_multinomial)
  
      pred_test_multinomial = predict_fn_multinomial(
          df_test_multinomial
      )
  
      expected_col_train = df_train_binary.columns.tolist() + [
          "prediction_0",
          "prediction_1",
          "prediction_2",
          "prediction",
      ]
      expected_col_test = df_test_binary.columns.tolist() + [
          "prediction_0",
          "prediction_1",
          "prediction_2",
          "prediction",
      ]
  
      assert Counter(expected_col_train) == Counter(
          pred_train_multinomial.columns.tolist()
      )
      assert Counter(expected_col_test) == Counter(
          pred_test_multinomial.columns.tolist()
      )
      assert (
          pred_test_multinomial.columns
          == pred_train_multinomial.columns
      ).all()
  
      # SHAP test multinomial
      pred_shap_multinomial = predict_fn_multinomial(
          df_test_multinomial, apply_shap=True
      )
  
      expected_col_shap = (
          expected_col_test
          + ["shap_values_0", "shap_values_1", "shap_values_2"]
          + [
              "shap_expected_value_0",
              "shap_expected_value_1",
              "shap_expected_value_2",
          ]
      )
      assert Counter(expected_col_shap) == Counter(
          pred_shap_multinomial.columns.tolist()
      )
      assert np.vstack(
          pred_shap_multinomial["shap_values_0"]
      ).shape == (6, 2)
  
file_path: fklearn/tests/training/test_classification.py
imports: from collections import Counter
  import numpy
  import pandas
  from fklearn.training.classification import logistic_classification_learner, xgb_classification_learner, nlp_logistic_classification_learner, lgbm_classification_learner, catboost_classification_learner
  from unittest.mock import MagicMock, patch, Mock
label_desc: Use image node_id 2 for calling a global function with example usage: test_xgb_classification_learner() without return types