<commit_before>require_relative "helper"
require_relative "db"

require "logger"
require "pp"
require "fileutils"

CDB = ConceptQL::Database.new(DB, :data_model=>ENV["CONCEPTQL_DATA_MODEL"].to_sym)
DB.extension :error_sql

class Minitest::Spec
  def annotate(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).annotate}
  end

  def scope_annotate(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).scope_annotate}
  end

  def domains(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).domains}
  end

  def results(test_name, statement=nil)
    load_check(test_name, statement) do |stmt, remove_window|
      ds = dataset(query(stmt)).from_self

      order_columns = [:person_id, :criterion_table, :criterion_domain, :start_date, :criterion_id]
      order_columns << :uuid if ds.columns.include?(:uuid)
      ds = ds.order(*order_columns)

      results = ds.all
      results.each do |h|
        h.transform_values! { |v| v.is_a?(Time) || v.is_a?(DateTime) ? v.to_date : v }
        h.delete(:window_id) if remove_window
        h
      end
    end
  end

  def query(statement)
    CDB.query(statement)
  end

  def dataset(statement)
    statement = query(statement) unless statement.is_a?(ConceptQL::Query)
    statement.query
  end

  def criteria_ids(test_name, statement=nil)
    load_check(test_name, statement){|stmt| hash_groups(stmt, :criterion_domain, :criterion_id)}
  end

  def code_check(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).code_list.map(&:to_s)}
  end

  # If no statement is passed, this function loads the statement from the specified test
  # file. If a statement is passed, it is written to the file.
  def load_statement(test_name, statement)
    path = "test/statements/#{test_name}"
    if statement
      FileUtils.mkdir_p(File.dirname(path))
      File.write(path, JSON.pretty_generate(statement))
      statement
    else
      JSON.parse(File.read(path))
    end
  end

  def check_output(test_name, results)
    path = "test/results/#{ENV["CONCEPTQL_DATA_MODEL"]}/#{test_name}"

    if ENV["OVERWRITE_CONCEPTQL_TEST_RESULTS"]
      save_results(path, results)
    end

    expected = begin
      File.read(path)
    rescue Errno::ENOENT
      save_results(path, { fail: true })
      '{ "fail": true }'
    end

    JSON.parse(results.to_json).must_equal(JSON.parse(expected))
    results
  end

  def save_results(path, results)
    FileUtils.mkdir_p(File.dirname(path))
    File.write(path, JSON.pretty_generate(results))
  end

  def numeric_values(test_name, statement=nil)
    load_check(test_name, statement){|stmt| hash_groups(stmt, :criterion_domain, :value_as_number)}
  end

  def criteria_counts(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).query.from_self.group_and_count(:criterion_domain).order(:criterion_domain).to_hash(:criterion_domain, :count)}
  end

  def optimized_criteria_counts(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).optimized.query.from_self.group_and_count(:criterion_domain).order(:criterion_domain).to_hash(:criterion_domain, :count)}
  end

  def hash_groups(statement, key, value)
    dataset(statement).from_self.distinct.order(key, *value).to_hash_groups(key, value)
  rescue
    puts $!.sql if $!.respond_to?(:sql)
    raise
  end

  def clock_time
    Process.clock_gettime(Process::CLOCK_MONOTONIC)
  end

  PERFORMANCE_TEST_TIMES = ENV["CONCEPTQL_PERFORMANCE_TEST_TIMES"].to_i
  def load_check(test_name, statement)
    statement = load_statement(test_name, statement)

    # Check without scope windows
    results = yield(statement, false)
    check_output(test_name, results)

    # Check with scope windows, unless the test is already a scope window test
    unless statement.first == 'window'
      sw_statement = ["window", statement, {'window_table' => [ 'date_range', { 'start' => '2007-01-01', 'end' => '2011-12-31' } ] } ]
      results = yield(sw_statement, true)
      check_output(test_name, results)
    end

    if PERFORMANCE_TEST_TIMES > 0
      times = PERFORMANCE_TEST_TIMES.times.map do 
        before = clock_time
        yield(statement, false)
        clock_time - before
      end
      avg_time = times.sum/times.length

      path = "test/performance/#{ENV["CONCEPTQL_DATA_MODEL"]}/#{test_name.sub(/\.json\z/, '.csv')}"
      sha1 = `git rev-parse HEAD`.chomp
      adapter = "#{DB.adapter_scheme}/#{DB.database_type}"

      FileUtils.mkdir_p(File.dirname(path))
      File.open(path, 'ab') do |file|
        file.puts([sha1, adapter, DB.opts[:database], Time.now, avg_time].join(','))
      end
    end<commit_msg>DbHelper: indicate when a test used windows or not<commit_after>require_relative "helper"
require_relative "db"

require "logger"
require "pp"
require "fileutils"

CDB = ConceptQL::Database.new(DB, :data_model=>ENV["CONCEPTQL_DATA_MODEL"].to_sym)
DB.extension :error_sql

class Minitest::Spec
  def annotate(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).annotate}
  end

  def scope_annotate(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).scope_annotate}
  end

  def domains(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).domains}
  end

  def results(test_name, statement=nil)
    load_check(test_name, statement) do |stmt, remove_window|
      ds = dataset(query(stmt)).from_self

      order_columns = [:person_id, :criterion_table, :criterion_domain, :start_date, :criterion_id]
      order_columns << :uuid if ds.columns.include?(:uuid)
      ds = ds.order(*order_columns)

      results = ds.all
      results.each do |h|
        h.transform_values! { |v| v.is_a?(Time) || v.is_a?(DateTime) ? v.to_date : v }
        h.delete(:window_id) if remove_window
        h
      end
    end
  end

  def query(statement)
    CDB.query(statement)
  end

  def dataset(statement)
    statement = query(statement) unless statement.is_a?(ConceptQL::Query)
    statement.query
  end

  def criteria_ids(test_name, statement=nil)
    load_check(test_name, statement){|stmt| hash_groups(stmt, :criterion_domain, :criterion_id)}
  end

  def code_check(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).code_list.map(&:to_s)}
  end

  # If no statement is passed, this function loads the statement from the specified test
  # file. If a statement is passed, it is written to the file.
  def load_statement(test_name, statement)
    path = "test/statements/#{test_name}"
    if statement
      FileUtils.mkdir_p(File.dirname(path))
      File.write(path, JSON.pretty_generate(statement))
      statement
    else
      JSON.parse(File.read(path))
    end
  end

  def check_output(test_name, results, has_windows = false)
    path = "test/results/#{ENV["CONCEPTQL_DATA_MODEL"]}/#{test_name}"

    if ENV["OVERWRITE_CONCEPTQL_TEST_RESULTS"]
      save_results(path, results)
    end

    expected = begin
      File.read(path)
    rescue Errno::ENOENT
      save_results(path, { fail: true })
      '{ "fail": true }'
    end

    JSON.parse(results.to_json).must_equal(JSON.parse(expected), "#{test_name}#{has_windows ? " (with windows)" : ""}")
    results
  end

  def save_results(path, results)
    FileUtils.mkdir_p(File.dirname(path))
    File.write(path, JSON.pretty_generate(results))
  end

  def numeric_values(test_name, statement=nil)
    load_check(test_name, statement){|stmt| hash_groups(stmt, :criterion_domain, :value_as_number)}
  end

  def criteria_counts(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).query.from_self.group_and_count(:criterion_domain).order(:criterion_domain).to_hash(:criterion_domain, :count)}
  end

  def optimized_criteria_counts(test_name, statement=nil)
    load_check(test_name, statement){|stmt| query(stmt).optimized.query.from_self.group_and_count(:criterion_domain).order(:criterion_domain).to_hash(:criterion_domain, :count)}
  end

  def hash_groups(statement, key, value)
    dataset(statement).from_self.distinct.order(key, *value).to_hash_groups(key, value)
  rescue
    puts $!.sql if $!.respond_to?(:sql)
    raise
  end

  def clock_time
    Process.clock_gettime(Process::CLOCK_MONOTONIC)
  end

  PERFORMANCE_TEST_TIMES = ENV["CONCEPTQL_PERFORMANCE_TEST_TIMES"].to_i
  def load_check(test_name, statement)
    statement = load_statement(test_name, statement)

    # Check without scope windows
    results = yield(statement, false)
    check_output(test_name, results, false)

    # Check with scope windows, unless the test is already a scope window test
    unless statement.first == 'window'
      sw_statement = ["window", statement, {'window_table' => [ 'date_range', { 'start' => '2007-01-01', 'end' => '2011-12-31' } ] } ]
      results = yield(sw_statement, true)
      check_output(test_name, results, true)
    end

    if PERFORMANCE_TEST_TIMES > 0
      times = PERFORMANCE_TEST_TIMES.times.map do 
        before = clock_time
        yield(statement, false)
        clock_time - before
      end
      avg_time = times.sum/times.length

      path = "test/performance/#{ENV["CONCEPTQL_DATA_MODEL"]}/#{test_name.sub(/\.json\z/, '.csv')}"
      sha1 = `git rev-parse HEAD`.chomp
      adapter = "#{DB.adapter_scheme}/#{DB.database_type}"

      FileUtils.mkdir_p(File.dirname(path))
      File.open(path, 'ab') do |file|
        file.puts([sha1, adapter, DB.opts[:database], Time.now, avg_time].join(','))
      end
    end
{'fraction_non_alphanumeric': 0.1224191866527633, 'fraction_numerical': 0.003962460896767466, 'mean_word_length': 4.109749600426212, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2594810', 'n_tokens_mistral': 3276, 'n_tokens_neox': 3193, 'n_words': 624}