🇺🇸 You can also read the English version of this article on AWS Community Builders.
Uma das coisas mais comuns em grandes projetos pelos quais já passei são pipelines com muitos testes, levando muito tempo para serem executados. E com isso eu pensei: por que não ter um agente de IA para analisar meus PRs e, a partir das modificações, selecionar apenas os testes que precisam ser executados? Nesse artigo vou te mostrar como eu fiz isso e já adianto que provavelmente isso vai virar uma série de artigos.
(Se você não conhece o termo, o que estou descrevendo aqui é uma versão simplificada de **Test Impact Analysis: em vez de rodar a suÃte inteira a cada PR, o objetivo é mapear quais testes realmente têm relação com o que foi alterado.)
Prompts utilizados
Um dos feedbacks que recebi nos meus últimos artigos foi para deixar mais explÃcitos os prompts que utilizei para criar o projeto, e nesse eu usei apenas dois:
claude poderia me ajudar a fazer um projeto onde vou ter um agente de IA para analisar as linhas de códigos que foram modificadas em um PR e selecionar testes relacionados a essa modificação?
Como resposta para essa primeira pergunta, ele criou pra mim um diagrama de como basicamente isso poderia funcionar, sem nenhuma integração a pipeline ou coisa do tipo:
Depois disso, simplesmente pedi para ele criar um prompt para o Claude Code (a IA que utilizo) e, depois de colar o prompt gerado:
A primeira versão do projeto foi criado no seguinte formato:
main.py CLI entrypoint — wires everything together
config.py Loads & validates ANTHROPIC_API_KEY / GITHUB_TOKEN
src/
github_client.py Fetches PR files & metadata from the GitHub API (with retries)
diff_parser.py Maps diff hunks to changed functions/classes via AST
test_discovery.py Indexes pytest test functions across the repo
ai_agent.py Calls Claude to select relevant tests
reporter.py Renders results as a table or JSON
tests/ Unit tests for the parser & discovery modules
Vale destacar o diff_parser.py: em vez de simplesmente olhar quais linhas mudaram no diff, ele usa AST (Abstract Syntax Tree) para entender quais funções e classes foram alteradas. Essa diferença importa bastante — duas linhas modificadas dentro da mesma função valem como uma única mudança de sÃmbolo, o que deixa a seleção de testes mais precisa do que uma comparação puramente textual.
Dando destaque ao nosso ai_agent.py, pois é nele que acontece a mágica, temos alguns argumentos interessantes de entender:
Relevance = Literal["low", "medium", "high"]
Ter um filtro de relevância é importante para auxiliar a IA no momento em que ela for analisar, entre os testes disponÃveis, quais realmente precisam ser executados. Na prática, esse filtro também vira um dial de confiança para quem está usando a ferramenta: em um pipeline de CI mais crÃtico, faz sentido rodar só os testes marcados como high; em uma verificação mais exploratória, pode valer incluir low e medium também.
O prompt utilizado foi esse:
_SYSTEM_PROMPT = """\
You are a senior software engineer who selects which pytest tests are relevant \
to review for a given pull request.
You will receive a JSON object with:
- "min_relevance": the minimum relevance level to include ("low", "medium", or "high")
- "changed_symbols": a list of functions/classes/methods (or whole files) modified in the PR,
each with file_path, symbol_name, symbol_type, and lines_changed
- "available_tests": a list of test functions discovered in the repository, each with
test_path, test_name, imports, and source_references
Decide which of the available tests are relevant to the changed symbols by considering:
- Whether a test's source_references directly name a changed symbol
- Whether a test's imports point to the module containing a changed symbol
- One level of transitivity: a test that references something which itself references
a changed symbol (e.g. a helper that calls the changed function)
- The semantic similarity between test names/paths and changed symbol names
Respond with ONLY a JSON array (no prose, no markdown fences, no text outside the array).
Each element must have exactly this shape:
{
"test_path": "<string>",
"test_name": "<string>",
"relevance": "<low|medium|high>",
"confidence_score": <float between 0.0 and 1.0>,
"reason": "<short string explaining why this test is relevant>"
}
Only include tests whose relevance is greater than or equal to the given min_relevance
("low" < "medium" < "high"). If no tests are relevant, return an empty array: []
"""
Um ponto que vale destacar nesse prompt é a "um nÃvel de transitividade": o agente não olha só se um teste referencia diretamente o sÃmbolo alterado, mas também se ele referencia algo que, por sua vez, referencia o sÃmbolo alterado (por exemplo, um helper chamado pela função modificada). Isso ajuda a pegar casos de teste que passariam despercebidos numa busca só por nome.
Seleção do AI Agent
A primeira mudança que fiz foi pedir para a IA adaptar o código de forma que fosse possÃvel passar qual IA utilizar como parâmetro, em vez de deixá-la fixa:
"""Thin wrappers so AIAgent can call Claude, Gemini, or ChatGPT interchangeably."""
from typing import Protocol
from config import Config
class LLMProvider(Protocol):
"""Something that can turn a system + user prompt into a text completion."""
def complete(self, system_prompt: str, user_prompt: str) -> str: ...
class ClaudeProvider:
"""Calls Anthropic's Claude models."""
def __init__(self, api_key: str, model: str) -> None:
from anthropic import Anthropic
self._client = Anthropic(api_key=api_key)
self._model = model
def complete(self, system_prompt: str, user_prompt: str) -> str:
response = self._client.messages.create(
model=self._model,
max_tokens=4096,
system=system_prompt,
messages=[{"role": "user", "content": user_prompt}],
)
return "".join(block.text for block in response.content if getattr(block, "type", None) == "text")
class GeminiProvider:
"""Calls Google's Gemini models."""
def __init__(self, api_key: str, model: str) -> None:
from google import genai
self._client = genai.Client(api_key=api_key)
self._model = model
def complete(self, system_prompt: str, user_prompt: str) -> str:
from google.genai import types
response = self._client.models.generate_content(
model=self._model,
contents=user_prompt,
config=types.GenerateContentConfig(system_instruction=system_prompt),
)
return response.text or ""
class ChatGPTProvider:
"""Calls OpenAI's ChatGPT models."""
def __init__(self, api_key: str, model: str) -> None:
from openai import OpenAI
self._client = OpenAI(api_key=api_key)
self._model = model
def complete(self, system_prompt: str, user_prompt: str) -> str:
response = self._client.chat.completions.create(
model=self._model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
return response.choices[0].message.content or ""
def create_provider(config: Config) -> LLMProvider:
"""Builds the LLMProvider for config.provider, using the matching API key and model."""
if config.provider == "claude":
assert config.anthropic_api_key is not None
return ClaudeProvider(api_key=config.anthropic_api_key, model=config.claude_model)
if config.provider == "gemini":
assert config.gemini_api_key is not None
return GeminiProvider(api_key=config.gemini_api_key, model=config.gemini_model)
if config.provider == "chatgpt":
assert config.openai_api_key is not None
return ChatGPTProvider(api_key=config.openai_api_key, model=config.chatgpt_model)
raise ValueError(f"Unknown provider: {config.provider!r}")
O uso do Protocol aqui é o que permite trocar de provedor sem alterar o restante do código: qualquer classe que implemente complete(system_prompt, user_prompt) serve, então AIAgent nunca precisa saber se está falando com Claude, Gemini ou ChatGPT.
Dessa forma, fica a critério de quem está utilizando o pacote selecionar qual IA será usada, podendo isso ser passado como parâmetro no momento da execução:
python main.py --pr 42 --repo owner/repo --provider gemini
python main.py --pr 42 --repo owner/repo --provider chatgpt
python main.py --pr 42 --repo owner/repo # default: claude
Validando se o que foi gerado realmente está certo
Para a primeira validação, pesquisei alguns repositórios públicos no GitHub e, nesse caso, selecionei este PR: https://github.com/pallets/flask/pull/6043
python main.py --repo pallets/flask --pr 6043 --provider gemini
E obtive o seguinte resultado:
Só que agora entra a parte complicada, sempre temos nos perguntar quando utilizamos IA:
Essa resposta faz sentido?
E, de cara, dá para ver que não — não está certo. Quando passamos um repositório diferente, ele continua buscando no repositório local, ignorando completamente o repositório do PR informado. Para corrigir isso, precisamos garantir que, sempre que for um repositório externo, o agente acesse as informações diretamente dele e para isso, vamos clonar temporariamente uma branch do repositório:
"""Clones a single branch of a GitHub repository into a temporary directory for local scanning."""
import logging
import shutil
import subprocess
import tempfile
logger = logging.getLogger(__name__)
class RepoFetchError(RuntimeError):
"""Raised when the target repository/branch cannot be cloned."""
class RepoFetcher:
"""Shallow-clones a single branch of a repository so its tests can be scanned locally."""
def __init__(self, token: str | None = None) -> None:
self._token = token
self._tmp_dir: str | None = None
def fetch(self, repo: str, ref: str) -> str:
"""Clones `repo` (owner/name) at branch `ref` and returns the local checkout path."""
self._tmp_dir = tempfile.mkdtemp(prefix="ai-test-selector-")
cmd = [
"git",
"clone",
"--depth",
"1",
"--branch",
ref,
"--single-branch",
self._clone_url(repo),
self._tmp_dir,
]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
shutil.rmtree(self._tmp_dir, ignore_errors=True)
self._tmp_dir = None
raise RepoFetchError(f"Failed to clone {repo}@{ref}: {result.stderr.strip()}")
return self._tmp_dir
def cleanup(self) -> None:
"""Removes the temporary checkout, if one was created."""
if self._tmp_dir:
shutil.rmtree(self._tmp_dir, ignore_errors=True)
self._tmp_dir = None
def _clone_url(self, repo: str) -> str:
if self._token:
return f"https://x-access-token:{self._token}@github.com/{repo}.git"
return f"https://github.com/{repo}.git"
Repare no --depth 1 e --single-branch do clone: como o objetivo aqui é só escanear os testes do repositório naquele ponto especÃfico, não faz sentido trazer todo o histórico de commits — um shallow clone resolve o problema mais rápido e consumindo bem menos espaço em disco.
Rodando o mesmo comando, temos o seguinte resultado:
Em resumo, ele selecionou:
Selected 192/403 tests (avg confidence: 0.76)
pytest tests/test_appctx.py tests/test_async.py tests/test_basic.py tests/test_blueprints.py tests/test_cli.py tests/test_converters.py
tests/test_helpers.py tests/test_json.py tests/test_logging.py tests/test_regression.py tests/test_reqctx.py tests/test_request.py
tests/test_session_interface.py tests/test_signals.py tests/test_subclassing.py tests/test_templating.py tests/test_testing.py
tests/test_user_error_handler.py tests/test_views.py -v
É possÃvel ver que com a IA, ao invés de rodarmos toda a suite, poderÃamos rodar apenas 48% dela, reduzindo o tempo de execução e retornando um feedback muito mais rápido para quem está desenvolvendo a feature.
Conclusão
Como o artigo já está ficando grande, vou parar por aqui nesta primeira etapa. Revisando, o foco desta primeira parte foi refinar o prompt utilizado pela IA para gerar o código e, a partir disso, revisar o que foi gerado. Nela tivemos duas mudanças significativas:
- Adaptação do código para utilizar diferentes provedores de IA, de forma que o projeto não ficasse preso a um único modelo.
- Ajuste na busca dos testes, para que ela considerasse os testes disponÃveis no repositório relacionado ao PR informado, e não no repositório local.
Próximos passos
Esse projeto ainda está bem no inÃcio, e pretendo continuar evoluindo ele nos próximos artigos. Já durante essa primeira parte, revisando o que foi gerado e testando em repositórios reais, surgiram algumas ideias novas que quero explorar:
- Integração com pipelines: hoje o projeto roda de forma isolada, via linha de comando, mas o próximo passo natural é pensar em como plugar isso direto no CI/CD — por exemplo, como uma action do GitHub Actions que roda automaticamente a cada PR aberto e comenta com os testes selecionados, ou até já dispara só o subconjunto relevante.
- Performance na parte de IA: com repositórios maiores, o número de testes disponÃveis para avaliar cresce bastante, e isso significa prompts maiores e mais tempo de resposta. Quero investigar formas de otimizar essa etapa — seja filtrando os candidatos antes de mandar para a IA, seja quebrando a avaliação em lotes menores, ou até testando modelos diferentes para essa tarefa especÃfica de classificação.
Se você quiser acompanhar a evolução do projeto (ou já colocar a mão na massa também), o código está aberto no meu repositório: ai-agent-test-selector
No próximo artigo, vou focar no prompt utilizado pela IA para selecionar os testes propriamente dito, além de adaptar o wrapper para ler repositórios que não sejam apenas em Python, mas também em outras linguagens.
Ficou alguma dúvida sobre alguma parte do código, do prompt ou da arquitetura? Deixa nos comentários — quero muito saber o que vocês acharam e trocar ideia sobre outros cenários onde isso poderia ser aplicado.
Top comments (1)
Maravilhoso.