<commit_before>"""Tests for client."""

import abc
import functools
from typing import Callable, Optional

from absl import logging
from vizier import pyvizier as vz
from vizier.client import client_abc

from absl.testing import parameterized

# Aliases are defined, so when you are developing a new client, you can
# swap it with your subclass. It makes your IDE understand which class
# you are using.
_StudyClient = client_abc.StudyInterface
_TrialClient = client_abc.TrialInterface


class VizierClientTestMixin(abc.ABC):

  @abc.abstractmethod
  def create_study(self, study_config: vz.ProblemStatement,
                   name: str) -> _StudyClient:
    """Create study given study config and study name."""
    pass


class MyMeta(type(VizierClientTestMixin), type(parameterized.TestCase)):
  pass


class TestCase(parameterized.TestCase, VizierClientTestMixin, metaclass=MyMeta):
  """Generic tests for cross-platform clients.

  This test provides basic coverage and it is not meant to be a comprehensive
  e2e test on all possible inputs to the client.

  Override `create_study` method.
  """

  def create_test_study(self, name: str) -> _StudyClient:
    """Creates a study."""
    logging.info('Creating study name=%s, testcasename=%s', name, self.id())
    problem = vz.ProblemStatement()
    problem.search_space.root.add_float_param('float', 0.0, 1.0)
    problem.metric_information.append(
        vz.MetricInformation(
            name='maximize_metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))
    study = self.create_study(problem, name)

    # TODO: Put this line back once we have __eq__ well-defined.
    # self.assertEqual(study.materialize_problem_statement(), problem)
    return study

  def test_create_or_load_study(self):
    study = self.create_test_study(self.id())
    study2 = study.from_resource_name(study.resource_name)
    self.assertEqual(study.resource_name, study2.resource_name)

    self.assertEqual(study.materialize_problem_statement(),
                     study2.materialize_problem_statement())

  def test_delete_study(self):
    study = self.create_test_study(self.id())
    resource_name = study.resource_name
    study.delete()

    with self.assertRaises(Exception):
      study.from_resource_name(resource_name)

  def _example_trials(self) -> list[vz.Trial]:
    trials = [
        vz.Trial().complete(vz.Measurement({'maximize_metric': 1.0})),
        vz.Trial().complete(vz.Measurement({'maximize_metric': 0.5})),
        vz.Trial(measurements=[vz.Measurement({'maximize_metric': 0.7})]),
        vz.Trial(is_requested=True)
    ]
    for idx, t in enumerate(trials):
      t.metadata['future_id'] = str(idx + 1)  # id to be assigned
    return trials

  def create_test_study_with_trials(self, name: str) -> _StudyClient:
    study = self.create_test_study(name)
    for t in self._example_trials():
      # TODO: Remove this.
      study._add_trial(t)  # pylint: disable=protected-access
    return study

  @parameterized.parameters(list(state for state in vz.StudyState))
  def test_set_state(self, target_state: vz.StudyState):
    study = self.create_test_study_with_trials(self.id())
    # TODO: Check that the study moved to the target state.
    # This test is currently a placeholder.
    try:
      study.set_state(target_state)
    except NotImplementedError:
      logging.exception('Set study state for %s is not implemented in %s',
                        target_state, type(self))

  def test_list_trials(self):
    study = self.create_test_study_with_trials(self.id())
    self.assertLen(list(study.trials()), 4)

  def test_trials_iter_and_get_are_equal(self):
    study = self.create_test_study_with_trials(self.id())
    all_trials = study.trials()
    self.assertEqual([t.id for t in all_trials],
                     [t.id for t in all_trials.get()])

  def test_optimal_trials_on_empty(self):
    study = self.create_test_study(self.id())
    self.assertEmpty(list(study.optimal_trials()))

  def test_optimal_trials_notempty(self):
    study = self.create_test_study_with_trials(self.id())
    optimal_trials = list(study.optimal_trials())
    self.assertLen(optimal_trials, 1)
    self.assertEqual(optimal_trials[0].materialize().status,
                     vz.TrialStatus.COMPLETED)

  def test_suggest_same_worker(self):
    # Given the same client id, should return the same trials.
    study = self.create_test_study(self.id())
    trials1 = study.suggest(count=5, client_id='worker1')
    self.assertLen(trials1, 5)
    trials2 = study.suggest(count=2, client_id='worker1')
    self.assertTrue({t.id for t in trials2}.issubset({t.id for t in trials1}))

  def test_suggest_different_workers(self):
    # Given different client ids, should generate new suggestions.
    study = self.create_test_study(self.id())
    worker1_trials = study.suggest(count=5, client_id='worker1')
    self.assertLen(worker1_trials, 5)
    worker1_trial_ids = {t.id for t in worker1_trials}
    worker2_trials = study.suggest(count=2, client_id='worker2')
    self.assertLen(worker2_trials, 2)
    worker2_trial_ids = {t.id for t in worker2_trials}
    self.assertEmpty(
        worker1_trial_ids.intersection(worker2_trial_ids),
        msg=(f'worker1_trial_ids={worker1_trial_ids}\n'
             f'worker2_trial_ids={worker2_trial_ids}'))

  def test_get_trial(self):
    study = self.create_test_study_with_trials(self.id())
    trial = study.get_trial(2).materialize()
    self.assertEqual(trial.metadata['future_id'], '2')
    self.assertEqual(trial.id, 2)

  def test_get_trials(self):
    study = self.create_test_study_with_trials(self.id())
    completed = study.trials(vz.TrialFilter(status={vz.TrialStatus.COMPLETED}))
    self.assertLen(list(completed), 2)

  def test_study_update_metadata(self):
    """Checks for correct merge behavior."""
    study = self.create_test_study(self.id())
    delta_metadata = vz.Metadata({'bar': 'bar_v'}, foo='foo_v')
    study.update_metadata(delta_metadata)
    self.assertEqual(study.materialize_problem_statement().metadata.get('bar'),
                     'bar_v')

    delta_metadata_2 = vz.Metadata({'bar': 'bar_w'})
    study.update_metadata(delta_metadata_2)

    problem_statement = study.materialize_problem_statement()
    self.assertEqual(problem_statement.metadata.get('bar'), 'bar_w')
    self.assertEqual(problem_statement.metadata.get('foo'), 'foo_v')

  ###########################################
  ## Tests on Trial class start here.
  ###########################################

  def test_delete_trial(self):
    study = self.create_test_study_with_trials(self.id())
    study.get_trial(2).delete()
    with self.assertRaises(Exception):
      study.get_trial(2)

  def test_complete_trial_no_measurements(self):
    study = self.create_test_study_with_trials(self.id())<commit_msg>Test two consecutive `parameters` access return the same results.<commit_after>"""Tests for client."""

import abc
import functools
from typing import Callable, Optional

from absl import logging
from vizier import pyvizier as vz
from vizier.client import client_abc

from absl.testing import parameterized

# Aliases are defined, so when you are developing a new client, you can
# swap it with your subclass. It makes your IDE understand which class
# you are using.
_StudyClient = client_abc.StudyInterface
_TrialClient = client_abc.TrialInterface


class VizierClientTestMixin(abc.ABC):

  @abc.abstractmethod
  def create_study(self, study_config: vz.ProblemStatement,
                   name: str) -> _StudyClient:
    """Create study given study config and study name."""
    pass


class MyMeta(type(VizierClientTestMixin), type(parameterized.TestCase)):
  pass


class TestCase(parameterized.TestCase, VizierClientTestMixin, metaclass=MyMeta):
  """Generic tests for cross-platform clients.

  This test provides basic coverage and it is not meant to be a comprehensive
  e2e test on all possible inputs to the client.

  Override `create_study` method.
  """

  def create_test_study(self, name: str) -> _StudyClient:
    """Creates a study."""
    logging.info('Creating study name=%s, testcasename=%s', name, self.id())
    problem = vz.ProblemStatement()
    problem.search_space.root.add_float_param('float', 0.0, 1.0)
    problem.metric_information.append(
        vz.MetricInformation(
            name='maximize_metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))
    study = self.create_study(problem, name)

    # TODO: Put this line back once we have __eq__ well-defined.
    # self.assertEqual(study.materialize_problem_statement(), problem)
    return study

  def test_create_or_load_study(self):
    study = self.create_test_study(self.id())
    study2 = study.from_resource_name(study.resource_name)
    self.assertEqual(study.resource_name, study2.resource_name)

    self.assertEqual(study.materialize_problem_statement(),
                     study2.materialize_problem_statement())

  def test_delete_study(self):
    study = self.create_test_study(self.id())
    resource_name = study.resource_name
    study.delete()

    with self.assertRaises(Exception):
      study.from_resource_name(resource_name)

  def _example_trials(self) -> list[vz.Trial]:
    trials = [
        vz.Trial().complete(vz.Measurement({'maximize_metric': 1.0})),
        vz.Trial().complete(vz.Measurement({'maximize_metric': 0.5})),
        vz.Trial(measurements=[vz.Measurement({'maximize_metric': 0.7})]),
        vz.Trial(is_requested=True)
    ]
    for idx, t in enumerate(trials):
      t.metadata['future_id'] = str(idx + 1)  # id to be assigned
    return trials

  def create_test_study_with_trials(self, name: str) -> _StudyClient:
    study = self.create_test_study(name)
    for t in self._example_trials():
      # TODO: Remove this.
      study._add_trial(t)  # pylint: disable=protected-access
    return study

  @parameterized.parameters(list(state for state in vz.StudyState))
  def test_set_state(self, target_state: vz.StudyState):
    study = self.create_test_study_with_trials(self.id())
    # TODO: Check that the study moved to the target state.
    # This test is currently a placeholder.
    try:
      study.set_state(target_state)
    except NotImplementedError:
      logging.exception('Set study state for %s is not implemented in %s',
                        target_state, type(self))

  def test_list_trials(self):
    study = self.create_test_study_with_trials(self.id())
    self.assertLen(list(study.trials()), 4)

  def test_trials_iter_and_get_are_equal(self):
    study = self.create_test_study_with_trials(self.id())
    all_trials = study.trials()
    self.assertEqual([t.id for t in all_trials],
                     [t.id for t in all_trials.get()])

  def test_optimal_trials_on_empty(self):
    study = self.create_test_study(self.id())
    self.assertEmpty(list(study.optimal_trials()))

  def test_optimal_trials_notempty(self):
    study = self.create_test_study_with_trials(self.id())
    optimal_trials = list(study.optimal_trials())
    self.assertLen(optimal_trials, 1)
    self.assertEqual(optimal_trials[0].materialize().status,
                     vz.TrialStatus.COMPLETED)

  def test_suggest_same_worker(self):
    # Given the same client id, should return the same trials.
    study = self.create_test_study(self.id())
    trials1 = study.suggest(count=5, client_id='worker1')
    self.assertLen(trials1, 5)
    trials2 = study.suggest(count=2, client_id='worker1')
    self.assertTrue({t.id for t in trials2}.issubset({t.id for t in trials1}))

  def test_suggest_different_workers(self):
    # Given different client ids, should generate new suggestions.
    study = self.create_test_study(self.id())
    worker1_trials = study.suggest(count=5, client_id='worker1')
    self.assertLen(worker1_trials, 5)
    worker1_trial_ids = {t.id for t in worker1_trials}
    worker2_trials = study.suggest(count=2, client_id='worker2')
    self.assertLen(worker2_trials, 2)
    worker2_trial_ids = {t.id for t in worker2_trials}
    self.assertEmpty(
        worker1_trial_ids.intersection(worker2_trial_ids),
        msg=(f'worker1_trial_ids={worker1_trial_ids}\n'
             f'worker2_trial_ids={worker2_trial_ids}'))

  def test_get_trial(self):
    study = self.create_test_study_with_trials(self.id())
    trial = study.get_trial(2).materialize()
    self.assertEqual(trial.metadata['future_id'], '2')
    self.assertEqual(trial.id, 2)

  def test_get_trials(self):
    study = self.create_test_study_with_trials(self.id())
    completed = study.trials(vz.TrialFilter(status={vz.TrialStatus.COMPLETED}))
    self.assertLen(list(completed), 2)

  def test_parameters(self):
    study = self.create_test_study_with_trials(self.id())
    self.assertEqual(
        study.get_trial(1).parameters,
        study.get_trial(1).parameters)

  def test_study_update_metadata(self):
    """Checks for correct merge behavior."""
    study = self.create_test_study(self.id())
    delta_metadata = vz.Metadata({'bar': 'bar_v'}, foo='foo_v')
    study.update_metadata(delta_metadata)
    self.assertEqual(study.materialize_problem_statement().metadata.get('bar'),
                     'bar_v')

    delta_metadata_2 = vz.Metadata({'bar': 'bar_w'})
    study.update_metadata(delta_metadata_2)

    problem_statement = study.materialize_problem_statement()
    self.assertEqual(problem_statement.metadata.get('bar'), 'bar_w')
    self.assertEqual(problem_statement.metadata.get('foo'), 'foo_v')

  ###########################################
  ## Tests on Trial class start here.
  ###########################################

  def test_delete_trial(self):
    study = self.create_test_study_with_trials(self.id())
    study.get_trial(2).delete()
    with self.assertRaises(Exception):
      study.get_trial(2)

  def test_complete_trial_no_measurements(self):
    study = self.create_test_study_with_trials(self.id())
{'fraction_non_alphanumeric': 0.12172095704814068, 'fraction_numerical': 0.008215624099164024, 'mean_word_length': 4.1434395848776875, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, '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': '3813659', 'n_tokens_mistral': 4659, 'n_tokens_neox': 4368, 'n_words': 884}