From 6c205e1d5a2ba38e3773bb4cab4bcc5aa43f5dcb Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Wed, 22 Jan 2025 20:31:18 -0800
Subject: [PATCH] Fix tool tests

---
 llama_stack/providers/tests/tools/conftest.py   | 7 +++----
 llama_stack/providers/tests/tools/test_tools.py | 2 +-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/llama_stack/providers/tests/tools/conftest.py b/llama_stack/providers/tests/tools/conftest.py
index 525abe8ab..0df547a9d 100644
--- a/llama_stack/providers/tests/tools/conftest.py
+++ b/llama_stack/providers/tests/tools/conftest.py
@@ -8,8 +8,8 @@ import pytest
 
 from ..conftest import get_provider_fixture_overrides
 from ..inference.fixtures import INFERENCE_FIXTURES
-from ..memory.fixtures import MEMORY_FIXTURES
 from ..safety.fixtures import SAFETY_FIXTURES
+from ..vector_io.fixtures import VECTOR_IO_FIXTURES
 from .fixtures import TOOL_RUNTIME_FIXTURES
 
 DEFAULT_PROVIDER_COMBINATIONS = [
@@ -17,7 +17,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
         {
             "inference": "together",
             "safety": "llama_guard",
-            "memory": "faiss",
+            "vector_io": "faiss",
             "tool_runtime": "memory_and_search",
         },
         id="together",
@@ -39,12 +39,11 @@ def pytest_generate_tests(metafunc):
         available_fixtures = {
             "inference": INFERENCE_FIXTURES,
             "safety": SAFETY_FIXTURES,
-            "memory": MEMORY_FIXTURES,
+            "vector_io": VECTOR_IO_FIXTURES,
             "tool_runtime": TOOL_RUNTIME_FIXTURES,
         }
         combinations = (
             get_provider_fixture_overrides(metafunc.config, available_fixtures)
             or DEFAULT_PROVIDER_COMBINATIONS
         )
-        print(combinations)
         metafunc.parametrize("tools_stack", combinations, indirect=True)
diff --git a/llama_stack/providers/tests/tools/test_tools.py b/llama_stack/providers/tests/tools/test_tools.py
index bb4265f94..281ea404d 100644
--- a/llama_stack/providers/tests/tools/test_tools.py
+++ b/llama_stack/providers/tests/tools/test_tools.py
@@ -88,7 +88,7 @@ class TestTools:
         tools_impl = tools_stack.impls[Api.tool_runtime]
 
         # Register memory bank
-        await vector_dbs_impl.register(
+        await vector_dbs_impl.register_vector_db(
             vector_db_id="test_bank",
             embedding_model="all-MiniLM-L6-v2",
             embedding_dimension=384,