# This file is distributed under New Relic's license terms.
# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
# frozen_string_literal: true

require 'mongo'
require 'newrelic_rpm'
require 'new_relic/agent/datastores/mongo'
require 'helpers/mongo_metric_builder'
require_relative '../../../helpers/docker'

def mongo_host
  docker? ? 'mongodb' : NewRelic::Agent::Hostname.get
end

if NewRelic::Agent::Datastores::Mongo.is_supported_version? &&
    NewRelic::Agent::Datastores::Mongo.is_monitoring_enabled?

  require File.join(File.dirname(__FILE__), 'helpers', 'mongo_helpers')

  module NewRelic
    module Agent
      module Instrumentation
        class Mongo2InstrumentationTest < Minitest::Test
          include Mongo
          include TestHelpers::MongoMetricBuilder
          include NewRelic::MongoHelpers

          def setup
            Mongo::Logger.logger = mongo_logger
            @database_name = 'multiverse'
            @client = Mongo::Client.new(
              ["#{$mongo.host}:#{$mongo.port}"],
              database: @database_name
            )
            @database = @client.database

            @collection_name = "tribbles-#{fake_guid(16)}"
            @collection = @database.collection(@collection_name)

            @tribbles = [{'name' => 'soterios johnson', 'count' => 1}, {'name' => 'wes mantooth', 'count' => 2}]
            @tribble = {:_id => 1, 'name' => 'soterios johnson'}

            NewRelic::Agent::Transaction.stubs(:recording_web_transaction?).returns(true)
            NewRelic::Agent.drop_buffered_data
          end

          def teardown
            NewRelic::Agent.drop_buffered_data
            @collection.drop
          end

          def test_noticed_error_at_segment_and_txn_when_violating_unique_constraints
            expected_error_class = /Mongo\:\:Error/
            txn = nil
            begin
              in_transaction do |db_txn|
                txn = db_txn
                @collection.insert_one(@tribble)
                @collection.insert_one(@tribble)
              end
            rescue StandardError => e
              # NOP -- allowing span and transaction to notice error
            end

            assert_segment_noticed_error txn, /insert/i, expected_error_class, /duplicate key error/i
            assert_transaction_noticed_error txn, expected_error_class
          end

          def test_noticed_error_only_at_segment_when_violating_unique_constraints
            expected_error_class = /Mongo\:\:Error/
            txn = nil
            in_transaction do |db_txn|
              begin
                txn = db_txn
                @collection.insert_one(_id: 1)
                @collection.insert_one(_id: 1)
              rescue Mongo::Error::OperationFailure => e
                # NOP -- allowing ONLY span to notice error
              end
            end

            assert_segment_noticed_error txn, /insert/i, expected_error_class, /duplicate key error/i
            refute_transaction_noticed_error txn, expected_error_class
          end

          def test_noticed_error_only_at_segment_when_command_fails
            skip if ENV['CI']
            expected_error_class = /Mongo\:\:Error/
            txn = nil
            in_transaction do |db_txn|
              begin
                txn = db_txn
                @database.collection('bogus').drop
              rescue Mongo::Error::OperationFailure => e
                # NOP -- allowing ONLY span to notice error
              end
            end

            assert_segment_noticed_error txn, /bogus\/drop/i, expected_error_class, /ns not found/i
            refute_transaction_noticed_error txn, expected_error_class
          end

          def test_records_metrics_for_insert_one
            in_transaction do
              @collection.insert_one(@tribbles.first)
            end

            metrics = build_test_metrics(:insert, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_insert_many
            in_transaction do
              @collection.insert_many(@tribbles)
            end

            metrics = build_test_metrics(:insert, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_delete_one
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.delete_one(@tribbles.first)
            end

            metrics = build_test_metrics(:delete, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_delete_many
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.delete_many(@tribbles.first)
            end

            metrics = build_test_metrics(:delete, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_replace_one
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.replace_one(@tribbles[0], @tribbles[1])
            end

            metrics = build_test_metrics(:update, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_update_one
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.update_one(@tribbles[0], '$set' => @tribbles[1])
            end

            metrics = build_test_metrics(:update, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_update_many
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.update_many(@tribbles[0], '$set' => @tribbles[1])
            end

            metrics = build_test_metrics(:update, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_find
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.find(@tribbles.first).to_a
            end

            metrics = build_test_metrics(:find, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_find_one_and_delete
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.find_one_and_delete(@tribbles.first)
            end

            metrics = build_test_metrics(:findAndModify, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_find_one_and_replace
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.find_one_and_replace(@tribbles[0], @tribbles[1])
            end

            metrics = build_test_metrics(:findAndModify, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_find_one_and_update
            @collection.insert_one(@tribbles.first)
            NewRelic::Agent.drop_buffered_data

            in_transaction do
              @collection.find_one_and_update(@tribbles[0], '$set' => @tribbles[1])
            end

            metrics = build_test_metrics(:findAndModify, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_distinct
            in_transaction do
              @collection.distinct('name')
            end

            metrics = build_test_metrics(:distinct, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_count
            in_transaction do
              @collection.count
            end

            metrics = build_test_metrics(:count, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_records_metrics_for_aggregate
            in_transaction do
              @collection.aggregate([
                {'$group' => {'_id' => 'name', 'max' => {'$max' => '$count'}}},
                {'$match' => {'max' => {'$gte' => 1}}}
              ]).to_a
            end

            metrics = build_test_metrics(:aggregate, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_aggregate_pipeline_obfuscated_by_default
            in_transaction do
              @collection.aggregate([
                {'$group' => {'_id' => 'name', 'max' => {'$max' => '$count'}}},
                {'$match' => {'max' => {'$gte' => 1}}}
              ]).to_a
            end

            sample = last_transaction_trace
            metric = "Datastore/statement/MongoDB/#{@collection_name}/aggregate"
            node = find_node_with_name(sample, metric)

            expected = [
              {'$group' => {'_id' => '?', 'max' => {'$max' => '?'}}},
              {'$match' => {'max' => {'$gte' => '?'}}}
            ]

            assert_equal expected, node[:statement]['pipeline']
          end

          def test_filter_obfuscated_by_default
            in_transaction do
              @collection.find('name' => 'Wes Mantooth', 'count' => {'$gte' => 1}).to_a
            end

            sample = last_transaction_trace
            metric = "Datastore/statement/MongoDB/#{@collection_name}/find"
            node = find_node_with_name(sample, metric)

            expected = {'name' => '?', 'count' => {'$gte' => '?'}}

            assert_equal expected, node[:statement]['filter']
          end

          def test_batched_queries
            25.times do |i|
              @collection.insert_one(:name => "test-#{i}", :active => true)
            end
            NewRelic::Agent.drop_buffered_data

            in_transaction('test_txn') do
              @collection.find(:active => true).batch_size(10).to_a
            end

            expected = {
              'test_txn' => {:call_count => 1},
              'OtherTransactionTotalTime' => {:call_count => 1},
              'OtherTransactionTotalTime/test_txn' => {:call_count => 1},
              ["Datastore/statement/MongoDB/#{@collection_name}/find", 'test_txn'] => {:call_count => 1},
              "Datastore/statement/MongoDB/#{@collection_name}/find" => {:call_count => 1},
              ["Datastore/statement/MongoDB/#{@collection_name}/getMore", 'test_txn'] => {:call_count => 2},
              "Datastore/statement/MongoDB/#{@collection_name}/getMore" => {:call_count => 2},
              'Datastore/operation/MongoDB/find' => {:call_count => 1},
              'Datastore/operation/MongoDB/getMore' => {:call_count => 2},
              "Datastore/instance/MongoDB/#{mongo_host}/27017" => {:call_count => 3},
              'Datastore/MongoDB/allWeb' => {:call_count => 3},
              'Datastore/MongoDB/all' => {:call_count => 3},
              'Datastore/allWeb' => {:call_count => 3},
              'Datastore/all' => {:call_count => 3},
              'Supportability/API/drop_buffered_data' => {:call_count => 1},
              'DurationByCaller/Unknown/Unknown/Unknown/Unknown/all' => {:call_count => 1},
              'DurationByCaller/Unknown/Unknown/Unknown/Unknown/allWeb' => {:call_count => 1}
            }

            assert_metrics_recorded_exclusive(expected,
              :ignore_filter => /^(Logging)/)
          end

          def test_batched_queries_have_node_per_query
            25.times do |i|
              @collection.insert_one(:name => "test-#{i}", :active => true)
            end
            NewRelic::Agent.drop_buffered_data
            in_transaction('webby') do
              @collection.find(:active => true).batch_size(10).to_a
            end

            expected = [
              "Datastore/statement/MongoDB/#{@collection_name}/find",
              "Datastore/statement/MongoDB/#{@collection_name}/getMore",
              "Datastore/statement/MongoDB/#{@collection_name}/getMore"
            ]

            trace = last_transaction_trace
            actual = []
            trace.each_node do |n|
              actual << n.metric_name if n.metric_name.start_with?('Datastore/statement/MongoDB')
            end

            assert_equal expected, actual
          end

          def test_trace_nodes_have_instance_attributes
            @collection.insert_one(:name => 'test', :active => true)
            NewRelic::Agent.drop_buffered_data
            in_transaction('webby') do
              @collection.find(:active => true).to_a
            end

            trace = last_transaction_trace
            node = find_node_with_name_matching(trace, /^Datastore\//)

            assert_equal mongo_host, node[:host]
            assert_equal '27017', node[:port_path_or_id]
            assert_equal @database_name, node[:database_name]
          end

          def test_drop_collection
            in_transaction do
              @collection.drop
            end

            metrics = build_test_metrics(:drop, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_web_scoped_metrics
            in_web_transaction('webby') do
              @collection.insert_one(@tribbles.first)
            end

            metric = statement_metric(:insert)

            assert_metrics_recorded([[metric, 'webby']])
          end

          def test_background_scoped_metrics
            in_background_transaction('backed-up') do
              @collection.insert_one(@tribbles.first)
            end

            metric = statement_metric(:insert)

            assert_metrics_recorded([[metric, 'backed-up']])
          end

          def test_notices_nosql
            node = nil

            in_transaction do
              @collection.insert_one(@tribbles.first)
            end

            node = find_last_transaction_node

            expected = {
              :database => @database_name,
              :collection => @collection_name,
              'insert' => @collection_name,
              :operation => :insert,
              'ordered' => true
            }

            result = node.params[:statement]

            # Depending on the Mongo DB version, we may get back
            # strings instead of symbols.  Let's adjust our
            # expectations accordingly.
            #
            if result[:operation].is_a?(String)
              expected[:operation] = expected[:operation].to_s
            end

            # The writeConcern is added by some, but not all versions
            # of the mongo driver, we don't care if it's present or
            # not, just that the statement is noticed
            #
            result.delete('writeConcern')

            if expected.is_a?(String)
              assert_equal expected, result
            else
              expected.each do |key, value|
                assert_equal value, result[key]
              end
            end
          end

          def test_noticed_nosql_includes_operation
            node = nil

            in_transaction do
              @collection.insert_one(@tribbles.first)
            end

            node = find_last_transaction_node
            query = node.params[:statement]

            assert_mongo_operation :insert, query
          end

          def test_noticed_nosql_includes_update_one_operation
            node = nil

            in_transaction do
              @collection.update_one(@tribbles[0], @tribbles[1])
            end

            node = find_last_transaction_node
            query = node.params[:statement]

            assert_mongo_operation :update, query
          end

          def test_noticed_nosql_includes_find_operation
            node = nil

            in_transaction do
              @collection.insert_one(@tribbles.first)
              @collection.find(@tribbles.first).to_a
            end

            node = find_last_transaction_node
            query = node.params[:statement]

            assert_mongo_operation 'find', query
          end

          def test_noticed_nosql_does_not_contain_documents
            node = nil

            in_transaction do
              @collection.insert_one(@tribbles.first)
            end

            node = find_last_transaction_node
            statement = node.params[:statement]

            refute statement.key?(:documents), "Noticed NoSQL should not include documents: #{statement}"
          end

          def test_noticed_nosql_does_not_contain_selector_values
            @collection.insert_one({'password' => '$ecret'})
            node = nil

            in_transaction do
              @collection.find({'password' => '$ecret'}).to_a
            end

            node = find_last_transaction_node
            statement = node.params[:statement]

            refute_includes statement.inspect, '$ecret'
            assert_equal '?', statement['filter']['password']
          end

          def test_web_requests_record_all_web_metric
            NewRelic::Agent::Transaction.stubs(:recording_web_transaction?).returns(true)
            in_web_transaction do
              @collection.insert_one(@tribbles.first)
            end

            metrics = build_test_metrics(:insert, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_web_requests_do_not_record_all_other_metric
            NewRelic::Agent::Transaction.stubs(:recording_web_transaction?).returns(true)
            in_background_transaction do
              @collection.insert_one(@tribbles.first)
            end

            assert_metrics_not_recorded(['Datastore/allOther'])
          end

          def test_other_requests_record_all_other_metric
            NewRelic::Agent::Transaction.stubs(:recording_web_transaction?).returns(false)
            in_transaction do
              @collection.insert_one(@tribbles.first)
            end

            metrics = build_test_metrics(:insert, true)
            expected = metrics_with_attributes(metrics)

            assert_metrics_recorded(expected)
          end

          def test_other_requests_do_not_record_all_web_metric
            NewRelic::Agent::Transaction.stubs(:recording_web_transaction?).returns(false)
            @collection.insert_one(@tribbles.first)

            assert_metrics_not_recorded(['Datastore/allWeb'])
          end

          def statement_metric(action)
            metrics = build_test_metrics(action, true)
            metrics.find { |m| m.start_with?('Datastore/statement') }
          end

          def assert_mongo_operation(expected_value, query)
            assert_equal expected_value.to_s, query[:operation].to_s
          end
        end
      end
    end
  end
end
