From 2121560fd158dde119a0284a771bfb8dca4d1537 Mon Sep 17 00:00:00 2001
From: Justin Traglia <jtraglia@pm.me>
Date: Fri, 11 Oct 2024 17:27:52 -0500
Subject: [PATCH] Add gradle task for downloading nightly reference tests

---
 build.gradle | 98 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 96 insertions(+), 2 deletions(-)

diff --git a/build.gradle b/build.gradle
index d59bfffd00f..35f8c4c7b3f 100644
--- a/build.gradle
+++ b/build.gradle
@@ -4,6 +4,8 @@ import tech.pegasys.teku.depcheck.DepCheckPlugin
 
 import java.text.SimpleDateFormat
 
+import groovy.json.JsonSlurper
+
 import static tech.pegasys.teku.repackage.Repackage.repackage
 
 buildscript {
@@ -318,7 +320,8 @@ allprojects {
   }
 }
 
-def refTestVersion = 'v1.5.0-alpha.8'
+def nightly = System.getenv("NIGHTLY") != null
+def refTestVersion = nightly ? "nightly" : "v1.5.0-alpha.8"
 def blsRefTestVersion = 'v0.1.2'
 def slashingProtectionInterchangeRefTestVersion = 'v5.3.0'
 def refTestBaseUrl = 'https://github.com/ethereum/consensus-spec-tests/releases/download'
@@ -329,7 +332,88 @@ def blsRefTestDownloadDir = "${buildDir}/blsRefTests/${blsRefTestVersion}"
 def slashingProtectionInterchangeRefTestDownloadDir = "${buildDir}/slashingProtectionInterchangeRefTests/${slashingProtectionInterchangeRefTestVersion}"
 def refTestExpandDir = "${project.rootDir}/eth-reference-tests/src/referenceTest/resources/consensus-spec-tests/"
 
-task downloadEthRefTests(type: Download) {
+def downloadFile(String url, String token, File outputFile) {
+  println "Download ${outputFile.getName()} (${url})"
+  def connection = new URL(url).openConnection()
+  connection.setRequestProperty("Authorization", "token ${token}")
+  connection.getInputStream().withCloseable { inputStream ->
+    outputFile.withOutputStream { outputStream ->
+      outputStream << inputStream
+    }
+  }
+}
+
+def downloadArtifacts(String repo, Long runId, String token, String downloadDir) {
+  def artifactsApiUrl = "https://api.github.com/repos/${repo}/actions/runs/${runId}/artifacts"
+  def connection = new URL(artifactsApiUrl).openConnection()
+  connection.setRequestProperty("Authorization", "token ${token}")
+  connection.setRequestProperty("Accept", "application/vnd.github.v3+json")
+
+  def response = new JsonSlurper().parse(connection.getInputStream())
+  if (response.artifacts && response.artifacts.size() > 0) {
+    response.artifacts.each { artifact ->
+      // We can skip the log file
+      if (artifact.name.contains("consensustestgen.log")) {
+        return
+      }
+
+      def fileOutput = new File(downloadDir, "${artifact.name}.zip")
+      downloadFile(artifact.archive_download_url, token, fileOutput)
+      ant.unzip(src: fileOutput, dest: downloadDir)
+      fileOutput.delete()
+    }
+    return true
+  }
+  return false
+}
+
+// Function to get the latest run ID
+static def getLatestRunId(String repo, String workflow, String branch, String token) {
+  def apiUrl = "https://api.github.com/repos/${repo}/actions/workflows/${workflow}/runs?branch=${branch}&status=success&per_page=1"
+  def connection = new URL(apiUrl).openConnection()
+  connection.setRequestProperty("Authorization", "token ${token}")
+  connection.setRequestProperty("Accept", "application/vnd.github.v3+json")
+
+  // Query & parse the ID out of the response
+  def response = new JsonSlurper().parse(connection.getInputStream())
+  if (response.workflow_runs && response.workflow_runs.size() > 0) {
+    return response.workflow_runs[0].id
+  }
+  return null
+}
+
+task downloadEthRefTestsNightly {
+  doLast {
+    def repo = "ethereum/consensus-specs"
+    def workflowFileName = "generate_vectors.yml"
+    def branch = "dev"
+
+    // We need a GitHub API token to download the artifacts
+    def githubToken = System.getenv("GITHUB_TOKEN")
+    if (!githubToken) {
+      println "Error: GITHUB_TOKEN environment variable is not set"
+      return
+    }
+
+    // Get the latest workflow run ID
+    def runId = getLatestRunId(repo, workflowFileName, branch, githubToken)
+    if (!runId) {
+      println "Error: Failed to get latest run ID"
+      return
+    }
+
+    // Create the download directory
+    file(refTestDownloadDir).mkdirs()
+
+    // Download artifacts for the run
+    def success = downloadArtifacts(repo, runId, githubToken, refTestDownloadDir)
+    if (!success) {
+      println "Error: Failed to download artifacts"
+    }
+  }
+}
+
+task downloadEthRefTestsStable(type: Download) {
   src([
       "${refTestBaseUrl}/${refTestVersion}/general.tar.gz",
       "${refTestBaseUrl}/${refTestVersion}/minimal.tar.gz",
@@ -339,6 +423,16 @@ task downloadEthRefTests(type: Download) {
   overwrite false
 }
 
+task downloadEthRefTests {
+  if (nightly) {
+    println "Downloading nightly reference tests"
+    dependsOn tasks.findByName("downloadEthRefTestsNightly")
+  } else {
+    println "Downloading stable reference tests"
+    dependsOn tasks.findByName("downloadEthRefTestsStable")
+  }
+}
+
 task downloadBlsRefTests(type: Download) {
   src([
       "${blsRefTestBaseUrl}/${blsRefTestVersion}/bls_tests_yaml.tar.gz"