added whitelist check and added registrant only flag to rake task

.gitignore

@@ -25,3 +25,4 @@ node_modules/
 /lib/tasks/data
 ettevotja_rekvisiidid__lihtandmed.csv.zip
 ./ettevotja_rekvisiidid__lihtandmed.csv.zip
+ai

Dockerfile

@@ -6,10 +6,28 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     postgresql-client \
   && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
+RUN apt-get remove -y google-chrome-stable
+RUN apt-get purge -y google-chrome-stable
+RUN apt-get autoremove -y && apt-get clean
+
+ENV CHROME_VERSION="128.0.6613.137"
+
+RUN wget -q "https://storage.googleapis.com/chrome-for-testing-public/${CHROME_VERSION}/linux64/chrome-linux64.zip" \
+    && unzip chrome-linux64.zip -d /opt/ \
+    && rm chrome-linux64.zip
+
+RUN wget -q "https://storage.googleapis.com/chrome-for-testing-public/${CHROME_VERSION}/linux64/chromedriver-linux64.zip" \
+    && unzip chromedriver-linux64.zip -d /opt/ \
+    && mv /opt/chromedriver-linux64/chromedriver /usr/local/bin/ \
+    && rm -rf chromedriver-linux64.zip /opt/chromedriver-linux64
+
+
 RUN mkdir -p /opt/webapps/app/tmp/pids
 WORKDIR /opt/webapps/app
 COPY Gemfile Gemfile.lock ./
 # ADD vendor/gems/omniauth-tara ./vendor/gems/omniauth-tara
 RUN gem install bundler && bundle install --jobs 20 --retry 5
 
+ENV PATH="/opt/chrome-linux64:${PATH}"
+
 EXPOSE 3000

Gemfile

@@ -107,3 +107,8 @@ gem 'pg_query', '>= 0.9.0'
 
 # token
 gem 'jwt'
+gem 'net-ftp'
+gem 'net-smtp', require: false
+gem 'net-imap', require: false
+gem 'net-pop', require: false
+gem 'matrix'

Gemfile.lock

@@ -311,6 +311,7 @@ GEM
       net-pop
       net-smtp
     marcel (1.0.2)
+    matrix (0.3.1)
     method_source (1.0.0)
     mime-types (3.3.1)
       mime-types-data (~> 3.2015)
@@ -332,6 +333,9 @@ GEM
       railties (>= 3.0)
     msgpack (1.7.2)
     mutex_m (0.2.0)
+    net-ftp (0.1.2)
+      net-protocol
+      time
     net-imap (0.4.10)
       date
       net-protocol
@@ -506,6 +510,8 @@ GEM
     temple (0.8.2)
     thor (1.2.2)
     tilt (2.0.11)
+    time (0.1.0)
+      date
     timeout (0.4.1)
     truemail (3.0.3)
       simpleidn (~> 0.2.1)
@@ -583,11 +589,16 @@ DEPENDENCIES
   jwt
   kaminari
   lhv!
+  matrix
   mime-types-data
   mimemagic (= 0.4.3)
   minitest (~> 5.17)
   minitest-stub_any_instance
   money-rails
+  net-ftp
+  net-imap
+  net-pop
+  net-smtp
   newrelic-infinite_tracing
   newrelic_rpm
   nokogiri (~> 1.16.0)

lib/tasks/company_status.rake

@@ -24,6 +24,7 @@ namespace :company_status do
     download_path = options[:download_path]
     soft_delete_enable = options[:soft_delete_enable]
     downloaded_filename = File.basename(URI(download_path).path)
+    registrants_only = options[:registrants_only]
 
     puts "*** Run 1 step. Downloading fresh open data file. ***"
     remove_old_file(DESTINATION + downloaded_filename)
@@ -34,7 +35,14 @@ namespace :company_status do
     company_data = collect_company_data(open_data_file_path)
 
     puts "*** Run 3 step. I process companies, update their information, and sort them into different files based on whether the companies are missing or removed from the business registry ***"
-    Registrant.where(ident_type: 'org', ident_country_code: 'EE').find_each do |contact|
+
+    whitelisted_companies = JSON.parse(ENV['whitelist_companies']) # ["12345678", "87654321"]    
+    contacts_query = Contact.where(ident_type: 'org', ident_country_code: 'EE')
+    contacts_query = contacts_query.where(type: 'Registrant') if registrants_only
+
+    contacts_query.find_each do |contact|
+      next if whitelisted_companies.include?(contact.ident)
+
       if company_data.key?(contact.ident)
         update_company_status(contact: contact, status: company_data[contact.ident][COMPANY_STATUS])
         puts "Company: #{contact.name} with ident: #{contact.ident} and ID: #{contact.id} has status: #{company_data[contact.ident][COMPANY_STATUS]}"
@@ -66,6 +74,7 @@ namespace :company_status do
       deleted_companies_output_path: deleted_companies_from_business_registry_path,
       download_path: url,
       soft_delete_enable: false,
+      registrants_only: false,
     }
 
     banner = 'Usage: rake companies:check_all -- [options]'
@@ -81,6 +90,7 @@ namespace :company_status do
       deleted_companies_output_path: ['-s [DELETED_COMPANIES_OUTPUT_PATH]', '--deleted_companies_output_path [DELETED_COMPANIES_OUTPUT_PATH]', String],
       download_path: ['-d [DOWNLOAD_PATH]', '--download_path [DOWNLOAD_PATH]', String],
       soft_delete_enable: ['-e [SOFT_DELETE_ENABLE]', '--soft_delete_enable [SOFT_DELETE_ENABLE]', FalseClass],
+      registrants_only: ['-r [REGISTRANTS_ONLY]', '--registrants_only [REGISTRANTS_ONLY]', TrueClass],
     }
   end
 
@@ -133,7 +143,7 @@ namespace :company_status do
   end
 
   def put_company_to_missing_file(contact:, path:)
-    write_to_csv_file(csv_file_path: path, headers: ["ID", "Ident", "Name"], attrs: [contact.id, contact.ident, contact.name])
+    write_to_csv_file(csv_file_path: path, headers: ["ID", "Ident", "Name", "Contact Type"], attrs: [contact.id, contact.ident, contact.name, determine_contact_type(contact)])
   end
 
   def sort_companies_to_files(contact:, missing_companies_in_business_registry_path:, deleted_companies_from_business_registry_path:, soft_delete_enable:)
@@ -152,8 +162,8 @@ namespace :company_status do
 
       if status == DELETED_FROM_REGISTRY_STATUS
         csv_file_path = deleted_companies_from_business_registry_path
-        headers = ["ID", "Ident", "Name", "Status", "Kandeliik Type", "Kandeliik Tekstina", "kande_kpv"]
-        attrs = [contact.id, contact.ident, contact.name, status, kandeliik_type, kandeliik_tekstina, kande_kpv]
+        headers = ["ID", "Ident", "Name", "Status", "Kandeliik Type", "Kandeliik Tekstina", "kande_kpv", "Contact Type"]
+        attrs = [contact.id, contact.ident, contact.name, status, kandeliik_type, kandeliik_tekstina, kande_kpv, determine_contact_type(contact)]
         write_to_csv_file(csv_file_path: csv_file_path, headers: headers, attrs: attrs)
 
         puts "Company: #{contact.name} with ident: #{contact.ident} and ID: #{contact.id} has status #{status}, company id: #{contact.id}"
@@ -162,6 +172,14 @@ namespace :company_status do
     end
   end
 
+  def determine_contact_type(contact)
+    roles = []
+    roles << 'Registrant' if contact.registrant_domains.any?
+    roles += contact.domain_contacts.pluck(:type).uniq if contact.domain_contacts.any?
+    roles << 'Unknown' if roles.empty?
+    roles.join(', ')
+  end
+
   def soft_delete_company(contact)
     contact.domains.reject { |domain| domain.force_delete_scheduled? }.each do |domain|
       domain.schedule_force_delete(type: :soft)

test/tasks/company_status_test.rb

@@ -3,144 +3,101 @@
 require 'minitest/mock'
 
 class CompanyStatusRakeTaskTest < ActiveSupport::TestCase
-  class MockRegistrant
-    attr_accessor :id, :ident, :ident_type, :ident_country_code, :name, :registrant_publishable
+  def setup
+    Rails.application.load_tasks
+    @task = Rake::Task['company_status:check_all']
+  end
 
-    def initialize(attributes = {})
-      attributes.each do |key, value|
-        send("#{key}=", value)
-      end
-    end
+  def teardown
+    Rake::Task.clear
+  end
 
-    def update(attributes)
-      attributes.each do |key, value|
-        send("#{key}=", value)
+  test "initialize_rake_task sets default options correctly" do
+    # Мокаем ENV для whitelisted_companies
+    ENV.stub :[], '["12345678", "87654321"]', ['whitelist_companies'] do
+      # Мокаем RakeOptionParserBoilerplate для избежания обработки аргументов командной строки
+      RakeOptionParserBoilerplate.stub :process_args, nil do
+        options = nil
+
+        # Перехватываем вызов метода collect_company_data, чтобы получить options
+        collect_company_data_method = lambda do |open_data_file_path|
+          options = open_data_file_path
+          {}  # Возвращаем пустой хэш, так как нам не нужны реальные данные для этого теста
+        end
+
+        Rake::Task['company_status:check_all'].enhance do
+          def collect_company_data(open_data_file_path)
+            collect_company_data_method.call(open_data_file_path)
+          end
+        end
+
+        # Выполняем задачу
+        silence_stream(STDOUT) { @task.execute }
+
+        # Проверяем, что опции установлены корректно
+        assert_equal 'tmp/ettevotja_rekvisiidid__lihtandmed.csv', options[:open_data_file_path]
+        assert_equal 'tmp/missing_companies_in_business_registry.csv', options[:missing_companies_output_path]
+        assert_equal 'tmp/deleted_companies_from_business_registry.csv', options[:deleted_companies_output_path]
+        assert_equal 'https://avaandmed.ariregister.rik.ee/sites/default/files/avaandmed/ettevotja_rekvisiidid__lihtandmed.csv.zip', options[:download_path]
+        assert_equal false, options[:soft_delete_enable]
+        assert_equal false, options[:registrants_only]
       end
-      true
     end
+  end
 
-    def publishable?
-      registrant_publishable
-    end
+  test "initialize_rake_task processes command line arguments" do
+    # Test that command line arguments are processed correctly
+  end
 
-    def epp_code_map
-      {}
-    end
+  test "download_open_data_file downloads file successfully" do
+    # Test the file download process
   end
 
-  def setup
-    Rails.application.load_tasks
-    @task = Rake::Task['company_status:check_all']
+  test "unzip_file extracts contents correctly" do
+    # Test the unzipping process
   end
 
-  def teardown
-    Rake::Task.clear
+  test "collect_company_data parses CSV file correctly" do
+    # Test the CSV parsing and data collection
+  end
+
+  test "update_company_status updates contact status correctly" do
+    # Test the contact status update process
+  end
+
+  test "sort_companies_to_files handles missing companies correctly" do
+    # Test the process for companies missing from the registry
+  end
+
+  test "sort_companies_to_files handles deleted companies correctly" do
+    # Test the process for companies deleted from the registry
+  end
+
+  test "determine_contact_type returns correct roles" do
+    # Test the contact type determination logic
+  end
+
+  test "soft_delete_company initiates force delete for domains" do
+    # Test the soft delete process for companies
+  end
+
+  test "write_to_csv_file creates and appends to CSV files correctly" do
+    # Test the CSV file writing process
+  end
+
+  test "check_all task processes contacts correctly with registrants_only flag" do
+    # Test the main task execution with the registrants_only flag
+  end
+
+  test "check_all task handles whitelisted companies correctly" do
+    # Test the whitelisted companies logic
+  end
+
+  test "check_all task updates existing companies correctly" do
+    # Test the process for updating existing companies
   end
 
-  test "check_all task runs successfully" do
-    # # Подготовка моков и заглушек
-    # mock_download = Minitest::Mock.new
-    # mock_download.expect :call, nil, [String, String]
-
-    # mock_unzip = Minitest::Mock.new
-    # mock_unzip.expect :call, nil, [String, String]
-
-    # mock_collect_data = Minitest::Mock.new
-    # mock_collect_data.expect :call, { "123456789" => { 'ettevotja_staatus' => "A" } }, [String]
-
-    # mock_update_status = Minitest::Mock.new
-    # mock_update_status.expect :call, true, [{ contact: MockRegistrant, status: String }]
-
-    # # Создание тестовых данных
-    # registrant = MockRegistrant.new(
-    #   id: 1, 
-    #   ident: "123456789", 
-    #   ident_type: 'org', 
-    #   ident_country_code: 'EE', 
-    #   name: 'Test Company',
-    #   registrant_publishable: true
-    # )
-
-    # # Подмена методов заглушками
-    # Rake::Task['company_status:check_all'].instance_eval do
-    #   define_method(:download_open_data_file) { |url, filename| mock_download.call(url, filename) }
-    #   define_method(:unzip_file) { |filename, destination| mock_unzip.call(filename, destination) }
-    #   define_method(:collect_company_data) { |path| mock_collect_data.call(path) }
-    #   define_method(:update_company_status) { |args| mock_update_status.call(args) }
-    # end
-
-    # # Имитация выборки Registrant.where
-    # Registrant.stub :where, [registrant] do
-    #   # Выполнение задачи
-    #   assert_nothing_raised { @task.execute }
-    # end
-
-    # # Проверка, что все моки были вызваны
-    # assert_mock mock_download
-    # assert_mock mock_unzip
-    # assert_mock mock_collect_data
-    # assert_mock mock_update_status
-  end
-
-  test "sort_companies_to_files handles missing company" do
-    # contact = MockRegistrant.new(id: 2, ident: "987654321", ident_type: 'org', ident_country_code: 'EE', name: 'Missing Company')
-
-    # mock_return_company_details = Minitest::Mock.new
-    # mock_return_company_details.expect :call, []
-
-    # contact.define_singleton_method(:return_company_details) { mock_return_company_details.call }
-
-    # missing_file = Tempfile.new(['missing', '.csv'])
-    # deleted_file = Tempfile.new(['deleted', '.csv'])
-
-    # Rake::Task['company_status:check_all'].send(:sort_companies_to_files,
-    #   contact: contact,
-    #   missing_companies_in_business_registry_path: missing_file.path,
-    #   deleted_companies_from_business_registry_path: deleted_file.path,
-    #   soft_delete_enable: false
-    # )
-
-    # assert_mock mock_return_company_details
-
-    # missing_content = File.read(missing_file.path)
-    # assert_includes missing_content, contact.ident
-    # assert_includes missing_content, contact.name
-
-    # missing_file.unlink
-    # deleted_file.unlink
-  end
-
-  test "sort_companies_to_files handles deleted company" do
-    # contact = MockRegistrant.new(id: 3, ident: "111222333", ident_type: 'org', ident_country_code: 'EE', name: 'Deleted Company')
-
-    # mock_return_company_details = Minitest::Mock.new
-    # mock_return_company_details.expect :call, [
-    #   OpenStruct.new(
-    #     status: 'K',
-    #     kandeliik: [['', OpenStruct.new(kandeliik: 'type', kandeliik_tekstina: 'text', kande_kpv: '2023-01-01')]]
-    #   )
-    # ]
-
-    # contact.define_singleton_method(:return_company_details) { mock_return_company_details.call }
-
-    # missing_file = Tempfile.new(['missing', '.csv'])
-    # deleted_file = Tempfile.new(['deleted', '.csv'])
-
-    # Rake::Task['company_status:check_all'].send(:sort_companies_to_files,
-    #   contact: contact,
-    #   missing_companies_in_business_registry_path: missing_file.path,
-    #   deleted_companies_from_business_registry_path: deleted_file.path,
-    #   soft_delete_enable: false
-    # )
-
-    # assert_mock mock_return_company_details
-
-    # deleted_content = File.read(deleted_file.path)
-    # assert_includes deleted_content, contact.ident
-    # assert_includes deleted_content, contact.name
-    # assert_includes deleted_content, 'K'
-
-    # missing_file.unlink
-    # deleted_file.unlink
+  test "check_all task handles companies not found in registry data" do
+    # Test the process for companies not found in the registry data
   end
 end