R/gen_tests.R

In gemini.R: Interface for 'Google Gemini' API

#' @title Generate Unit Tests for R Functions
#'
#' @description Generates unit test code for an R function using the Gemini AI model.
#'
#' @param prompt A character string specifying additional instructions for the Gemini model.
#'        If NULL, a default prompt requesting unit tests is used.
#'
#' @return Invisibly returns the generated test code, but primarily inserts it into the RStudio console.
#'
#' @examples
#' \dontrun{
#' # Select your function code in the editor, then run:
#' gen_tests()
#' 
#' # For custom instructions:
#' gen_tests("Generate comprehensive testthat tests with edge cases")
#' }
#'
#' @importFrom rstudioapi getActiveDocumentContext executeCommand insertText
#' @importFrom cli cli_alert_danger cli_status cli_status_clear
#'
#' @export
gen_tests <- function(prompt = NULL) {
  if (is.null(prompt)) {
    prompt <-
      paste0(
        "Generate unit test code for an R function.",
        "\nReturn only R code without markdown code blocks.",
        "\nDo not include ```r and ``` markers in the response.",
        "\nPreferably use testthat framework for testing."
      )
  }

  # Get selected code from the editor
  context <- tryCatch({
    getActiveDocumentContext()
  }, error = function(e) {
    cli_alert_danger("Failed to get active document context: ", e$message)
    return(NULL)
  })
  
  if (is.null(context)) {
    return(invisible(NULL))
  }
  
  selectedCode <- context$selection[[1]]$text
  
  # Check if code is selected
  if (is.null(selectedCode) || nchar(trim(selectedCode)) == 0) {
    cli_alert_danger("No code selected. Please select an R function to generate tests for.")
    return(invisible(NULL))
  }

  # Show status message
  sb <- cli_status("Generating unit tests...")
  
  # API call and error handling
  test_code <- tryCatch({
    result <- gemini(
      prompt = paste0(
        prompt,
        "\n---\n",
        selectedCode
      )
    )
    
    if (is.null(result) || length(result) == 0 || nchar(trim(result)) == 0) {
      cli_status_clear(id = sb)
      cli_alert_danger("Failed to generate test code or received empty response.")
      return(invisible(NULL))
    }
    
    result
  }, 
  error = function(e) {
    cli_status_clear(id = sb)
    cli_alert_danger(paste0("Error generating test code: ", e$message))
    return(invisible(NULL))
  })
  
  # Early exit if error occurred
  if (is.null(test_code)) {
    return(invisible(NULL))
  }
  
  cli_status_clear(id = sb)
  
  # Insert test code into the console
  tryCatch({
    executeCommand("activateConsole")
    insertText(text = test_code)
  }, error = function(e) {
    cli_alert_danger(paste0("Failed to insert text into console: ", e$message))
  })
  
  # Return generated test code invisibly
  return(invisible(test_code))
}