-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
✅ Added Tests for EnvironmentBasedModelMapper (#422)
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
Showing
2 changed files
with
79 additions
and
5 deletions.
There are no files selected for viewing
22 changes: 17 additions & 5 deletions
22
...odder-plugin-llm/src/main/java/io/codemodder/plugins/llm/EnvironmentBasedModelMapper.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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"; | ||
} |
62 changes: 62 additions & 0 deletions
62
...r-plugin-llm/src/test/java/io/codemodder/plugins/llm/EnvironmentBasedModelMapperTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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()); | ||
} | ||
} |