gateway_service_bedrock_model_support_test.go 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. package service
  2. import "testing"
  3. func TestGatewayServiceIsModelSupportedByAccount_BedrockDefaultMappingRestrictsModels(t *testing.T) {
  4. svc := &GatewayService{}
  5. account := &Account{
  6. Platform: PlatformAnthropic,
  7. Type: AccountTypeBedrock,
  8. Credentials: map[string]any{
  9. "aws_region": "us-east-1",
  10. },
  11. }
  12. if !svc.isModelSupportedByAccount(account, "claude-sonnet-4-5") {
  13. t.Fatalf("expected default Bedrock alias to be supported")
  14. }
  15. if svc.isModelSupportedByAccount(account, "claude-3-5-sonnet-20241022") {
  16. t.Fatalf("expected unsupported alias to be rejected for Bedrock account")
  17. }
  18. }
  19. func TestGatewayServiceIsModelSupportedByAccount_BedrockCustomMappingStillActsAsAllowlist(t *testing.T) {
  20. svc := &GatewayService{}
  21. account := &Account{
  22. Platform: PlatformAnthropic,
  23. Type: AccountTypeBedrock,
  24. Credentials: map[string]any{
  25. "aws_region": "eu-west-1",
  26. "model_mapping": map[string]any{
  27. "claude-sonnet-*": "claude-sonnet-4-6",
  28. },
  29. },
  30. }
  31. if !svc.isModelSupportedByAccount(account, "claude-sonnet-4-6") {
  32. t.Fatalf("expected matched custom mapping to be supported")
  33. }
  34. if !svc.isModelSupportedByAccount(account, "claude-opus-4-6") {
  35. t.Fatalf("expected default Bedrock alias fallback to remain supported")
  36. }
  37. if svc.isModelSupportedByAccount(account, "claude-3-5-sonnet-20241022") {
  38. t.Fatalf("expected unsupported model to still be rejected")
  39. }
  40. }