Android/Gradle Run a task after Ui tests end

In my gradle file I wanted to move some files from my phone to build directory after the execution of ui tests.

So behavior I want is - assembleDebugAndroidTest - Delete specific file in build directory - Tests end - Move files from phone to build directory

What ends up happening is assembleDebugAndroidTest - Delete specific file in build directory - Move files from phone to build directory - Tests end

I understand the reasoning here, as I add the task assembleDebugAndroidTest, all tasks that are related to it gets executed immediately. But is there a way to withhold the task(or do something else) to get the files after tests end?

Here is my gradle file

def location = "$buildDir/reports/androidTests/connected"
def taskTestName = "assembleDebugAndroidTest"

task wipeReport(type: Delete) {
    delete location
}

def getScreenshots = task('createScreenshotFolder', type: Exec) {
    executable "${android.getAdbExe().toString()}"
    args 'shell', 'mkdir', '-p', '/sdcard/newLocation/screenshots'

    dependsOn {
        wipeReport
    }
}


def pullScreenshots= task('pullScreenshots', type: Exec) { dependsOn
    executable "${android.getAdbExe().toString()}"
    args 'pull', '/sdcard/newLocation/screenshots/.', location

    dependsOn {
        getScreenshots
    }
    finalizedBy {
        //TODO delete screenshots in phone
    }
    doFirst {
        new File(location).mkdirs()
    }
}

tasks.whenTaskAdded { task ->
    if (task.name == taskTestName) {
        task.finalizedBy {
            pullScreenshots
        }
    }
}

Read more here: https://stackoverflow.com/questions/65722874/android-gradle-run-a-task-after-ui-tests-end