Esempio n. 1
0
    def setUp(self):
        super(GRRBaseTest, self).setUp()

        self.temp_dir = temp.TempDirPath()
        config.CONFIG.SetWriteBack(
            os.path.join(self.temp_dir, "writeback.yaml"))

        logging.info("Starting test: %s.%s", self.__class__.__name__,
                     self._testMethodName)
        self.last_start_time = time.time()

        data_store.DB.ClearTestDB()
        # Each datastore is wrapped with DatabaseValidationWrapper, so we have
        # to access the delegate directly (assuming it's an InMemoryDB
        # implementation).
        data_store.REL_DB.delegate.ClearTestDB()

        aff4.FACTORY.Flush()

        # Create a Foreman and Filestores, they are used in many tests.
        aff4_grr.GRRAFF4Init().Run()
        filestore.FileStoreInit().Run()
        hunts_results.ResultQueueInitHook().Run()
        email_alerts.EmailAlerterInit().RunOnce()
        audit.AuditEventListener._created_logs.clear()

        # Stub out the email function
        self.emails_sent = []

        def SendEmailStub(to_user, from_user, subject, message,
                          **unused_kwargs):
            self.emails_sent.append((to_user, from_user, subject, message))

        self.mail_stubber = utils.MultiStubber(
            (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub),
            (email.utils, "make_msgid", lambda: "<message id stub>"))
        self.mail_stubber.Start()

        # We don't want to send actual email in our tests
        self.smtp_patcher = mock.patch("smtplib.SMTP")
        self.mock_smtp = self.smtp_patcher.start()

        def DisabledSet(*unused_args, **unused_kw):
            raise NotImplementedError(
                "Usage of Set() is disabled, please use a configoverrider in tests."
            )

        self.config_set_disable = utils.Stubber(config.CONFIG, "Set",
                                                DisabledSet)
        self.config_set_disable.Start()

        if self.use_relational_reads:
            self.relational_read_stubber = utils.Stubber(
                data_store, "RelationalDBReadEnabled", lambda: True)
            self.relational_read_stubber.Start()

        self._SetupFakeStatsContext()
Esempio n. 2
0
  def setUp(self):
    super(GRRBaseTest, self).setUp()

    self.temp_dir = temp.TempDirPath()
    config.CONFIG.SetWriteBack(os.path.join(self.temp_dir, "writeback.yaml"))
    self.addCleanup(lambda: shutil.rmtree(self.temp_dir, ignore_errors=True))

    data_store.DB.ClearTestDB()
    # Each datastore is wrapped with DatabaseValidationWrapper, so we have
    # to access the delegate directly (assuming it's an InMemoryDB
    # implementation).
    data_store.REL_DB.delegate.ClearTestDB()

    aff4.FACTORY.Flush()

    # Create a Foreman and Filestores, they are used in many tests.
    aff4_grr.GRRAFF4Init().Run()
    filestore.FileStoreInit().Run()
    hunts_results.ResultQueueInitHook().Run()
    email_alerts.EmailAlerterInit().RunOnce()
    audit.AuditEventListener._created_logs.clear()

    # Stub out the email function
    self.emails_sent = []

    def SendEmailStub(to_user, from_user, subject, message, **unused_kwargs):
      self.emails_sent.append((to_user, from_user, subject, message))

    self.mail_stubber = utils.MultiStubber(
        (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub),
        (email.utils, "make_msgid", lambda: "<message id stub>"))
    self.mail_stubber.Start()
    self.addCleanup(self.mail_stubber.Stop)

    # We don't want to send actual email in our tests
    self.smtp_patcher = mock.patch("smtplib.SMTP")
    self.mock_smtp = self.smtp_patcher.start()
    self.addCleanup(self.smtp_patcher.stop)

    def DisabledSet(*unused_args, **unused_kw):
      raise NotImplementedError(
          "Usage of Set() is disabled, please use a configoverrider in tests.")

    self.config_set_disable = utils.Stubber(config.CONFIG, "Set", DisabledSet)
    self.config_set_disable.Start()
    self.addCleanup(self.config_set_disable.Stop)

    if self.use_relational_reads:
      self.relational_read_stubber = utils.Stubber(
          data_store, "RelationalDBEnabled", lambda: True)
      self.relational_read_stubber.Start()
      self.addCleanup(self.relational_read_stubber.Stop)

    self._SetupFakeStatsContext()

    # Turn off WithLimitedCallFrequency-based caching in tests. Tests that need
    # to test caching behavior explicitly, should turn it on explicitly.
    with_limited_call_frequency_stubber = utils.Stubber(
        cache, "WITH_LIMITED_CALL_FREQUENCY_PASS_THROUGH", True)
    with_limited_call_frequency_stubber.Start()
    self.addCleanup(with_limited_call_frequency_stubber.Stop)