code: def test_drain_agent() -> None:
      """
      Start an experiment, `disable --drain` the agent once the trial is running,
      make sure the experiment still finishes, but the new ones won't schedule.
      """
  
      slots = _fetch_slots()
      assert len(slots) == 1
      agent_id = slots[0]["agent_id"]
  
      experiment_id = exp.create_experiment(
          conf.fixtures_path("no_op/single-medium-train-step.yaml"),
          conf.fixtures_path("no_op"),
          [
              "--config",
              "hyperparameters.training_batch_seconds=0.15",
          ],  # Take 15 seconds.
      )
      exp.wait_for_experiment_state(
          experiment_id, experimentv1State.RUNNING, max_wait_secs=300
      )
      exp.wait_for_experiment_active_workload(experiment_id)
      exp.wait_for_experiment_workload_progress(experiment_id)
  
      # Disable and quickly enable it back.
      with _disable_agent(agent_id, drain=True):
          pass
  
      # Try to launch another experiment. It shouldn't get scheduled because the
      # slot is still busy with the first experiment.
      experiment_id_no_start = exp.create_experiment(
          conf.fixtures_path("no_op/single-medium-train-step.yaml"),
          conf.fixtures_path("no_op"),
          None,
      )
      time.sleep(5)
      exp.wait_for_experiment_state(
          experiment_id_no_start, experimentv1State.QUEUED
      )
  
      with _disable_agent(agent_id, drain=True):
          # Ensure the first one has finished with the correct number of workloads.
          exp.wait_for_experiment_state(
              experiment_id, experimentv1State.COMPLETED
          )
          trials = exp.experiment_trials(experiment_id)
          assert len(trials) == 1
          assert len(trials[0].workloads) == 7
  
          # Check for 15 seconds it doesn't get scheduled into the same slot.
          for _ in range(15):
              assert (
                  exp.experiment_state(experiment_id_no_start)
                  == experimentv1State.QUEUED
              )
              time.sleep(1)
  
          # Ensure the slot is empty.
          slots = _fetch_slots()
          assert len(slots) == 1
          assert slots[0]["enabled"] is False
          assert slots[0]["draining"] is True
          assert slots[0]["allocation_id"] == "FREE"
  
          # Check agent state.
          command = [
              "det",
              "-m",
              conf.make_master_url(),
              "agent",
              "list",
              "--json",
          ]
          output = subprocess.check_output(command).decode()
          agent_data = cast(List[Dict[str, Any]], json.loads(output))[0]
          assert agent_data["id"] == agent_id
          assert agent_data["enabled"] is False
          assert agent_data["draining"] is True
  
          exp.kill_single(experiment_id_no_start)
  
imports: import contextlib
  import json
  import subprocess
  import time
  from typing import Any, Dict, Iterator, List, Optional, cast
  import pytest
  from determined.common.api.bindings import experimentv1State, get_GetSlot
  from tests import api_utils
  from tests import config
  from tests import experiment
  from .test_users import logged_in_user
  from .utils import assert_command_succeeded, run_zero_slot_command, wait_for_command_state
label_desc: Use image node_id 7 for calling a global function with example usage: test_drain_agent() without return types