Skip to content

Commit

Permalink
✅ Added Tests for EnvironmentBasedModelMapper (#422)
Browse files Browse the repository at this point in the history
While working to confirm that the logic in codemodder-java and the
platform is in sync, I found opportunities to add missing tests and fix
some pitfalls.

/towards #work
  • Loading branch information
gilday authored Jul 18, 2024
1 parent 4a98fa1 commit 901144f
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 5 deletions.
Original file line number Diff line number Diff line change
@@ -1,22 +1,34 @@
package io.codemodder.plugins.llm;

import java.util.HashMap;
import java.util.Map;

/** Mapper that maps models to their deployment names based on environment variables. */
final class EnvironmentBasedModelMapper implements ModelMapper {
private static final String DEPLOYMENT_TEMPLATE = "CODEMODDER_AZURE_OPENAI_%s_DEPLOYMENT";

private final HashMap<Model, String> map = new HashMap<>();

EnvironmentBasedModelMapper() {
for (Model m : StandardModel.values()) {
final var deployment = System.getenv(String.format(DEPLOYMENT_TEMPLATE, m));
map.put(m, deployment == null ? m.id() : deployment);
this(System.getenv());
}

EnvironmentBasedModelMapper(final Map<String, String> environment) {
for (final Model model : StandardModel.values()) {
final var name = String.format(DEPLOYMENT_TEMPLATE, toEnvironmentVariableCase(model.id()));
final var deployment = environment.getOrDefault(name, model.id());
map.put(model, deployment);
}
}

@Override
public String getModelName(Model model) {
return map.get(model);
return map.getOrDefault(model, model.id());
}

/** Converts a model ID to environment variable casing. */
private static String toEnvironmentVariableCase(String input) {
return input.toUpperCase().replace('-', '_').replace('.', '_');
}

private static final String DEPLOYMENT_TEMPLATE = "CODEMODDER_AZURE_OPENAI_%s_DEPLOYMENT";
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
package io.codemodder.plugins.llm;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.withSettings;

import java.util.Map;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;

/** Unit tests for {@link EnvironmentBasedModelMapper}. */
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
final class EnvironmentBasedModelMapperTest {

private EnvironmentBasedModelMapper mapper;

@BeforeAll
void before() {
final var environment =
Map.of(
"CODEMODDER_AZURE_OPENAI_GPT_3_5_TURBO_0125_DEPLOYMENT",
"my-gpt-3.5-turbo",
"CODEMODDER_AZURE_OPENAI_GPT_4_0613_DEPLOYMENT",
"my-gpt-4",
"CODEMODDER_AZURE_OPENAI_GPT_4_TURBO_2024_04_09_DEPLOYMENT",
"my-gpt-4-turbo",
"CODEMODDER_AZURE_OPENAI_GPT_4O_2024_05_13_DEPLOYMENT",
"my-gpt-4o");
mapper = new EnvironmentBasedModelMapper(environment);
}

/** Spot checks one of the standard models to make sure the mapping works as expected */
@Test
void it_maps_model_name_to_deployment() {
final var name = mapper.getModelName(StandardModel.GPT_3_5_TURBO_0125);
assertThat(name).isEqualTo("my-gpt-3.5-turbo");
}

/**
* This is a meta-test that fails when we add a new standard model but forget to update the
* mapping in {@link #before()} to ensure that all standard models are covered.
*/
@EnumSource(StandardModel.class)
@ParameterizedTest
void it_looks_up_all_standard_models(final Model model) {
final var name = mapper.getModelName(model);
assertThat(name).isNotEqualTo(model.id()).startsWith("my-gpt");
}

@Test
void it_returns_model_id_when_no_mapping_exists() {
// GIVEN some model that doesn't have a mapping
final var model = mock(Model.class, withSettings().stubOnly());
when(model.id()).thenReturn("test");
final var name = mapper.getModelName(model);
assertThat(name).isEqualTo(model.id());
}
}

0 comments on commit 901144f

Please sign in to comment.