ruby-sdk/test/model_context_protocol/prompt_test.rb

# typed: strict
# frozen_string_literal: true

require "test_helper"

module ModelContextProtocol
  class PromptTest < ActiveSupport::TestCase
    class TestPrompt < Prompt
      description "Test prompt"
      arguments [
        Prompt::Argument.new(name: "test_argument", description: "Test argument", required: true),
      ]

      def template(args)
        Prompt::Result.new(
          description: "Hello, world!",
          messages: [
            Prompt::Message.new(role: "user", content: Content::Text.new("Hello, world!")),
            Prompt::Message.new(role: "assistant", content: Content::Text.new("Hello, friend!")),
          ],
        )
      end
    end

    test "#template returns a Result with description and messages" do
      prompt = TestPrompt.new

      expected_template_result = {
        description: "Hello, world!",
        messages: [
          { role: "user", content: { text: "Hello, world!" } },
          { role: "assistant", content: { text: "Hello, friend!" } },
        ],
      }

      result = prompt.template({ "test_argument" => "Hello, friend!" })

      assert_equal expected_template_result, result.to_h
    end

    test "allows declarative definition of prompts as classes" do
      class MockPrompt < Prompt
        prompt_name "my_mock_prompt"
        description "a mock prompt for testing"
        arguments [
          Prompt::Argument.new(name: "test_argument", description: "Test argument", required: true),
        ]

        def template(args)
          Prompt::Result.new(
            description: "Hello, world!",
            messages: [
              Prompt::Message.new(role: "user", content: Content::Text.new("Hello, world!")),
              Prompt::Message.new(role: "assistant", content: Content::Text.new(args["test_argument"])),
            ],
          )
        end
      end

      prompt = MockPrompt.new

      assert_equal "my_mock_prompt", prompt.name
      assert_equal "a mock prompt for testing", prompt.description
      assert_equal "test_argument", prompt.arguments.first.name
      assert_equal "Test argument", prompt.arguments.first.description
      assert_equal true, prompt.arguments.first.required

      expected_template_result = {
        description: "Hello, world!",
        messages: [
          { role: "user", content: { text: "Hello, world!" } },
          { role: "assistant", content: { text: "Hello, friend!" } },
        ],
      }

      result = prompt.template({ "test_argument" => "Hello, friend!" })
      assert_equal expected_template_result, result.to_h
    end

    test "defaults to class name as prompt name" do
      class DefaultNamePrompt < Prompt
        description "a mock prompt for testing"
        arguments [
          Prompt::Argument.new(name: "test_argument", description: "Test argument", required: true),
        ]

        def template(args)
          Prompt::Result.new(
            description: "Hello, world!",
            messages: [
              Prompt::Message.new(role: "user", content: Content::Text.new("Hello, world!")),
              Prompt::Message.new(role: "assistant", content: Content::Text.new(args["test_argument"])),
            ],
          )
        end
      end

      prompt = DefaultNamePrompt.new

      assert_equal "default_name_prompt", prompt.name
      assert_equal "a mock prompt for testing", prompt.description
      assert_equal "test_argument", prompt.arguments.first.name
    end

    test ".define allows definition of simple prompts with a block" do
      prompt = Prompt.define(
        name: "mock_prompt",
        description: "a mock prompt for testing",
        arguments: [
          Prompt::Argument.new(name: "test_argument", description: "Test argument", required: true),
        ],
      ) do |args|
        Prompt::Result.new(
          description: "Hello, world!",
          messages: [
            Prompt::Message.new(role: "user", content: Content::Text.new("Hello, world!")),
            Prompt::Message.new(role: "assistant", content: Content::Text.new(args["test_argument"])),
          ],
        )
      end

      assert_equal "mock_prompt", prompt.name
      assert_equal "a mock prompt for testing", prompt.description
      assert_equal "test_argument", prompt.arguments.first.name

      expected = {
        description: "Hello, world!",
        messages: [
          { role: "user", content: { text: "Hello, world!" } },
          { role: "assistant", content: { text: "Hello, friend!" } },
        ],
      }

      result = prompt.template({ "test_argument" => "Hello, friend!" })
      assert_equal expected, result.to_h
    end
  end
end