jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-03-02 23:51:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix model-router plugin (#1432)

Browse Source
This commit is contained in:
澄潭
2024-10-25 16:48:37 +08:00
committed by GitHub
parent 496d365a95
commit d309bf2e25
2 changed files with 53 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
plugins/wasm-cpp/extensions/model_router/plugin_test.cc
View File

@@ -40,6 +40,11 @@ class MockContext : public proxy_wasm::ContextBase {
MOCK_METHOD(WasmResult, getHeaderMapValue,
(WasmHeaderMapType /* type */, std::string_view /* key */,
std::string_view* /*result */));
MOCK_METHOD(WasmResult, replaceHeaderMapValue,
(WasmHeaderMapType /* type */, std::string_view /* key */,
std::string_view /* value */));
MOCK_METHOD(WasmResult, removeHeaderMapValue,
(WasmHeaderMapType /* type */, std::string_view /* key */));
MOCK_METHOD(WasmResult, addHeaderMapValue,
(WasmHeaderMapType, std::string_view, std::string_view));
MOCK_METHOD(WasmResult, getProperty, (std::string_view, std::string*));
@@ -87,6 +92,16 @@ class ModelRouterTest : public ::testing::Test {
}
return WasmResult::Ok;
});
ON_CALL(*mock_context_,
replaceHeaderMapValue(WasmHeaderMapType::RequestHeaders, testing::_,
testing::_))
.WillByDefault([&](WasmHeaderMapType, std::string_view key,
std::string_view value) { return WasmResult::Ok; });
ON_CALL(*mock_context_,
removeHeaderMapValue(WasmHeaderMapType::RequestHeaders, testing::_))
.WillByDefault([&](WasmHeaderMapType, std::string_view key) {
return WasmResult::Ok;
});
ON_CALL(*mock_context_, addHeaderMapValue(WasmHeaderMapType::RequestHeaders,
testing::_, testing::_))
.WillByDefault([&](WasmHeaderMapType, std::string_view header,
@@ -128,10 +143,10 @@ TEST_F(ModelRouterTest, RewriteModelAndHeader) {
return WasmResult::Ok;
});
EXPECT_CALL(
*mock_context_,
addHeaderMapValue(testing::_, std::string_view("x-higress-llm-provider"),
std::string_view("qwen")));
EXPECT_CALL(*mock_context_,
replaceHeaderMapValue(testing::_,
std::string_view("x-higress-llm-provider"),
std::string_view("qwen")));
body_.set(request_json);
EXPECT_EQ(context_->onRequestHeaders(0, false),
@@ -139,6 +154,39 @@ TEST_F(ModelRouterTest, RewriteModelAndHeader) {
EXPECT_EQ(context_->onRequestBody(28, true), FilterDataStatus::Continue);
}
TEST_F(ModelRouterTest, RouteLevelRewriteModelAndHeader) {
std::string configuration = R"(
{
"_rules_": [
{
"_match_route_": ["route-a"],
"enable": true
}
]})";
config_.set(configuration);
EXPECT_TRUE(root_context_->configure(configuration.size()));
std::string request_json = R"({"model": "qwen/qwen-long"})";
EXPECT_CALL(*mock_context_,
setBuffer(testing::_, testing::_, testing::_, testing::_))
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
EXPECT_EQ(body, R"({"model":"qwen-long"})");
return WasmResult::Ok;
});
EXPECT_CALL(*mock_context_,
replaceHeaderMapValue(testing::_,
std::string_view("x-higress-llm-provider"),
std::string_view("qwen")));
body_.set(request_json);
route_name_ = "route-a";
EXPECT_EQ(context_->onRequestHeaders(0, false),
FilterHeadersStatus::StopIteration);
EXPECT_EQ(context_->onRequestBody(28, true), FilterDataStatus::Continue);
}
} // namespace model_router
} // namespace null_plugin
} // namespace proxy_wasm