From 9515467439255d9868af9a509e0ccb4f5feee925 Mon Sep 17 00:00:00 2001
From: Andrei Betlen <abetlen@gmail.com>
Date: Wed, 22 Nov 2023 06:02:21 -0500
Subject: [PATCH] tests: add mock_kv_cache placeholder functions

---
 tests/test_llama.py | 50 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/tests/test_llama.py b/tests/test_llama.py
index 9ce2a2a..c98148e 100644
--- a/tests/test_llama.py
+++ b/tests/test_llama.py
@@ -96,6 +96,56 @@ def mock_llama(monkeypatch):
         monkeypatch.setattr("llama_cpp.llama_cpp.llama_decode", mock_decode)
         monkeypatch.setattr("llama_cpp.llama_cpp.llama_get_logits", mock_get_logits)
 
+        def mock_kv_cache_clear(ctx: llama_cpp.llama_context_p):
+            # Test some basic invariants of this mocking technique
+            assert ctx == llama._ctx.ctx, "context does not match mock_llama"
+            return
+
+        def mock_kv_cache_seq_rm(
+            ctx: llama_cpp.llama_context_p,
+            seq_id: llama_cpp.llama_seq_id,
+            pos0: llama_cpp.llama_pos,
+            pos1: llama_cpp.llama_pos,
+        ):
+            # Test some basic invariants of this mocking technique
+            assert ctx == llama._ctx.ctx, "context does not match mock_llama"
+            return
+
+        def mock_kv_cache_seq_cp(
+            ctx: llama_cpp.llama_context_p,
+            seq_id_src: llama_cpp.llama_seq_id,
+            seq_id_dst: llama_cpp.llama_seq_id,
+            pos0: llama_cpp.llama_pos,
+            pos1: llama_cpp.llama_pos,
+        ):
+            # Test some basic invariants of this mocking technique
+            assert ctx == llama._ctx.ctx, "context does not match mock_llama"
+            return
+    
+        def mock_kv_cache_seq_keep(
+            ctx: llama_cpp.llama_context_p,
+            seq_id: llama_cpp.llama_seq_id,
+        ):
+            # Test some basic invariants of this mocking technique
+            assert ctx == llama._ctx.ctx, "context does not match mock_llama"
+            return
+
+        def mock_kv_cache_seq_shift(
+            ctx: llama_cpp.llama_context_p,
+            seq_id: llama_cpp.llama_seq_id,
+            pos0: llama_cpp.llama_pos,
+            pos1: llama_cpp.llama_pos,
+        ):
+            # Test some basic invariants of this mocking technique
+            assert ctx == llama._ctx.ctx, "context does not match mock_llama"
+            return
+
+        monkeypatch.setattr("llama_cpp.llama_cpp.llama_kv_cache_clear", mock_kv_cache_clear)
+        monkeypatch.setattr("llama_cpp.llama_cpp.llama_kv_cache_seq_rm", mock_kv_cache_seq_rm)
+        monkeypatch.setattr("llama_cpp.llama_cpp.llama_kv_cache_seq_cp", mock_kv_cache_seq_cp)
+        monkeypatch.setattr("llama_cpp.llama_cpp.llama_kv_cache_seq_keep", mock_kv_cache_seq_keep)
+        monkeypatch.setattr("llama_cpp.llama_cpp.llama_kv_cache_seq_shift", mock_kv_cache_seq_shift)
+
     return setup_mock