From 0da322f8e9000a9804e782daf367ec8e7ee868aa Mon Sep 17 00:00:00 2001
From: Jack Tysoe <jack.tysoe@konghq.com>
Date: Wed, 27 Nov 2024 00:40:35 +0000
Subject: [PATCH] fix(AG-178): add logging tests for streaming

(cherry picked from commit bd526e5ef1bcce2e641970da9ea4adf8f2f8c9a0)
---
 .../09-streaming_integration_spec.lua         | 61 +++++++++++++++++++
 .../11-gemini_integration_spec.lua            |  2 +-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/spec/03-plugins/38-ai-proxy/09-streaming_integration_spec.lua b/spec/03-plugins/38-ai-proxy/09-streaming_integration_spec.lua
index d7c89b79b0b7..5cb8a94ed531 100644
--- a/spec/03-plugins/38-ai-proxy/09-streaming_integration_spec.lua
+++ b/spec/03-plugins/38-ai-proxy/09-streaming_integration_spec.lua
@@ -1,12 +1,59 @@
 local helpers = require "spec.helpers"
 local cjson = require "cjson.safe"
 local pl_file = require "pl.file"
+local strip = require("kong.tools.string").strip
 
 local http = require("resty.http")
 
 local PLUGIN_NAME = "ai-proxy"
 local MOCK_PORT = helpers.get_available_port()
 
+local FILE_LOG_PATH_WITH_PAYLOADS = os.tmpname()
+
+local _EXPECTED_CHAT_STATS = {
+  meta = {
+    plugin_id = '6e7c40f6-ce96-48e4-a366-d109c169e444',
+    provider_name = 'openai',
+    request_model = 'gpt-3.5-turbo',
+    response_model = 'gpt-3.5-turbo',
+    llm_latency = 1
+  },
+  usage = {
+    prompt_tokens = 18,
+    completion_tokens = 7,
+    total_tokens = 25,
+    time_per_token = 1,
+    cost = 0.00037,
+  },
+}
+
+local truncate_file = function(path)
+  local file = io.open(path, "w")
+  file:close()
+end
+
+local function wait_for_json_log_entry(FILE_LOG_PATH)
+  local json
+
+  assert
+    .with_timeout(10)
+    .ignore_exceptions(true)
+    .eventually(function()
+      local data = assert(pl_file.read(FILE_LOG_PATH))
+
+      data = strip(data)
+      assert(#data > 0, "log file is empty")
+
+      data = data:match("%b{}")
+      assert(data, "log file does not contain JSON")
+
+      json = cjson.decode(data)
+    end)
+    .has_no_error("log file contains a valid JSON entry")
+
+  return json
+end
+
 for _, strategy in helpers.all_strategies() do
   describe(PLUGIN_NAME .. ": (access) [#" .. strategy  .. "]", function()
     local client
@@ -353,6 +400,7 @@ for _, strategy in helpers.all_strategies() do
       })
       bp.plugins:insert {
         name = PLUGIN_NAME,
+        id = "6e7c40f6-ce96-48e4-a366-d109c169e444",
         route = { id = openai_chat_partial.id },
         config = {
           route_type = "llm/v1/chat",
@@ -360,6 +408,10 @@ for _, strategy in helpers.all_strategies() do
             header_name = "Authorization",
             header_value = "Bearer openai-key",
           },
+          logging = {
+            log_payloads = true,
+            log_statistics = true,
+          },
           model = {
             name = "gpt-3.5-turbo",
             provider = "openai",
@@ -371,6 +423,13 @@ for _, strategy in helpers.all_strategies() do
           },
         },
       }
+      bp.plugins:insert {
+        name = "file-log",
+        route = { id = openai_chat_partial.id },
+        config = {
+          path = FILE_LOG_PATH_WITH_PAYLOADS,
+        },
+      }
       --
 
       -- 200 chat cohere
@@ -497,10 +556,12 @@ for _, strategy in helpers.all_strategies() do
 
     lazy_teardown(function()
       helpers.stop_kong()
+      os.remove(FILE_LOG_PATH_WITH_PAYLOADS)
     end)
 
     before_each(function()
       client = helpers.proxy_client()
+      truncate_file(FILE_LOG_PATH_WITH_PAYLOADS)
     end)
 
     after_each(function()
diff --git a/spec/03-plugins/38-ai-proxy/11-gemini_integration_spec.lua b/spec/03-plugins/38-ai-proxy/11-gemini_integration_spec.lua
index 76c80d541ebb..1e1bce8081bb 100644
--- a/spec/03-plugins/38-ai-proxy/11-gemini_integration_spec.lua
+++ b/spec/03-plugins/38-ai-proxy/11-gemini_integration_spec.lua
@@ -219,4 +219,4 @@ for _, strategy in helpers.all_strategies() do
       end)
     end)
   end
-end
+end
\ No newline at end of file