From a8eba33f9d9ab3c5cbf40ec8e8570a1961ab3642 Mon Sep 17 00:00:00 2001
From: Zac Davies <zac@databricks.com>
Date: Tue, 24 Dec 2024 16:29:05 +1100
Subject: [PATCH] Minor fixes and test adjustment

---
 R/chattr-use.R                   | 4 ++--
 tests/testthat/test-chattr-use.R | 6 +++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/R/chattr-use.R b/R/chattr-use.R
index 4184a69..71335f2 100644
--- a/R/chattr-use.R
+++ b/R/chattr-use.R
@@ -97,8 +97,8 @@ ch_get_ymls <- function(menu = TRUE) {
   }
 
   if (!dbrx_exists) {
-    prep_files$`databricks-meta-llama-3-1-405b-instruct` <- NULL
-    prep_files$`databricks-meta-llama-3-3-70b-instruct` <- NULL
+    prep_files$`databricks-meta-llama-3-1-405b` <- NULL
+    prep_files$`databricks-meta-llama-3-3-70b` <- NULL
   }
 
   if (!llama_exists) {
diff --git a/tests/testthat/test-chattr-use.R b/tests/testthat/test-chattr-use.R
index 5ecbc8b..76b7e3e 100644
--- a/tests/testthat/test-chattr-use.R
+++ b/tests/testthat/test-chattr-use.R
@@ -31,7 +31,11 @@ test_that("Missing token prevents showing the option", {
 test_that("Menu works", {
   skip_on_cran()
   withr::with_envvar(
-    new = c("OPENAI_API_KEY" = "test", "DATABRICKS_TOKEN" = NA),
+    new = c(
+      "OPENAI_API_KEY" = "test",
+      "DATABRICKS_HOST" = NA,
+      "DATABRICKS_TOKEN" = NA
+      ),
     {
       local_mocked_bindings(
         menu = function(...) {