1  use std::path::{Path, PathBuf};
  2  
  3  use apibara_sink_common::{ScriptOptions, StreamOptions};
  4  use clap::Args;
  5  use error_stack::{Result, ResultExt};
  6  use tracing::warn;
  7  
  8  use crate::error::CliError;
  9  
 10  mod error;
 11  mod run;
 12  mod snapshot;
 13  
 14  pub const SNAPSHOTS_DIR: &str = "snapshots";
 15  
 16  #[derive(Args, Debug)]
 17  pub struct TestArgs {
 18      /// An indexer script (.js/.ts), a snapshot file (.json) or a folder of snapshots.
 19      path: Option<PathBuf>,
 20      /// The number of blocks to stream.
 21      #[arg(long, short = 'b')]
 22      num_batches: Option<usize>,
 23      /// Regenerate the snapshot even if it already exists.
 24      #[arg(long, short = 'o', default_value_t = false)]
 25      r#override: bool,
 26      /// The name of the snapshot.
 27      #[arg(long, short = 'n')]
 28      name: Option<String>,
 29      /// Override the starting block from the script.
 30      #[arg(long, short, env)]
 31      starting_block: Option<u64>,
 32      #[clap(flatten)]
 33      stream_options: StreamOptions,
 34      #[clap(flatten)]
 35      dotenv_options: ScriptOptions,
 36  }
 37  
 38  fn validate_args(args: &TestArgs) -> Result<(), CliError> {
 39      if let Some(name) = &args.name {
 40          let is_invalid = name.contains(|c| {
 41              c == '/'
 42                  || c == '.'
 43                  || c == '\\'
 44                  || c == ':'
 45                  || c == '*'
 46                  || c == '?'
 47                  || c == '"'
 48                  || c == '<'
 49                  || c == '>'
 50                  || c == '|'
 51          });
 52          if is_invalid {
 53              return Err(CliError).attach_printable(
 54                  r#"Invalid name `{name}`, name should not contain  /, ., \, :, *, ?, ", <, >, or | as it'll be used to construct the snapshot path"#
 55              );
 56          }
 57      }
 58  
 59      if let Some(num_batches) = args.num_batches {
 60          if num_batches < 1 {
 61              return Err(CliError).attach_printable_lazy(|| {
 62                  format!("Invalid number of blocks `{num_batches}`, it should be > 0")
 63              });
 64          }
 65      }
 66  
 67      if let Some(path) = &args.path {
 68          // Think about using try_exists instead
 69          if !path.exists() {
 70              return Err(CliError).attach_printable_lazy(|| {
 71                  format!(
 72                      "Invalid path: `{}`, no such file or directory",
 73                      &path.display()
 74                  )
 75              });
 76          }
 77      }
 78  
 79      Ok(())
 80  }
 81  
 82  pub fn warn_ignored_args(args: &TestArgs) {
 83      if args.starting_block.is_some()
 84          || args.num_batches.is_some()
 85          || args.r#override
 86          || args.stream_options.stream_url.is_some()
 87          || args.stream_options.max_message_size.is_some()
 88          || args.stream_options.metadata.is_some()
 89      {
 90          warn!(
 91              "The following arguments are ignored: --starting-block, --num-batches, \
 92              --override, --stream-url, --max-message-size, --metadata when running tests, \
 93              if you want to generate a snapshot with different options, \
 94              use the --override flag or give it a name with --name"
 95          )
 96      }
 97  }
 98  
 99  pub async fn run(args: TestArgs) -> Result<(), CliError> {
100      validate_args(&args)?;
101  
102      match &args.path {
103          Some(path) => {
104              // Think about using try_exists instead
105              if !path.exists() {
106                  return Err(CliError).attach_printable_lazy(|| {
107                      format!(
108                          "Invalid path: `{}`, no such file or directory",
109                          &path.display()
110                      )
111                  });
112              }
113  
114              if path.is_dir() {
115                  warn_ignored_args(&args);
116                  run::run_all_tests(&path, &args.dotenv_options, None).await?;
117                  return Ok(());
118              }
119  
120              let extension = Path::new(&path)
121                  .extension()
122                  .ok_or(CliError)
123                  .attach_printable_lazy(|| format!("Invalid path: `{}`", path.display()))?;
124  
125              match extension.to_str().unwrap() {
126                  "json" => {
127                      warn_ignored_args(&args);
128                      run::run_single_test(path, None, None, &args.dotenv_options).await?;
129                  },
130                  "js" | "ts" => {
131                      let snapshot_path = args.name.clone()
132                          .map(|name| Path::new(SNAPSHOTS_DIR).join(name).with_extension("json"))
133                          .unwrap_or(snapshot::get_snapshot_path(path)?);
134  
135                      if args.r#override || !snapshot_path.exists() {
136                          run::run_generate_snapshot(
137                              path,
138                              &snapshot_path,
139                              args.starting_block,
140                              args.num_batches,
141                              &args.stream_options,
142                              &args.dotenv_options,
143                          ).await?;
144                      } else {
145                          warn_ignored_args(&args);
146                          if args.name.is_some() {
147                              run::run_single_test(&snapshot_path, None, Some(path), &args.dotenv_options).await?;
148                          } else {
149                              run::run_all_tests(SNAPSHOTS_DIR, &args.dotenv_options, Some(path)).await?;
150                          }
151                      }
152                  }
153                  _ => return Err(CliError).attach_printable_lazy(|| format!(
154                      "Invalid file extension: `{}`, must be a .json for snapshots or .js / .ts for scripts",
155                      path.display()
156                  )),
157              }
158          }
159          None => {
160              warn_ignored_args(&args);
161              run::run_all_tests(SNAPSHOTS_DIR, &args.dotenv_options, None).await?;
162          }
163      }
164  
165      Ok(())
166  }