Generate agents dynamically inside stages in jenkins pipelines

I have this configuration in my pipeline job

def k8sTestPodTemplate(docker_image) {
 return """
 apiVersion: v1
 kind: Pod
 metadata:
   name: my-agent
   labels:
     name: my-agent
 spec:
   serviceAccountName: jenkins
   containers:
     - name: python
       image: ${docker_image}
       command: ["/bin/bash", "-c", "cat"]
       tty: true
"""
}

pipeline {
 agent none

 stages {
   stage('Run tests') {
     parallel {
       stage('Tests Python 3.5') {
         agent {
           kubernetes {
             defaultContainer 'jnlp'
             yaml k8sTestPodTemplate('python:3.5')
           }
         }
         steps {
           container('python') {
             sh "echo 'Hello from Python 3.5'"
           }
         }
       }
       stage('Tests Python 3.6') {
         agent {
           kubernetes {
             defaultContainer 'jnlp'
             yaml k8sTestPodTemplate('python:3.6')
           }
         }
         steps {
           container('python') {
             sh "echo 'Hello from Python 3.6'"
           }
         }
       }
       stage('Tests Python 3.7') {
         agent {
           kubernetes {
             defaultContainer 'jnlp'
             yaml k8sTestPodTemplate('python:3.7')
           }
         }
         steps {
           container('python') {
             sh "echo 'Hello from Python 3.7'"
           }
         }
       }
     }
   }
 }
}

But as you can see I could easily improve this code to something like that:

def k8sTestPodTemplate(docker_image) {
  return """
  apiVersion: v1
  kind: Pod
  metadata:
    name: my-agent
    labels:
      name: my-agent
  spec:
    serviceAccountName: jenkins
    containers:
      - name: python
        image: ${docker_image}
        command: ["/bin/bash", "-c", "cat"]
        tty: true
"""
}

def generateStage(docker_image) {
  return {
    stage("Tests ${docker_image}") {
      agent {
        kubernetes {
          defaultContainer 'jnlp'
          yaml k8sTestPodTemplate("${docker_image}")
        }
      }
      steps {
        container('python') {
          sh "echo ${docker_image}"
        }
      }
    }
  }
}

pipeline {
  agent none

  stages {
    stage('Run tests') {
      parallel {
        generateStage("python:3.5")
        generateStage("python:3.6")
        generateStage("python:3.7")
      }
    }
  }
}

But I cannot get this to work. The problem is that Jenkins is raising an error

No such DSL method 'agent' found among steps

I am using the "agent" directive inside the "step" directive and the agent is being generated dynamically.