chunks.txt

Chunk ID: alembic/env.py::1
from logging.config import fileConfig

from sqlalchemy import engine_from_config
from sqlalchemy import pool

from alembic import context
from src.database.core import Base  # Wherever your Base is defined

from src.auth.models import CowboyUser
from src.repo.models import RepoConfig

from src.ast.models import NodeModel
from src.test_modules.models import TestModuleModel
from src.target_code.models import TargetCodeModel
from src.coverage.models import CoverageModel
from src.test_gen.models import AugmentTestResult
from src.stats.models import RepoStats


from src.config import SQLALCHEMY_DATABASE_URI


# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config

# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = Base.metadata

# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.


--------------------------------------------------------------------------------

Chunk ID: alembic/env.py::2
def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode.

    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.

    Calls to context.execute() here emit the given string to the
    script output.

    """
    url = SQLALCHEMY_DATABASE_URI
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()


--------------------------------------------------------------------------------

Chunk ID: alembic/env.py::3
def patch_sql_url(config):
    config["sqlalchemy.url"] = SQLALCHEMY_DATABASE_URI

    return config


def run_migrations_online() -> None:
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    connectable = engine_from_config(
        patch_sql_url(config.get_section(config.config_ini_section, {})),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(connection=connection, target_metadata=target_metadata)

        with context.begin_transaction():
            context.run_migrations()


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()


--------------------------------------------------------------------------------

Chunk ID: versions/13e10adf8e27_added_repo_id.py::1
"""added repo_id

Revision ID: 13e10adf8e27
Revises: 8a518189662f
Create Date: 2024-05-30 02:50:32.704308

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '13e10adf8e27'
down_revision: Union[str, None] = '8a518189662f'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_stats', sa.Column('repo_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'repo_stats', 'repo_config', ['repo_id'], ['id'], ondelete='CASCADE')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'repo_stats', type_='foreignkey')
    op.drop_column('repo_stats', 'repo_id')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/2bc1adebc872_added_repoid_to_test_results_model.py::1
"""added repoid to test_results_model

Revision ID: 2bc1adebc872
Revises: bf5d46d450f5
Create Date: 2024-05-20 00:12:28.376639

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "2bc1adebc872"
down_revision: Union[str, None] = "bf5d46d450f5"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/2bc1adebc872_added_repoid_to_test_results_model.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "augment_test_results", sa.Column("repo_id", sa.Integer(), nullable=True)
    )
    op.create_foreign_key(
        None, "augment_test_results", "repo_config", ["repo_id"], ["id"]
    )
    op.add_column("coverage", sa.Column("test_result_id", sa.Integer(), nullable=True))
    op.drop_constraint("coverage_test_result_fkey", "coverage", type_="foreignkey")
    op.create_foreign_key(
        None, "coverage", "augment_test_results", ["test_result_id"], ["id"]
    )
    op.drop_column("coverage", "test_result")
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/2bc1adebc872_added_repoid_to_test_results_model.py::3
def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "coverage",
        sa.Column("test_result", sa.INTEGER(), autoincrement=False, nullable=True),
    )
    op.drop_constraint(None, "coverage", type_="foreignkey")
    op.create_foreign_key(
        "coverage_test_result_fkey",
        "coverage",
        "augment_test_results",
        ["test_result"],
        ["id"],
    )
    op.drop_column("coverage", "test_result_id")
    op.drop_constraint(None, "augment_test_results", type_="foreignkey")
    op.drop_column("augment_test_results", "repo_id")
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/2ee3ecc05c40_added_filepath_to_nodemodel_and_name_to_.py::1
"""added filepath to NodeModel and name to TestModuleModel

Revision ID: 2ee3ecc05c40
Revises: df1d20e2c832
Create Date: 2024-05-01 11:12:35.483141

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '2ee3ecc05c40'
down_revision: Union[str, None] = 'df1d20e2c832'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('nodes', sa.Column('testfilepath', sa.String(), nullable=True))
    op.add_column('test_modules', sa.Column('name', sa.String(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('test_modules', 'name')
    op.drop_column('nodes', 'testfilepath')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/3c7b1822da96_added_repostats.py::1
"""added RepoStats

Revision ID: 3c7b1822da96
Revises: 4bf808259948
Create Date: 2024-05-30 01:31:04.314037

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '3c7b1822da96'
down_revision: Union[str, None] = '4bf808259948'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/3e1fd387befe_fixed_some_relations.py::1
"""fixed some relations

Revision ID: 3e1fd387befe
Revises: d0c885df6cfa
Create Date: 2024-05-02 16:20:44.582061

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '3e1fd387befe'
down_revision: Union[str, None] = 'd0c885df6cfa'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/3e1fd387befe_fixed_some_relations.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('nodes', sa.Column('target_code_id', sa.Integer(), nullable=True))
    op.drop_constraint('nodes_repo_id_fkey', 'nodes', type_='foreignkey')
    op.drop_constraint('nodes_test_module_id_fkey', 'nodes', type_='foreignkey')
    op.create_foreign_key(None, 'nodes', 'target_code', ['target_code_id'], ['id'], ondelete='CASCADE')
    op.create_foreign_key(None, 'nodes', 'repo_config', ['repo_id'], ['id'], ondelete='CASCADE')
    op.create_foreign_key(None, 'nodes', 'test_modules', ['test_module_id'], ['id'], ondelete='CASCADE')
    op.add_column('target_code', sa.Column('test_module_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'target_code', 'test_modules', ['test_module_id'], ['id'], ondelete='CASCADE')
    op.drop_column('target_code', 'class_scope')
    op.drop_column('target_code', 'func_scope')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/3e1fd387befe_fixed_some_relations.py::3
def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('target_code', sa.Column('func_scope', sa.INTEGER(), autoincrement=False, nullable=True))
    op.add_column('target_code', sa.Column('class_scope', sa.INTEGER(), autoincrement=False, nullable=True))
    op.drop_constraint(None, 'target_code', type_='foreignkey')
    op.drop_column('target_code', 'test_module_id')
    op.drop_constraint(None, 'nodes', type_='foreignkey')
    op.drop_constraint(None, 'nodes', type_='foreignkey')
    op.drop_constraint(None, 'nodes', type_='foreignkey')
    op.create_foreign_key('nodes_test_module_id_fkey', 'nodes', 'test_modules', ['test_module_id'], ['id'])
    op.create_foreign_key('nodes_repo_id_fkey', 'nodes', 'repo_config', ['repo_id'], ['id'])
    op.drop_column('nodes', 'target_code_id')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/4182504b989a_renamed_feedback_to_decide.py::1
"""renamed feedback to decide

Revision ID: 4182504b989a
Revises: 2bc1adebc872
Create Date: 2024-05-20 02:06:40.051434

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '4182504b989a'
down_revision: Union[str, None] = '2bc1adebc872'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('augment_test_results', sa.Column('decide', sa.Integer(), nullable=True))
    op.drop_column('augment_test_results', 'feedback')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('augment_test_results', sa.Column('feedback', sa.INTEGER(), autoincrement=False, nullable=True))
    op.drop_column('augment_test_results', 'decide')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/486ee62cfdb2_added_main_and_remote_to_repoconfig.py::1
"""added main and remote to repoconfig

Revision ID: 486ee62cfdb2
Revises: 53cfcb81976f
Create Date: 2024-05-13 18:00:08.309289

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '486ee62cfdb2'
down_revision: Union[str, None] = '53cfcb81976f'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('remote', sa.String(), nullable=True))
    op.add_column('repo_config', sa.Column('main', sa.String(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('repo_config', 'main')
    op.drop_column('repo_config', 'remote')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/4b2308ac52fa_finally_working.py::1
"""finally working ... ?

Revision ID: 4b2308ac52fa
Revises: f8709f78a9a4
Create Date: 2024-04-19 22:49:15.298246

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '4b2308ac52fa'
down_revision: Union[str, None] = 'f8709f78a9a4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/4b2308ac52fa_finally_working.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('cowboy_user',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('email', sa.String(), nullable=True),
    sa.Column('password', sa.LargeBinary(), nullable=False),
    sa.Column('last_mfa_time', sa.DateTime(), nullable=True),
    sa.Column('experimental_features', sa.Boolean(), nullable=True),
    sa.Column('created_at', sa.DateTime(), nullable=True),
    sa.Column('updated_at', sa.DateTime(), nullable=True),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('email')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('cowboy_user')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/4be689188190_added_auto_gen_to_test_modules.py::1
"""added auto_gen to test_modules


Revision ID: 4be689188190
Revises: e3cae2f0fb3e
Create Date: 2024-05-16 22:49:42.800773

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '4be689188190'
down_revision: Union[str, None] = 'e3cae2f0fb3e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('test_modules', sa.Column('auto_gen', sa.Boolean(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('test_modules', 'auto_gen')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/4bf808259948_added_session_id_to_testresults.py::1
"""added session_id to TestResults

Revision ID: 4bf808259948
Revises: 7ab0269403a8
Create Date: 2024-05-27 15:02:17.441346

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '4bf808259948'
down_revision: Union[str, None] = '7ab0269403a8'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('augment_test_results', sa.Column('session_id', sa.String(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('augment_test_results', 'session_id')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/50c4f4a1bd1a_changed_cowboyusers_table.py::1
"""Changed CowboyUsers table

Revision ID: 50c4f4a1bd1a
Revises: e63f0f3ddd8b
Create Date: 2024-04-19 22:33:56.432302

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '50c4f4a1bd1a'
down_revision: Union[str, None] = 'e63f0f3ddd8b'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/529e2719b577_added_relationship_from_cowboyuser_to_.py::1
"""added relationship from CowboyUser to RepoConfig, cascade deletes

Revision ID: 529e2719b577
Revises: 13e10adf8e27
Create Date: 2024-05-31 17:37:21.980761

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '529e2719b577'
down_revision: Union[str, None] = '13e10adf8e27'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/53cfcb81976f_repair_add_cloned_folders_deleteforked_.py::1
"""REPAIR: add cloned_folders, deleteforked_urls

Revision ID: 53cfcb81976f
Revises: 68e7bf63522d
Create Date: 2024-05-12 18:57:15.357469

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '53cfcb81976f'
down_revision: Union[str, None] = '68e7bf63522d'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('cloned_folders', sa.String(), nullable=True))
    op.drop_column('repo_config', 'forked_url')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('forked_url', sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_column('repo_config', 'cloned_folders')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/5a9b1afb76e3_added_repo_model.py::1
"""Added repo_model

Revision ID: 5a9b1afb76e3
Revises: e635c218de2a
Create Date: 2024-04-23 21:45:56.527438

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '5a9b1afb76e3'
down_revision: Union[str, None] = 'e635c218de2a'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/5a9b1afb76e3_added_repo_model.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('repo_config',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('repo_name', sa.String(), nullable=True),
    sa.Column('url', sa.String(), nullable=True),
    sa.Column('forked_url', sa.String(), nullable=True),
    sa.Column('cloned_folders', sa.String(), nullable=True),
    sa.Column('source_folder', sa.String(), nullable=True),
    sa.Column('python_conf', sa.JSON(), nullable=True),
    sa.Column('user_id', sa.Integer(), nullable=True),
    sa.ForeignKeyConstraint(['user_id'], ['cowboy_user.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('repo_config')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/68e7bf63522d_updated_range_on_targecodemodel.py::1
"""updated range on TargeCodeModel

Revision ID: 68e7bf63522d
Revises: 3e1fd387befe
Create Date: 2024-05-02 16:26:07.546110

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '68e7bf63522d'
down_revision: Union[str, None] = '3e1fd387befe'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('target_code', sa.Column('start', sa.Integer(), nullable=True))
    op.add_column('target_code', sa.Column('end', sa.Integer(), nullable=True))
    op.drop_column('target_code', 'range')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('target_code', sa.Column('range', sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_column('target_code', 'end')
    op.drop_column('target_code', 'start')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/7ab0269403a8_added_exp_id_to_repoconfig.py::1
"""added exp_id to RepoConfig

Revision ID: 7ab0269403a8
Revises: 8cbcddf29c9e
Create Date: 2024-05-25 04:32:44.055230

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '7ab0269403a8'
down_revision: Union[str, None] = '8cbcddf29c9e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('is_experiment', sa.Boolean(), nullable=True))
    op.add_column('test_modules', sa.Column('experiment_id', sa.String(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('test_modules', 'experiment_id')
    op.drop_column('repo_config', 'is_experiment')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/7f8df6212797_added_augmenttestresult_models.py::1
"""added augmentTestResult models

Revision ID: 7f8df6212797
Revises: 4be689188190
Create Date: 2024-05-19 14:22:05.027883

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '7f8df6212797'
down_revision: Union[str, None] = '4be689188190'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/7f8df6212797_added_augmenttestresult_models.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('augment_test_results',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=True),
    sa.Column('test_case', sa.String(), nullable=True),
    sa.Column('cov_plus', sa.Integer(), nullable=True),
    sa.Column('feedback', sa.Integer(), nullable=True),
    sa.Column('commit_hash', sa.String(), nullable=True),
    sa.Column('testfile', sa.String(), nullable=True),
    sa.Column('classname', sa.String(), nullable=True),
    sa.Column('test_module_id', sa.Integer(), nullable=True),
    sa.ForeignKeyConstraint(['test_module_id'], ['test_modules.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('augment_test_results')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/87a33ce1322a_added_cascade_delete_from_testmodule_.py::1
"""added cascade delete from TestModule -> AugmentTestResults

Revision ID: 87a33ce1322a
Revises: 529e2719b577
Create Date: 2024-06-20 19:51:54.939523

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '87a33ce1322a'
down_revision: Union[str, None] = '529e2719b577'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint('augment_test_results_test_module_id_fkey', 'augment_test_results', type_='foreignkey')
    op.create_foreign_key(None, 'augment_test_results', 'test_modules', ['test_module_id'], ['id'], ondelete='CASCADE')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'augment_test_results', type_='foreignkey')
    op.create_foreign_key('augment_test_results_test_module_id_fkey', 'augment_test_results', 'test_modules', ['test_module_id'], ['id'])
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/8a518189662f_added_repostats_actually.py::1
"""added RepoStats (actually)

Revision ID: 8a518189662f
Revises: 3c7b1822da96
Create Date: 2024-05-30 01:32:54.606182

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '8a518189662f'
down_revision: Union[str, None] = '3c7b1822da96'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('repo_stats',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('total_tests', sa.Integer(), nullable=True),
    sa.Column('accepted_tests', sa.Integer(), nullable=True),
    sa.Column('rejected_tests', sa.Integer(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('repo_stats')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/8cbcddf29c9e_addde_language_field_to_repoconfig_model.py::1
"""addde language field to RepoConfig model

Revision ID: 8cbcddf29c9e
Revises: 4182504b989a
Create Date: 2024-05-20 18:57:41.319324

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '8cbcddf29c9e'
down_revision: Union[str, None] = '4182504b989a'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('language', sa.String(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('repo_config', 'language')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/91a43f59cf7d_removed_cloned_folders_from_repoconfig.py::1
"""removed cloned_folders from RepoConfig

Revision ID: 91a43f59cf7d
Revises: baca7f3acb5c
Create Date: 2024-04-30 23:40:02.904914

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '91a43f59cf7d'
down_revision: Union[str, None] = 'baca7f3acb5c'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('repo_config', 'cloned_folders')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('repo_config', sa.Column('cloned_folders', sa.VARCHAR(), autoincrement=False, nullable=True))
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/94e8e80b2fc8_mistake.py::1
"""mistake

Revision ID: 94e8e80b2fc8
Revises: d428a766b0ef
Create Date: 2024-05-01 21:43:01.088713

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '94e8e80b2fc8'
down_revision: Union[str, None] = 'd428a766b0ef'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/b4b0ffad026d_added_testmodules.py::1
"""Added testmodules

Revision ID: b4b0ffad026d
Revises: 5a9b1afb76e3
Create Date: 2024-04-30 22:08:23.733495

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'b4b0ffad026d'
down_revision: Union[str, None] = '5a9b1afb76e3'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/b4b0ffad026d_added_testmodules.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('target_code',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('range', sa.String(), nullable=True),
    sa.Column('lines', sa.String(), nullable=True),
    sa.Column('filepath', sa.String(), nullable=True),
    sa.Column('func_scope', sa.String(), nullable=True),
    sa.Column('class_scope', sa.String(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('test_modules',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('testfilepath', sa.String(), nullable=True),
    sa.Column('commit_sha', sa.String(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('nodes',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=True),
    sa.Column('node_type', sa.String(), nullable=True),
    sa.Column('test_module_id', sa.Integer(), nullable=True),
    sa.ForeignKeyConstraint(['test_module_id'], ['test_modules.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('nodes')
    op.drop_table('test_modules')
    op.drop_table('target_code')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/b9671daa1ec7_added_coverage_table.py::1
"""added coverage table

Revision ID: b9671daa1ec7
Revises: 486ee62cfdb2
Create Date: 2024-05-16 01:06:36.599380

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'b9671daa1ec7'
down_revision: Union[str, None] = '486ee62cfdb2'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/b9671daa1ec7_added_coverage_table.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('coverage',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('filename', sa.String(), nullable=False),
    sa.Column('covered_lines', sa.String(), nullable=False),
    sa.Column('missing_lines', sa.String(), nullable=False),
    sa.Column('stmts', sa.Integer(), nullable=False),
    sa.Column('misses', sa.Integer(), nullable=False),
    sa.Column('covered', sa.Integer(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )
    op.add_column('target_code', sa.Column('coverage_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'target_code', 'coverage', ['coverage_id'], ['id'], ondelete='CASCADE')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'target_code', type_='foreignkey')
    op.drop_column('target_code', 'coverage_id')
    op.drop_table('coverage')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/baca7f3acb5c_added_repoid_to_testmodulemodel.py::1
"""added repoId to TestModuleModel

Revision ID: baca7f3acb5c
Revises: b4b0ffad026d
Create Date: 2024-04-30 23:27:30.931847

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'baca7f3acb5c'
down_revision: Union[str, None] = 'b4b0ffad026d'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


--------------------------------------------------------------------------------

Chunk ID: versions/baca7f3acb5c_added_repoid_to_testmodulemodel.py::2
def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('nodes', sa.Column('repo_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'nodes', 'repo_config', ['repo_id'], ['id'])
    op.add_column('test_modules', sa.Column('repo_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'test_modules', 'repo_config', ['repo_id'], ['id'])
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'test_modules', type_='foreignkey')
    op.drop_column('test_modules', 'repo_id')
    op.drop_constraint(None, 'nodes', type_='foreignkey')
    op.drop_column('nodes', 'repo_id')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/bf5d46d450f5_added_coverage_test_result_relation.py::1
"""added coverage/test_result relation

Revision ID: bf5d46d450f5
Revises: 7f8df6212797
Create Date: 2024-05-19 23:59:00.090831

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'bf5d46d450f5'
down_revision: Union[str, None] = '7f8df6212797'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('augment_test_results', 'cov_plus')
    op.add_column('coverage', sa.Column('test_result', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'coverage', 'augment_test_results', ['test_result'], ['id'])
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'coverage', type_='foreignkey')
    op.drop_column('coverage', 'test_result')
    op.add_column('augment_test_results', sa.Column('cov_plus', sa.INTEGER(), autoincrement=False, nullable=True))
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/d0c885df6cfa_change_targetcode_tables_from_varchar_.py::1
"""Change TargetCode tables from varchar to int

Revision ID: d0c885df6cfa
Revises: 94e8e80b2fc8
Create Date: 2024-05-02 04:15:17.169023

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "d0c885df6cfa"
down_revision: Union[str, None] = "94e8e80b2fc8"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.execute(
        "ALTER TABLE target_code ALTER COLUMN func_scope SET DATA TYPE int4 USING (func_scope::int);"
    )
    op.execute(
        "ALTER TABLE target_code ALTER COLUMN class_scope SET DATA TYPE int4 USING (class_scope::int);"
    )


def downgrade() -> None:
    pass


--------------------------------------------------------------------------------

Chunk ID: versions/d428a766b0ef_added_backref_from_tmmodel_to_nodes.py::1
"""added backref from TMModel to Nodes

Revision ID: d428a766b0ef
Revises: 2ee3ecc05c40
Create Date: 2024-05-01 21:37:52.458142

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'd428a766b0ef'
down_revision: Union[str, None] = '2ee3ecc05c40'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/df1d20e2c832_added_cascade_delete_from_repo_to_tms_.py::1
"""added cascade delete from repo to tms and nodes

Revision ID: df1d20e2c832
Revises: 91a43f59cf7d
Create Date: 2024-04-30 23:53:52.137184

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'df1d20e2c832'
down_revision: Union[str, None] = '91a43f59cf7d'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/e3cae2f0fb3e_added_repoid_to_coverage.py::1
"""added repoid to coverage

Revision ID: e3cae2f0fb3e
Revises: b9671daa1ec7
Create Date: 2024-05-16 02:32:20.376283

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'e3cae2f0fb3e'
down_revision: Union[str, None] = 'b9671daa1ec7'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('coverage', sa.Column('repo_id', sa.Integer(), nullable=True))
    op.create_foreign_key(None, 'coverage', 'repo_config', ['repo_id'], ['id'])
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'coverage', type_='foreignkey')
    op.drop_column('coverage', 'repo_id')
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/e635c218de2a_added_repo_model_and_added_relation_.py::1
"""Added repo model and added relation from user to many repos

Revision ID: 84e9f0d53973
Revises: 4b2308ac52fa
Create Date: 2024-04-23 21:30:25.896998

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '84e9f0d53973'
down_revision: Union[str, None] = '4b2308ac52fa'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/e635c218de2a_test2.py::1
"""test2

Revision ID: e635c218de2a
Revises: 84e9f0d53973
Create Date: 2024-04-23 21:40:25.046636

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'e635c218de2a'
down_revision: Union[str, None] = '84e9f0d53973'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/e63f0f3ddd8b_added_initial_tables.py::1
"""Added initial tables

Revision ID: e63f0f3ddd8b
Revises: 
Create Date: 2024-04-19 22:25:32.303804

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'e63f0f3ddd8b'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/f8709f78a9a4_updated_base.py::1
"""updated base

Revision ID: f8709f78a9a4
Revises: 50c4f4a1bd1a
Create Date: 2024-04-19 22:44:00.590963

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'f8709f78a9a4'
down_revision: Union[str, None] = '50c4f4a1bd1a'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###


--------------------------------------------------------------------------------

Chunk ID: versions/faccdf9696ed_.py::1
"""empty message

Revision ID: faccdf9696ed
Revises: 87a33ce1322a
Create Date: 2024-11-27 15:37:38.654668

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'faccdf9696ed'
down_revision: Union[str, None] = '87a33ce1322a'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    pass


def downgrade() -> None:
    pass


--------------------------------------------------------------------------------

Chunk ID: main.py::1
from typing import Optional, Final
from contextvars import ContextVar

from fastapi import FastAPI, status
from fastapi.responses import JSONResponse
from pydantic import ValidationError

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response, StreamingResponse

from sqlalchemy.orm import sessionmaker

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse

import uvicorn
from logging import getLogger, Filter

from src.queue.core import TaskQueue
from src.auth.views import auth_router
from src.repo.views import repo_router
from src.test_modules.views import tm_router
from src.queue.views import task_queue_router
from src.test_gen.views import test_gen_router
from src.target_code.views import tgtcode_router
from src.experiments.views import exp_router
from src.health.views import health_router
from src.exceptions import CowboyRunTimeException

from src.extensions import init_sentry
from src.config import PORT

import uuid

from src.sync_repos import start_sync_thread
from src.database.core import engine

# NEWTODO: this very likely gonna be broken when we start running multiple worker
# threads in uvicorn since each thread will get its own copy.. which should actually 
# be fine since we just issue a new authentication request
# global var token_registry
from src.token_registry import token_registry


log = getLogger(__name__)


async def not_found(request, exc):
    return JSONResponse(
        status_code=status.HTTP_404_NOT_FOUND,
        content={"detail": [{"msg": "Not Found."}]},
    )


--------------------------------------------------------------------------------

Chunk ID: main.py::2
class ExceptionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        try:
            response = await call_next(request)
            return response
        except CowboyRunTimeException as e:
            return JSONResponse(
                status_code=status.HTTP_400_BAD_REQUEST,
                content={"detail": [{"msg": str(e)}]},
            )
        except ValidationError as e:
            return JSONResponse(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                content={"detail": e.errors()},
            )
        except Exception as e:
            log.exception(f"Unexpected error: {str(e)}")
            return JSONResponse(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                content={"detail": [{"msg": "Internal server error"}]},
            )


--------------------------------------------------------------------------------

Chunk ID: main.py::3
def create_app(task_queue: TaskQueue, engine) -> FastAPI:
    exception_handlers = {404: not_found}
    app = FastAPI(exception_handlers=exception_handlers, openapi_url="/docs/openapi.json")

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    class DBMiddleware(BaseHTTPMiddleware):
        async def dispatch(self, request: Request, call_next):
            try:
                request.state.session_id = str(uuid.uuid4())
                task_auth_token = request.headers.get("x-task-auth", None)
                if not task_auth_token or not task_auth_token in token_registry:
                    session = sessionmaker(bind=engine)
                    request.state.db = session()
                    request.state.db.id = str(uuid.uuid4())

                response = await call_next(request)
            except Exception as e:
                raise e from None
            finally:
                db = getattr(request.state, "db", None)
                if db:
                    db.close()
            return response
    # ... other code


--------------------------------------------------------------------------------

Chunk ID: main.py::4
def create_app(task_queue: TaskQueue, engine) -> FastAPI:
    # ... other code

    class AddTaskQueueMiddleware(BaseHTTPMiddleware):
        async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
        ) -> Response:
            request.state.task_queue = task_queue
            response = await call_next(request)
            return response

    app.add_middleware(ExceptionMiddleware)
    app.add_middleware(DBMiddleware)
    app.add_middleware(AddTaskQueueMiddleware)

    app.include_router(auth_router)
    app.include_router(repo_router)
    app.include_router(tm_router)
    app.include_router(task_queue_router)
    app.include_router(test_gen_router)
    app.include_router(tgtcode_router)
    app.include_router(exp_router)
    app.include_router(health_router)

    return app


--------------------------------------------------------------------------------

Chunk ID: main.py::5
def create_server(app: FastAPI, host, port, workers = 2):
    # init_sentry()

    # start the repo sync thread
    # Session = sessionmaker(bind=engine)
    # db_session = Session()
    # start_sync_thread(db_session, task_queue)

    class CustomFilter(Filter):
        def filter(self, record):
            return "GET /task/get" not in record.getMessage()

    # Move the logger configuration into a function that will be called for each worker
    async def configure_worker_logger():
        uvicorn_logger = getLogger("uvicorn.access")
        uvicorn_logger.addFilter(CustomFilter())

    app = create_app(TaskQueue(), engine)

    # Configure the logging when the worker starts
    config = uvicorn.Config(
        app,
        host=host, 
        port=port,
        workers=workers,
        callback_notify=configure_worker_logger # This will run for each worker
    )
    server = uvicorn.Server(config)
    server.run()


if __name__ == "__main__":
    app = create_app(TaskQueue(), engine)
    create_server(app, "0.0.0.0", int(PORT))


--------------------------------------------------------------------------------

Chunk ID: ast/models.py::1
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean

from cowboy_lib.ast.code import ASTNode
from cowboy_lib.repo.source_repo import SourceRepo
from src.database.core import Base


class NodeModel(Base):
    """
    An AST node, either a class or a function, that holds a single or a group of unit tests
    """

    __tablename__ = "nodes"
    id = Column(Integer, primary_key=True)
    name = Column(String)
    node_type = Column(String)
    testfilepath = Column(String)

    # start_line = Column(Integer)
    # end_line = Column(Integer)
    # lines = Column(String)
    # is_test = Column(Boolean)
    # scope_id = Column(Integer, ForeignKey("nodes.id"))
    # children = relationship(
    #     "NodeModel",
    #     backref=backref("parent", remote_side=[id]),
    #     cascade="all, delete-orphan",
    # )

    # # decorators = Column

    repo_id = Column(Integer, ForeignKey("repo_config.id", ondelete="CASCADE"))
    # Node is either related to test_modules or target_code
    test_module_id = Column(
        Integer,
        ForeignKey("test_modules.id", ondelete="CASCADE"),
        nullable=True,
    )
    target_code_id = Column(
        Integer,
        ForeignKey("target_code.id", ondelete="CASCADE"),
        nullable=True,
    )

    # chunks = relationship("TargetCodeModel", backref="nodes")

    def to_astnode(self, source_repo: SourceRepo):
        node = source_repo.find_node(self.name, self.testfilepath, self.node_type)
        return node

    @classmethod
    def from_astnode(cls, node: ASTNode) -> "NodeModel":
        return cls(
            name=node.name,
            node_type=node.node_type,
        )


--------------------------------------------------------------------------------

Chunk ID: ast/service.py::1
from cowboy_lib.ast.code import ASTNode
from cowboy_lib.test_modules.test_module import TestModule

from src.test_modules.models import TestModuleModel

from sqlalchemy.orm import Session
from .models import NodeModel


def get_node(
    *, db_session: Session, node_name: str, repo_id: int, node_type: str, filepath: str
):
    return (
        db_session.query(NodeModel)
        .filter(
            NodeModel.name == node_name,
            NodeModel.repo_id == repo_id,
            NodeModel.node_type == node_type,
            NodeModel.testfilepath == filepath,
        )
        .one_or_none()
    )


def create_node(
    *,
    db_session: Session,
    node: ASTNode,
    repo_id: int,
    filepath: str,
    test_module_id: str = None,
):
    node = NodeModel(
        name=node.name,
        node_type=node.node_type.value,
        repo_id=repo_id,
        test_module_id=test_module_id,
        testfilepath=filepath,
    )

    db_session.add(node)
    db_session.commit()

    return node


--------------------------------------------------------------------------------

Chunk ID: ast/service.py::2
def create_or_update_node(
    *, db_session: Session, repo_id: str, node: ASTNode, filepath: str
):
    old_node = get_node(
        db_session=db_session,
        node_name=node.name,
        repo_id=repo_id,
        node_type=node.node_type.value,
        filepath=filepath,
    )

    if old_node:
        # NOTE: there is actually no point in updating node right now
        # because none of the node attributes should change ..
        print("Node exists: ", node.name)
        # node_model = (
        #     db_session.query(NodeModel)
        #     .filter(
        #         NodeModel.name == node.name
        #         and NodeModel.repo_id == repo_id
        #         and NodeModel.node_type == node.node_type
        #         and NodeModel.testfilepath == filepath
        #     )
        #     .update(node)
        # )
        return old_node
    else:
        node_model = create_node(
            db_session=db_session, node=node, repo_id=repo_id, filepath=filepath
        )

    return node_model


--------------------------------------------------------------------------------

Chunk ID: auth/models.py::1
from src.config import COWBOY_JWT_ALG, COWBOY_JWT_EXP, COWBOY_JWT_SECRET
from src.database.core import Base
from src.models import TimeStampMixin, CowboyBase, PrimaryKey

import string
import secrets
import bcrypt
from jose import jwt
from typing import Optional
from pydantic import validator, Field, BaseModel
from pydantic.networks import EmailStr
from sqlalchemy import (
    ForeignKey,
    DateTime,
    Column,
    String,
    LargeBinary,
    Integer,
    Boolean,
)
from sqlalchemy.orm import relationship
from typing import List
from datetime import datetime, timedelta


def generate_token(email):
    now = datetime.utcnow()
    exp = (now + timedelta(seconds=COWBOY_JWT_EXP)).timestamp()
    data = {
        "exp": exp,
        "email": email,
    }
    return jwt.encode(data, COWBOY_JWT_SECRET, algorithm=COWBOY_JWT_ALG)


def generate_password():
    """Generates a reasonable password if none is provided."""
    alphanumeric = string.ascii_letters + string.digits
    while True:
        password = "".join(secrets.choice(alphanumeric) for i in range(10))
        if (
            any(c.islower() for c in password)
            and any(c.isupper() for c in password)  # noqa
            and sum(c.isdigit() for c in password) >= 3  # noqa
        ):
            break
    return password


def hash_password(password: str):
    """Generates a hashed version of the provided password."""
    pw = bytes(password, "utf-8")
    salt = bcrypt.gensalt()
    return bcrypt.hashpw(pw, salt)


--------------------------------------------------------------------------------

Chunk ID: auth/models.py::2
class CowboyUser(Base, TimeStampMixin):
    __tablename__ = "cowboy_user"

    id = Column(Integer, primary_key=True)
    email = Column(String, unique=True)
    password = Column(LargeBinary, nullable=False)
    last_mfa_time = Column(DateTime, nullable=True)
    experimental_features = Column(Boolean, default=False)

    repos = relationship(
        "RepoConfig", backref="cowboy_user", cascade="all, delete-orphan"
    )

    # search_vector = Column(
    #     TSVectorType("email", regconfig="pg_catalog.simple", weights={"email": "A"})
    # )

    def check_password(self, password):
        return bcrypt.checkpw(password.encode("utf-8"), self.password)

    @property
    def token(self):
        return generate_token(self.email)


--------------------------------------------------------------------------------

Chunk ID: auth/models.py::3
class UserBase(CowboyBase):
    email: EmailStr

    @validator("email")
    def email_required(cls, v):
        if not v:
            raise ValueError("Must not be empty string and must be a email")
        return v


class UserLogin(UserBase):
    password: str

    @validator("password")
    def password_required(cls, v):
        if not v:
            raise ValueError("Must not be empty string")
        return v


class UserRegister(UserLogin):
    openai_api_key: str
    password: Optional[str] = Field(None, nullable=True)

    @validator("password", pre=True, always=True)
    def password_required(cls, v):
        # we generate a password for those that don't have one
        password = v or generate_password()
        return hash_password(password)


class UserLoginResponse(CowboyBase):
    token: Optional[str] = Field(None, nullable=True)


class UserRead(UserBase):
    id: PrimaryKey
    role: Optional[str] = Field(None, nullable=True)
    experimental_features: Optional[bool]


class UserUpdate(CowboyBase):
    id: PrimaryKey
    password: Optional[str] = Field(None, nullable=True)

    @validator("password", pre=True)
    def hash(cls, v):
        return hash_password(str(v))


class UserCreate(CowboyBase):
    email: EmailStr
    password: Optional[str] = Field(None, nullable=True)

    @validator("password", pre=True)
    def hash(cls, v):
        return hash_password(str(v))


class UserRegisterResponse(CowboyBase):
    token: Optional[str] = Field(None, nullable=True)


class UpdateOAIKey(BaseModel):
    openai_api_key: str


--------------------------------------------------------------------------------

Chunk ID: auth/permissions.py::1
from starlette.requests import Request

import logging
from abc import ABC, abstractmethod

from fastapi import HTTPException
from starlette.requests import Request
from starlette.status import HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND

from .service import get_current_user


class BasePermission(ABC):
    """
    Abstract permission that all other Permissions must be inherited from.

    Defines basic error message, status & error codes.

    Upon initialization, calls abstract method  `has_required_permissions`
    which will be specific to concrete implementation of Permission class.

    You would write your permissions like this:

    .. code-block:: python

        class TeapotUserAgentPermission(BasePermission):

            def has_required_permissions(self, request: Request) -> bool:
                return request.headers.get('User-Agent') == "Teapot v1.0"

    """

    user_error_msg = [{"msg": "User not found"}]
    user_error_code = HTTP_404_NOT_FOUND

    role = None

    # @abstractmethod
    # def has_required_permissions(self, request: Request) -> bool: ...

    def __init__(self, request: Request):
        user = get_current_user(request=request)
        if not user:
            raise HTTPException(
                status_code=self.user_error_code, detail=self.user_error_msg
            )

        # if not self.has_required_permissions(request):
        #     raise HTTPException(
        #         status_code=self.user_role_error_code, detail=self.user_role_error_msg
        #     )
#     """


--------------------------------------------------------------------------------

Chunk ID: auth/permissions.py::2
class PermissionsDependency(object):
    """
    Permission dependency that is used to define and check all the permission
    classes from one place inside route definition.

    Use it as an argument to FastAPI's `Depends` as follows:

    .. code-block:: python

        app = FastAPI()

        @app.get(
            "/teapot/",
            dependencies=[Depends(
                PermissionsDependency([TeapotUserAgentPermission]))]
        )
        async def teapot() -> dict:
            return {"teapot": True}
    """

    def __init__(self, permissions_classes: list):
        self.permissions_classes = permissions_classes

    def __call__(self, request: Request):
        for permission_class in self.permissions_classes:
            permission_class(request=request)


--------------------------------------------------------------------------------

Chunk ID: auth/permissions.py::3
from starlette.requests import Request


# class InternalClientPermissions(BasePermission):
#     """
#     Permissions used by our internal client to get ahold of user JWT tokens
#     for testing and other purposes
#     """

#     def __init__(self, request: Request):
#         user = get_current_user(request=request)
#         if not user:
#             raise HTTPException(
#                 status_code=self.user_error_code, detail=self.user_error_msg
#             )

#         if not user.is_admin:
#             raise HTTPException(
#                 status_code=HTTP_403_FORBIDDEN, detail="User is not an admin"
#             )


--------------------------------------------------------------------------------

Chunk ID: auth/service.py::1
import logging

from typing import Optional

from fastapi import HTTPException
from fastapi.security.utils import get_authorization_scheme_param
from jose import JWTError, jwt
from jose.exceptions import JWKError
from starlette.requests import Request
from starlette.status import HTTP_401_UNAUTHORIZED

from src.config import COWBOY_JWT_SECRET
from src.database.core import DBNotSetException, get_db

from .models import generate_token


log = logging.getLogger(__name__)

from .sm import SecretManager
from .models import CowboyUser, UserRegister, UserCreate

InvalidCredentialException = HTTPException(
    status_code=HTTP_401_UNAUTHORIZED,
    detail=[{"msg": "Could not validate credentials"}],
)


def get(*, db_session, user_id: int) -> Optional[CowboyUser]:
    """Returns a user based on the given user id."""
    return db_session.query(CowboyUser).filter(CowboyUser.id == user_id).one_or_none()


def get_by_email(*, db_session, email: str) -> Optional[CowboyUser]:
    """Returns a user object based on user email."""
    return db_session.query(CowboyUser).filter(CowboyUser.email == email).one_or_none()


--------------------------------------------------------------------------------

Chunk ID: auth/service.py::2
def create(*, db_session, user_in: UserRegister | UserCreate) -> CowboyUser:
    """Creates a new dispatch user."""
    # pydantic forces a string password, but we really want bytes
    password = bytes(user_in.password, "utf-8")

    # create the user
    user = CowboyUser(
        **user_in.dict(exclude={"password", "openai_api_key"}),
        password=password,
    )

    # create the credentials
    store_oai_key(user_in.openai_api_key, user.id)

    db_session.add(user)
    db_session.commit()

    return user


def get_user_token(*, db_session, user_id):
    user = get(db_session=db_session, user_id=user_id)
    return generate_token(user.email)


--------------------------------------------------------------------------------

Chunk ID: auth/service.py::3
def extract_user_email_jwt(request: Request, **kwargs):
    try:
        authorization: str = request.headers.get("Authorization")
        scheme, param = get_authorization_scheme_param(authorization)
        if not authorization or scheme.lower() != "bearer":
            log.exception(
                f"Malformed authorization header. Scheme: {scheme} Param: {param} Authorization: {authorization}"
            )
            return

        token = authorization.split()[1]
        data = jwt.decode(token, COWBOY_JWT_SECRET)
    except (JWKError, JWTError, IndexError, KeyError):
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=[{"msg": "Could not validate credentials"}],
        ) from None
    return data["email"]


--------------------------------------------------------------------------------

Chunk ID: auth/service.py::4
# def get_current_user(request: Request) -> CowboyUser:
#     user_email = extract_user_email_jwt(request=request)

#     if not user_email:
#         log.exception(f"Failed to extract user email")
#         raise InvalidCredentialException

#     # kinda of strange ... if user exists, we generate a random password
#     # for the user here ...
#     return get_or_create(
#         db_session=request.state.db,
#         user_in=UserRegister(email=user_email),
#     )


def get_current_user(request: Request) -> CowboyUser:
    user_email = extract_user_email_jwt(request=request)

    if not user_email:
        log.exception(f"Failed to extract user email")
        raise InvalidCredentialException

    # kinda of strange ... if user exists, we generate a random password
    # for the user here ...
    try:
        user = get_by_email(
            db_session=get_db(request),
            email=user_email,
        )
    # this is special case for requests polling the /task/get endpoint
    # where we are not passed a db session, and we want to proceed with the rest
    # of endpoint logic
    except DBNotSetException:
        return None

    # generic case for user not existing
    if not user:
        print("No user")
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED, detail=[{"msg": "User not found"}]
        )

    return user


def store_oai_key(api_key, user_id):
    sm = SecretManager()
    sm.store_parameter("OAI_KEY_" + str(user_id), api_key)


def retrieve_oai_key(user_id):
    sm = SecretManager()
    return sm.retrieve_parameter("OAI_KEY_" + str(user_id))


--------------------------------------------------------------------------------

Chunk ID: auth/sm.py::1
import boto3
from botocore.exceptions import ClientError

from src.config import AWS_REGION


class SecretManager:
    def __init__(self, region_name=AWS_REGION):
        """Initialize the ParameterManager with a specific AWS region."""
        self.client = boto3.client("ssm", region_name=region_name)

    def store_parameter(self, name, value, description="", key_id=None):
        """
        Store a parameter in AWS Parameter Store.

        :param name: Name of the parameter.
        :param value: Value of the parameter.
        :param description: Description of the parameter.
        :param key_id: The ID of the KMS key to use for encryption. If None, the default KMS key is used.
        :return: Response from the put_parameter call.
        """
        try:
            params = {
                "Name": name,
                "Value": value,
                "Type": "SecureString" if key_id else "String",
                "Description": description,
                "Overwrite": True,
            }
            if key_id:
                params["KeyId"] = key_id

            response = self.client.put_parameter(**params)
            return response
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None


--------------------------------------------------------------------------------

Chunk ID: auth/sm.py::2
class SecretManager:

    def retrieve_parameter(self, name, with_decryption=True):
        """
        Retrieve a parameter from AWS Parameter Store.

        :param name: Name of the parameter.
        :param with_decryption: Whether to decrypt the parameter value.
        :return: The parameter value.
        """
        try:
            response = self.client.get_parameter(
                Name=name, WithDecryption=with_decryption
            )
            return response["Parameter"]["Value"]
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None


--------------------------------------------------------------------------------

Chunk ID: auth/sm.py::3
# Example usage:
if __name__ == "__main__":
    region = "us-east-2"
    param_manager = SecretManager(region_name=region)

    # Store a parameter
    param_name = "MySecret"
    param_value = "SuperSecretValue"
    response = param_manager.store_parameter(
        name=param_name, value=param_value, description="A test secret parameter"
    )
    if response:
        print(f"Parameter {param_name} stored successfully.")

    # Retrieve a parameter
    retrieved_value = param_manager.retrieve_parameter(name=param_name)
    if retrieved_value:
        print(f"Retrieved value: {retrieved_value}")


--------------------------------------------------------------------------------

Chunk ID: auth/views.py::1
from src.database.core import get_db
from src.auth.models import CowboyUser
from src.models import HTTPSuccess

from src.runner.service import RunServiceArgs, shutdown_client
from src.queue.core import get_queue, TaskQueue

from .service import get_current_user, get, get_by_email, create, store_oai_key
from .models import UserLoginResponse, UserRegister, UpdateOAIKey

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session

auth_router = APIRouter()


@auth_router.post("/user/register", response_model=UserLoginResponse)
def register_user(
    user_in: UserRegister,
    db_session: Session = Depends(get_db),
):
    user = get_by_email(db_session=db_session, email=user_in.email)
    if user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="A user with this email already exists.",
        )

    user = create(db_session=db_session, user_in=user_in)

    return user


--------------------------------------------------------------------------------

Chunk ID: auth/views.py::2
@auth_router.get("/user/delete")
async def delete_user(
    curr_user: CowboyUser = Depends(get_current_user),
    db_session: Session = Depends(get_db),
    task_queue: TaskQueue = Depends(get_queue),
):
    user = get(db_session=db_session, user_id=curr_user.id)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="User not found",
        )

    # resets the client to get it to sync with the deleted user
    args = RunServiceArgs(user_id=user.id, task_queue=task_queue)
    await shutdown_client(args)

    db_session.delete(user)
    db_session.commit()

    return HTTPSuccess()


--------------------------------------------------------------------------------

Chunk ID: auth/views.py::3
@auth_router.post("/user/update/openai-key")
def update_oai_key(
    request: UpdateOAIKey,
    curr_user: CowboyUser = Depends(get_current_user),
    db_session: Session = Depends(get_db),
):
    user = get(db_session=db_session, user_id=curr_user.id)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="User not found",
        )

    store_oai_key(user_id=user.id, api_key=request.openai_api_key)

    return HTTPSuccess()


--------------------------------------------------------------------------------

Chunk ID: aws/sm.py::1
import boto3
from botocore.exceptions import ClientError

from src.config import AWS_REGION


class SecretManager:
    def __init__(self, region_name=AWS_REGION):
        """Initialize the ParameterManager with a specific AWS region."""
        self.client = boto3.client("ssm", region_name=region_name)

    def store_parameter(self, name, value, description="", key_id=None):
        """
        Store a parameter in AWS Parameter Store.

        :param name: Name of the parameter.
        :param value: Value of the parameter.
        :param description: Description of the parameter.
        :param key_id: The ID of the KMS key to use for encryption. If None, the default KMS key is used.
        :return: Response from the put_parameter call.
        """
        try:
            params = {
                "Name": name,
                "Value": value,
                "Type": "SecureString" if key_id else "String",
                "Description": description,
                "Overwrite": True,
            }
            if key_id:
                params["KeyId"] = key_id

            response = self.client.put_parameter(**params)
            return response
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None


--------------------------------------------------------------------------------

Chunk ID: aws/sm.py::2
class SecretManager:

    def retrieve_parameter(self, name, with_decryption=True):
        """
        Retrieve a parameter from AWS Parameter Store.

        :param name: Name of the parameter.
        :param with_decryption: Whether to decrypt the parameter value.
        :return: The parameter value.
        """
        try:
            response = self.client.get_parameter(
                Name=name, WithDecryption=with_decryption
            )
            return response["Parameter"]["Value"]
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None


--------------------------------------------------------------------------------

Chunk ID: aws/sm.py::3
# Example usage:
if __name__ == "__main__":
    region = "us-east-2"
    param_manager = SecretManager(region_name=region)

    # Store a parameter
    param_name = "MySecret"
    param_value = "SuperSecretValue"
    response = param_manager.store_parameter(
        name=param_name, value=param_value, description="A test secret parameter"
    )
    if response:
        print(f"Parameter {param_name} stored successfully.")

    # Retrieve a parameter
    retrieved_value = param_manager.retrieve_parameter(name=param_name)
    if retrieved_value:
        print(f"Retrieved value: {retrieved_value}")


--------------------------------------------------------------------------------

Chunk ID: aws/sqs.py::1
import boto3
from botocore.exceptions import ClientError

from src.config import AWS_REGION


class SQS:
    def __init__(self, queue_name, region_name=AWS_REGION):
        self.sqs = boto3.client("sqs", region_name=region_name)
        self.queue_name = queue_name
        self.queue_url = self.create_fifo_queue(queue_name)

    def create_fifo_queue(self, queue_name):
        try:
            response = self.sqs.create_queue(
                QueueName=f"{queue_name}.fifo",
                Attributes={"FifoQueue": "true", "ContentBasedDeduplication": "true"},
            )
            return response["QueueUrl"]
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None

    def put(self, message_body, message_deduplication_id=None):
        try:
            response = self.sqs.send_message(
                QueueUrl=self.queue_url,
                MessageBody=message_body,
                MessageDeduplicationId=message_deduplication_id or message_body,
            )
            return response
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None

    def receive(self, max_number_of_messages=1, wait_time_seconds=0):
        try:
            response = self.sqs.receive_message(
                QueueUrl=self.queue_url,
                MaxNumberOfMessages=max_number_of_messages,
                WaitTimeSeconds=wait_time_seconds,
            )
            return response.get("Messages", [])
        except ClientError as e:
            print(f"An error occurred: {e}")
            return None

    def peek(self, max_number_of_messages=1):
        return self.receive(
            max_number_of_messages=max_number_of_messages, wait_time_seconds=0
        )


--------------------------------------------------------------------------------

Chunk ID: src/config.py::1
from starlette.config import Config
from enum import Enum

config = Config(".env")

ENV = config("ENV", default="dev")
PORT = config("PORT")
API_ENDPOINT = "http://18.223.150.134:" + PORT

# JWT settings
COWBOY_JWT_SECRET = config("DISPATCH_JWT_SECRET", default="")
COWBOY_JWT_ALG = config("DISPATCH_JWT_ALG", default="HS256")
COWBOY_JWT_EXP = config("DISPATCH_JWT_EXP", cast=int, default=308790000)  # Seconds

COWBOY_OPENAI_API_KEY = config("OPENAI_API_KEY")

DB_PASS = config("DB_PASS")
DB_NAME = config("DB_NAME")
DB_USER = config("DB_USER")
SQLALCHEMY_DATABASE_URI = (
    f"postgresql://{DB_USER}:{DB_PASS}@127.0.0.1:5432/{DB_NAME}"
)
SQLALCHEMY_ENGINE_POOL_SIZE = 50

ALEMBIC_INI_PATH = "."
ALEMBIC_CORE_REVISION_PATH = "alembic"

# LLM settings and test gen settings
AUGMENT_ROUNDS = 3 if ENV == "release" else 1
LLM_RETRIES = 3
AUTO_GEN_SIZE = 7
MAX_CTXT_SIZE = 10000

# Run test settings 
RUN_TEST = "app"
TESTCONFIG_ROOT = "tests/configs"

LOG_DIR = "log"
REPOS_ROOT = "repos"
AWS_REGION = "us-east-2"


class Language(str, Enum):
    """
    Currently supported languages
    """
    python = "python"


--------------------------------------------------------------------------------

Chunk ID: coverage/__init__.py::1
from .models import CoverageModel


--------------------------------------------------------------------------------

Chunk ID: coverage/models.py::1
from cowboy_lib.coverage import Coverage, TestCoverage

from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.orm import relationship
from src.database.core import Base
from typing import List


class CoverageModel(Base):
    __tablename__ = "coverage"

    id = Column(Integer, primary_key=True)
    filename = Column(String, nullable=False)
    covered_lines = Column(String, nullable=False)
    missing_lines = Column(String, nullable=False)

    stmts = Column(Integer, nullable=False)
    misses = Column(Integer, nullable=False)
    covered = Column(Integer, nullable=False)

    # test_coverage_id = Column(Integer, ForeignKey("test_coverage.id"))
    target_code_list = relationship(
        "TargetCodeModel", back_populates="coverage", cascade="all, delete-orphan"
    )
    repo_id = Column(Integer, ForeignKey("repo_config.id"))
    test_result_id = Column(Integer, ForeignKey("augment_test_results.id"))

    def deserialize(self) -> Coverage:
        return Coverage(
            filename=self.filename,
            covered_lines=[int(l) for l in self.covered_lines.split(",") if l],
            missing_lines=[int(l) for l in self.missing_lines.split(",") if l],
        )

def coverage_to_model(cov: Coverage) -> CoverageModel:
    return CoverageModel(
        filename=cov.filename,
        covered_lines=",".join(map(str, cov.covered_lines)),
        missing_lines=",".join(map(str, cov.missing_lines)), 
        stmts=cov.stmts,
        misses=cov.misses,
        covered=cov.covered
    )


--------------------------------------------------------------------------------

Chunk ID: coverage/service.py::1
from cowboy_lib.coverage import Coverage
from src.test_gen.models import AugmentTestResult

from .models import CoverageModel

from sqlalchemy.orm import Session
from typing import List


def upsert_coverage(
    *,
    db_session: Session,
    repo_id: int,
    cov_list: List[Coverage],
    test_result_id: int = None
):
    """Deletes old coverage and insert new"""

    db_session.query(CoverageModel).filter(CoverageModel.repo_id == repo_id).delete()
    for coverage in cov_list:
        # if it does not exist, create
        cov_model = CoverageModel(
            filename=coverage.filename,
            covered_lines=",".join(map(str, coverage.covered_lines)),
            missing_lines=",".join(map(str, coverage.missing_lines)),
            stmts=coverage.stmts,
            misses=coverage.misses,
            covered=coverage.covered,
            repo_id=repo_id,
            test_result_id=test_result_id,
        )
        db_session.add(cov_model)

    db_session.commit()
    return cov_model


--------------------------------------------------------------------------------

Chunk ID: coverage/service.py::2
def get_cov_by_filename(
    *, db_session: Session, repo_id: int, filename: str
) -> CoverageModel:
    """Get a coverage model by filename."""

    return (
        db_session.query(CoverageModel)
        .filter(CoverageModel.filename == filename, CoverageModel.repo_id == repo_id)
        .one_or_none()
    )


def get_cov_by_id(*, db_session: Session, repo_id: int, id: int) -> CoverageModel:
    """Get a coverage model by id."""

    return (
        db_session.query(CoverageModel)
        .filter(CoverageModel.id == id, CoverageModel.repo_id == repo_id)
        .one_or_none()
    )


--------------------------------------------------------------------------------

Chunk ID: coverage/stats.py::1
from src.coverage.models import CoverageModel
from cowboy_lib.coverage import Coverage

from src.test_modules.models import TestModuleModel, TestModule

from pathlib import Path
from typing import List


def get_coverage_stats(tm: TestModule, cov: Coverage):
    """
    Calculate stats on what % of coverage is our baseline'd TestModule covering
    """
    total_covered = cov.covered
    tgt_covered = 0
    missing = cov.misses

    for chunk in tm.chunks:
        if Path(chunk.filepath) == cov.filename:
            tgt_covered += len(chunk.lines)

    score = get_score(tgt_covered, total_covered, missing)
    return score


def get_score(tgt_covered, total_covered, missing):
    return tgt_covered + missing / total_covered if total_covered else 0


--------------------------------------------------------------------------------

Chunk ID: database/core.py::1
import functools
import re

from sqlalchemy import create_engine, inspect
from sqlalchemy.ext.declarative import declarative_base
from starlette.requests import Request

import src.config as config

engine = create_engine(
    config.SQLALCHEMY_DATABASE_URI,
    # pool_size=config.SQLALCHEMY_ENGINE_POOL_SIZE,
    # max_overflow=config.DATABASE_ENGINE_MAX_OVERFLOW,
    # pool_pre_ping=config.DATABASE_ENGINE_POOL_PING,
)


def resolve_table_name(name):
    """Resolves table names to their mapped names."""
    names = re.split("(?=[A-Z])", name)
    return "_".join([x.lower() for x in names if x])


# nested level get() function
def resolve_attr(obj, attr, default=None):
    """Attempts to access attr via dotted notation, returns none if attr does not exist."""
    try:
        return functools.reduce(getattr, attr.split("."), obj)
    except AttributeError:
        return default


--------------------------------------------------------------------------------

Chunk ID: database/core.py::2
class CustomBase:
    __repr_attrs__ = []
    __repr_max_length__ = 15

    # @declared_attr
    # def __tablename__(self):
    #     return resolve_table_name(self.__name__)

    def dict(self):
        """Returns a dict representation of a model."""
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}

    def update(self, obj):
        """Updates a model with values from another model."""
        for key, value in obj.dict().items():
            if key in self.dict():
                setattr(self, key, value)

    @property
    def _id_str(self):
        ids = inspect(self).identity
        if ids:
            return "-".join([str(x) for x in ids]) if len(ids) > 1 else str(ids[0])
        else:
            return "None"


--------------------------------------------------------------------------------

Chunk ID: database/core.py::3
class CustomBase:

    @property
    def _repr_attrs_str(self):
        max_length = self.__repr_max_length__

        values = []
        single = len(self.__repr_attrs__) == 1
        for key in self.__repr_attrs__:
            if not hasattr(self, key):
                raise KeyError(
                    "{} has incorrect attribute '{}' in "
                    "__repr__attrs__".format(self.__class__, key)
                )
            value = getattr(self, key)
            wrap_in_quote = isinstance(value, str)

            value = str(value)
            if len(value) > max_length:
                value = value[:max_length] + "..."

            if wrap_in_quote:
                value = "'{}'".format(value)
            values.append(value if single else "{}:{}".format(key, value))

        return " ".join(values)

    def __repr__(self):
        # get id like '#123'
        id_str = ("#" + self._id_str) if self._id_str else ""
        # join class name, id and repr_attrs
        return "<{} {}{}>".format(
            self.__class__.__name__,
            id_str,
            " " + self._repr_attrs_str if self._repr_attrs_str else "",
        )


Base = declarative_base(cls=CustomBase)


class DBNotSetException(Exception):
    pass


def get_db(request: Request):
    try:
        return request.state.db
    except AttributeError:
        raise DBNotSetException("Database not set on request.")


# Triggers initial response field validation error
# DbSession = Annotated[Session, Depends(get_db)]


--------------------------------------------------------------------------------

Chunk ID: database/manage.py::1
import os
import logging

from alembic import command as alembic_command
from alembic.config import Config as AlembicConfig

from sqlalchemy import text
from sqlalchemy.schema import CreateSchema
from sqlalchemy_utils import create_database, database_exists

import src.config as config

from .core import Base, sessionmaker


log = logging.getLogger(__file__)


def version_schema(script_location: str):
    """Applies alembic versioning to schema."""

    # add it to alembic table
    alembic_cfg = AlembicConfig(config.ALEMBIC_INI_PATH)
    alembic_cfg.set_main_option("script_location", script_location)
    alembic_command.stamp(alembic_cfg, "head")


def get_core_tables():
    """Fetches tables that belong to the 'dispatch_core' schema."""
    core_tables = []
    for _, table in Base.metadata.tables.items():
        if table.schema == "dispatch_core":
            core_tables.append(table)
    return core_tables


def get_tenant_tables():
    """Fetches tables that belong to their own tenant tables."""
    tenant_tables = []
    for _, table in Base.metadata.tables.items():
        if not table.schema:
            tenant_tables.append(table)
    return tenant_tables


--------------------------------------------------------------------------------

Chunk ID: database/manage.py::2
def init_database(engine):
    """Initializes the database."""

    print(engine.__dict__)

    if not database_exists(str(config.SQLALCHEMY_DATABASE_URI)):
        create_database(str(config.SQLALCHEMY_DATABASE_URI))

    # schema_name = "dispatch_core"
    # if not engine.dialect.has_schema(engine, schema_name):
    #     with engine.connect() as connection:
    #         connection.execute(CreateSchema(schema_name))

    tables = get_core_tables()

    print(tables)

    # Base.metadata.create_all(engine, tables=tables)

    # version_schema(script_location=config.ALEMBIC_CORE_REVISION_PATH)
    # setup_fulltext_search(engine, tables)

    # # setup an required database functions
    # session = sessionmaker(bind=engine)
    # db_session = session()

    # # we create the default organization if it doesn't exist
    # organization = (
    #     db_session.query(Organization)
    #     .filter(Organization.name == "default")
    #     .one_or_none()
    # )
    # if not organization:
    #     print("Creating default organization...")
    #     organization = Organization(
    #         name="default",
    #         slug="default",
    #         default=True,
    #         description="Default Dispatch organization.",
    #     )

    #     db_session.add(organization)
    #     db_session.commit()

    # # we initialize the database schema
    # init_schema(engine=engine, organization=organization)

    # # we install all plugins
    # from dispatch.common.utils.cli import install_plugins
    # from dispatch.plugins.base import plugins

    # install_plugins()

    # for p in plugins.all():
    #     plugin = Plugin(
    #         title=p.title,
    #         slug=p.slug,
    #         type=p.type,
    #         version=p.version,


--------------------------------------------------------------------------------

Chunk ID: database/manage.py::3
# def init_schema(*, engine, organization: Organization):
#     """Initializes a new schema."""
#     schema_name = f"test_schema"

#     if not engine.dialect.has_schema(engine, schema_name):
#         with engine.connect() as connection:
#             connection.execute(CreateSchema(schema_name))

#     # set the schema for table creation
#     tables = get_tenant_tables()

#     schema_engine = engine.execution_options(
#         schema_translate_map={
#             None: schema_name,
#         }
#     )

#     Base.metadata.create_all(schema_engine, tables=tables)

#     # put schema under version control
#     version_schema(script_location=config.ALEMBIC_TENANT_REVISION_PATH)

#     with engine.connect() as connection:
#         # we need to map this for full text search as it uses sql literal strings
#         # and schema translate map does not apply
#         for t in tables:
#             t.schema = schema_name

#         setup_fulltext_search(connection, tables)

#     session = sessionmaker(bind=schema_engine)
#     db_session = session()

#     organization = db_session.merge(organization)
#     db_session.add(organization)
#     db_session.commit()
#     return organization


# def setup_fulltext_search(connection, tables):
#     """Syncs any required fulltext table triggers and functions."""
#     # parsing functions
#     function_path = os.path.join(
#         os.path.dirname(os.path.abspath(fulltext.__file__)), "expressions.sql"
#     )
#     connection.execute(text(open(function_path).read()))

#     for table in tables:
#         table_triggers = []
#         for column in table.columns:
#             if column.name.endswith("search_vector"):
#                 if hasattr(column.type, "columns"):
#                     table_triggers.append(
#                         {
#                             "conn": connection,
#                             "table": table,
#                             "tsvector_column": "search_vector",
#                             "indexed_columns": column.type.columns,
#                         }
#                     )
#                 else:
#                     log.warning(
#                         f"Column search_vector defined but no index columns found. Table: {table.name}"
#                     )

#         for trigger in table_triggers:
#             sync_trigger(**trigger)


--------------------------------------------------------------------------------

Chunk ID: eval/base.py::1
from typing import Dict

from cowboy_lib.repo import SourceRepo
from braintrust import Dataset, init_dataset

from dataclasses import dataclass

@dataclass
class TMRecord:
    name: str
    coverage: float


--------------------------------------------------------------------------------

Chunk ID: src/exceptions.py::1
from pydantic.errors import PydanticUserError


class CowboyRunTimeException(Exception):
    def __init__(self, message: str):
        self.message = message
        super().__init__(message)


--------------------------------------------------------------------------------

Chunk ID: experiments/augment_test.py::1
from cowboy_lib.repo import SourceRepo, GitRepo
from cowboy_lib.coverage import TestCoverage, Coverage
from cowboy_lib.test_modules.test_module import TestModule
from cowboy_lib.ast.code import NodeType
from cowboy_lib.repo.source_file import NodeNotFound, SameNodeException

from src.test_modules.models import TestModuleModel
from src.repo.models import RepoConfig

from pathlib import Path
from git import Repo
from logging import getLogger
from typing import List, Tuple


logger = getLogger("test_results")


class ExpVarDeleteFuncs(Exception):
    pass


# TODO: implement on TestCoverage
def find_file_cov_for_tm(tm: TestModule, base_cov: TestCoverage) -> Coverage:
    try:
        return next(
            filter(
                lambda x: x.filename == str(tm.targeted_files(base_path=False)[0]),
                base_cov.cov_list,
            )
        )
    except StopIteration:
        return None


--------------------------------------------------------------------------------

Chunk ID: experiments/augment_test.py::2
def num_funcs_to_delete(tm: TestModule, to_keep: int = 0, to_delete: int = 0) -> int:
    """
    Returns the number of functions to delete or keep from a test module
    """
    print("to_keep: ", to_keep, "to_delete:", to_delete)
    if to_keep and to_delete:
        raise Exception("Cannot have both values > 0")

    # always leave at least one test
    if to_keep:
        num_to_del = max(0, len(tm.tests) - to_keep)
        return num_to_del
    elif to_delete:
        num_to_del = min(len(tm.tests) - 1, to_delete)
        return num_to_del
    else:
        raise ExpVarDeleteFuncs


# for every experiment, we potentially create two new branches
# br1: keep2
# br2: keep2/mod1 or keep2/expid
def nuke_name_br(to_keep, to_delete):
    branch = f"{'keep' + str(to_keep) if to_keep else 'del' + str(to_delete)}"
    return branch


def exp_name_br(exp_id: str):
    return exp_id


--------------------------------------------------------------------------------

Chunk ID: experiments/augment_test.py::3
def create_nuked_branch(
    tms: List[TestModule],
    src_repo: SourceRepo,
    git_repo: GitRepo,
    branch_name: str,
    to_keep: int = 0,
    to_delete: int = 0,
):
    """
    Modifies all test files by adding or deleting a set number of tests, then commits
    the changes to
    """
    total_deleted = 0
    for tm in tms:
        # TODO: would need to change this loop to handle files
        logger.info(f"Generating augmented tests for {tm.name} in {tm.path}")

        to_exclude = []
        # BUG: tm.tests gets modified somehow
        num_to_del = num_funcs_to_delete(tm, to_keep=to_keep, to_delete=to_delete)
        total_tests = len(tm.tests)

        logger.info(f"Deleting {num_to_del}/{total_tests} tests from {tm.name}")

        for func in tm.tests[:num_to_del]:
            try:
                to_exclude.append((func, tm.test_file.path))
                # CARE: this operation has changes state of src_repo,
                # which is then propagated to strategy below
                src_repo.find_file(tm.path).delete(
                    func.name, node_type=NodeType.Function
                )
                # mirror the modifications made to src_repo
                # CARE: this only modifies TestFile, but not the nodes
                # TODO: figure out why the double delete here didnt work
                # tm.test_file.delete(func.name, node_type=NodeType.Function)
                src_repo.write_file(tm.path)
                total_deleted += 1
            except (NodeNotFound, SameNodeException) as e:
                print(e)
                continue

    git_repo.checkout(branch_name, new=True)
    git_repo.add_n_commit([str(tm.path) for tm in tms], f"Delete {total_deleted} tests")


--------------------------------------------------------------------------------

Chunk ID: experiments/augment_test.py::4
def run_experiment(
    repo: RepoConfig,
    test_modules: List[TestModuleModel],
    to_keep: int = 0,
    to_delete: int = 0,
):
    src_repo = SourceRepo(Path(repo.source_folder))
    git_repo = GitRepo(Path(repo.source_folder))
    tms = [tm.serialize(src_repo) for tm in test_modules]

    base_cov = repo.base_cov()
    nuked_branch = nuke_name_br(to_keep, to_delete)

    if not git_repo.branch_exists(nuked_branch):
        print("Creating new branch")
        create_nuked_branch(tms, src_repo, git_repo, nuked_branch, to_keep, to_delete)

    git_repo.checkout(nuked_branch)

    print(f"Finished modifying repo: {src_repo.repo_path}")

    # for tm in test_modules:
    #     # we need to update fs so we can collect the coverage on the deleted tests
    #     # to compare it
    #     del_res, *_ = self.repo_ctxt.runner.run_test(exclude_tests=to_exclude)
    #     del_cov = base_cov.coverage - del_res.coverage
    #     logger.info(f"Deleted coverage: {del_cov.total_cov.covered}")

    #     if del_cov.total_cov.covered <= 0:
    #         logger.info(
    #             f"Skipping {tm.name}, no coverage difference: {del_cov.total_cov.misses}"
    #         )
    #         zero_improvement_tests += f"{tm.name}\n"
    #         continue

    #     # TODO: we need to use base_path from RepoPath here
    #     # we may have just deleted the targeted coverage from abvoe
    #     tgt_file = tm.targeted_files()[0]
    #     tgt_filecov_before = del_cov.get_file_cov(
    #         tgt_file, self.repo_ctxt.repo_path
    #     )
    #     tgt_filecov_before.read_line_contents(self.repo_ctxt.repo_path)
    #     logger.info(
    #         f"Target file coverage: {tgt_filecov_before} : {tgt_filecov_before.filename}"
    #     )

    #     total_coverage += del_cov.total_cov.covered

    #     evaluator_inst: Evaluator = evaluator(
    #         self.repo_ctxt.runner,
    #         self.repo_ctxt.git_repo,
    #         self.repo_ctxt.src_repo,
    #     )
    #     strat_instance: BaseTestStrategy = strategy(
    #         tm,
    #         self.repo_ctxt,
    #         evaluator_inst,
    #         del_res,
    #         tgt_filecov_before,
    #     )
    #     improved, failed, no_improve = await strat_instance.generate_test(
    #         n_times=1
    #     )

    #     print("Improved: ")
    #     for imp in improved:
    #         print(imp)

    #     self.save_results(tm, del_cov, improved, failed, no_improve)
    #     num_improve = len(improved)
    #     num_failed = len(failed)
    #     num_noimprove = len(no_improve)

    #     logger.info(
    #         f"Results for TM: {tm.name} => Improve: {num_improve}, Failed: {num_failed}, NoImprove: {num_noimprove}"
    #     )

    #     total_improvement += sum(
    #         [test[1].total_cov.covered for test in improved]
    #     )

    # except LintException as e:
    #     logger.info(f"Linting error: {e} on {tm.path}")
    #     continue

    # except SameNodeException as e:
    #     logger.info(f"SameNodeException on class: {tm.name}")
    #     continue

    # except Exception as e:
    #     logger.info(
    #         f"Exception on class: {tm.name} in {tm.path}",
    #         exc_info=True,
    #     )

    # logger.info(f"Final Results: {results}")
    # logger.info(f"Total improvement: {total_improvement}/{total_coverage}")
    # logger.info(f"Zero improvement tests: {zero_improvement_tests}")
    # logger.info(f"Run complete for {self.repo_ctxt.exp_id} ")


--------------------------------------------------------------------------------

Chunk ID: experiments/models.py::1
from pydantic import BaseModel
from typing import List


class ExperimentRequest(BaseModel):
    repo_name: str
    tms: List[str]
    to_keep: int
    to_delete: int


--------------------------------------------------------------------------------

Chunk ID: experiments/views.py::1
from src.database.core import get_db
from src.auth.service import get_current_user, CowboyUser
from src.repo.service import get_experiment
from src.test_modules.service import get_tms_by_names

from .models import ExperimentRequest
from .augment_test import run_experiment

from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session


exp_router = APIRouter()


@exp_router.post("/experiment/create")
def create_experiment(
    exp_config: ExperimentRequest,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    repo = get_experiment(
        db_session=db_session, curr_user=current_user, repo_name=exp_config.repo_name
    )
    tm_models = get_tms_by_names(
        db_session=db_session, repo_id=repo.id, tm_names=exp_config.tms
    )

    run_experiment(
        repo=repo,
        test_modules=tm_models,
        to_keep=int(exp_config.to_keep),
        to_delete=int(exp_config.to_delete),
    )


--------------------------------------------------------------------------------

Chunk ID: extensions/__init__.py::1
from .sentry import init_sentry


--------------------------------------------------------------------------------

Chunk ID: extensions/sentry.py::1
import sentry_sdk


def init_sentry():
    sentry_sdk.init(
        dsn="https://0de14048c02d4d10a00dafb70966c33c@o4507195649294336.ingest.us.sentry.io/4507195650736128",
        # Set traces_sample_rate to 1.0 to capture 100%
        # of transactions for performance monitoring.
        traces_sample_rate=1.0,
        # Set profiles_sample_rate to 1.0 to profile 100%
        # of sampled transactions.
        # We recommend adjusting this value in production.
        profiles_sample_rate=1.0,
    )


--------------------------------------------------------------------------------

Chunk ID: health/views.py::1
from fastapi import APIRouter, Depends
from src.models import HTTPSuccess
from src.auth.service import get_current_user
from src.auth.models import CowboyUser


health_router = APIRouter()


@health_router.get("/health")
async def health():
    return HTTPSuccess()


--------------------------------------------------------------------------------

Chunk ID: src/logger.py::1
from src.config import LOG_DIR

import logging
import os
from datetime import datetime
import pytz
import sys

# class CowboyLogger:
#     def __init__(self, logger: logging.Logger):
#         self.logger = logging.getLogger("testgen_logger")
#         self.logger.setLevel(logging.INFO)
#         self.logger.addHandler(get_file_handler(file_prefix="testgen"))
#         self.logger.addHandler(get_console_handler())

#     def _parse_args(self, args):
#         return ", ".join([str(a) for a in args])

#     def info(self, *args):
#         self.logger.info(self._parse_args(args))

#     def warning(self, *args):
#         self.logger.warning(self._parse_args(args))

#     def error(self, *args):
#         self.logger.error(self._parse_args(args))

#     def debug(self, *args):
#         self.logger.debug(self._parse_args(args))

#     def critical(self, *args):
#         self.logger.critical(self._parse_args(args))


def converter(timestamp):
    dt = datetime.fromtimestamp(timestamp, tz=pytz.utc)
    return dt.astimezone(pytz.timezone("US/Eastern")).timetuple()


formatter = logging.Formatter(
    "%(asctime)s - %(name)s:%(levelname)s: %(filename)s:%(lineno)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
formatter.converter = converter


def get_file_handler(log_dir=LOG_DIR, file_prefix: str = ""):
    """
    Returns a file handler for logging.
    """
    os.makedirs(log_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y-%m-%d")
    file_name = f"{file_prefix}_{timestamp}.log"
    file_handler = logging.FileHandler(os.path.join(log_dir, file_name))
    file_handler.setFormatter(formatter)
    return file_handler


def get_console_handler():
    """
    Returns a console handler for logging.
    """
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    return console_handler

testgen_logger = logging.getLogger("testgen_logger")
testgen_logger.setLevel(logging.INFO)
testgen_logger.addHandler(get_file_handler(file_prefix="testgen"))
# testgen_logger.addHandler(get_console_handler())

buildtm_logger = logging.getLogger("buildtm_logger")
buildtm_logger.setLevel(logging.INFO)
buildtm_logger.addHandler(get_file_handler(file_prefix="buildtm"))
buildtm_logger.addHandler(get_console_handler())

sync_repo = logging.getLogger("sync_repo")
sync_repo.setLevel(logging.INFO)
sync_repo.addHandler(get_file_handler(file_prefix="syncrepo"))
# sync_repo.addHandler(get_console_handler())

loggers = [testgen_logger, sync_repo, buildtm_logger]


def set_log_level(level=logging.INFO):
    """
    Sets the logging level for all defined loggers.
    """
    for logger in loggers:
        logger.setLevel(level)
        for handler in logger.handlers:
            handler.setLevel(level)


def configure_uvicorn_logger():
    uvicorn_error_logger = logging.getLogger("uvicorn.error")
    uvicorn_error_logger.addHandler(get_file_handler())
    uvicorn_error_logger.addHandler(get_console_handler())


# LOGFIRE METRICS
# accepted_count = logfire.metric_counter("accepted_tests", unit="1")
# failed_count = logfire.metric_counter("failed_tests", unit="1")
# total_count = logfire.metric_counter("total_tests", unit="1")


--------------------------------------------------------------------------------

Chunk ID: src/models.py::1
from datetime import datetime
from sqlalchemy import Column, DateTime, event
from pydantic import BaseModel, Field
from pydantic.types import SecretStr

from typing import Annotated

PrimaryKey = Annotated[int, Field(gt=0, lt=2147483647)]
NameStr = Annotated[
    str, Field(pattern=r"^(?!\s*$).+", strip_whitespace=True, min_length=3)
]


class TimeStampMixin(object):
    """Timestamping mixin"""

    created_at = Column(DateTime, default=datetime.utcnow)
    created_at._creation_order = 9998
    updated_at = Column(DateTime, default=datetime.utcnow)
    updated_at._creation_order = 9998

    @staticmethod
    def _updated_at(mapper, connection, target):
        target.updated_at = datetime.utcnow()

    @classmethod
    def __declare_last__(cls):
        event.listen(cls, "before_update", cls._updated_at)


class CowboyBase(BaseModel):
    class Config:
        from_attributes = True
        validate_assignment = True
        arbitrary_types_allowed = True
        str_strip_whitespace = True

        json_encoders = {
            # custom output conversion for datetime
            datetime: lambda v: v.strftime("%Y-%m-%dT%H:%M:%SZ") if v else None,
            SecretStr: lambda v: v.get_secret_value() if v else None,
        }


class HTTPSuccess(BaseModel):
    msg: str = "Success"


--------------------------------------------------------------------------------

Chunk ID: queue/core.py::1
from cowboy_lib.api.runner.shared import Task, TaskStatus

from fastapi import Request
from threading import Lock
from collections import defaultdict
from typing import List, Dict
from asyncio import Event, wait_for


class TaskEvent:
    def __init__(self, task: Task):
        self.event = Event()
        self.task = task
        self.result = None

    async def wait(self, timeout: float = None):
        try:
            if timeout:
                await wait_for(self.event.wait(), timeout)
            else:
                await self.event.wait()
        except TimeoutError:
            return None

        return self.result

    def complete(self, result):
        """
        Complete with result and signal event to wake up
        """
        self.result = result
        self.event.set()

    @property
    def task_id(self):
        return self.task.task_id

    def __eq__(self, other):
        return self.task_id == other.task_id

    # def __hash__(self):
    #     return sum([ord(c) for c in self.task_id])


--------------------------------------------------------------------------------

Chunk ID: queue/core.py::2
class TaskQueue:
    """
    A set of queues separated by user_id
    """

    _instance = None

    def __new__(cls, *args, **kwargs):
        if not isinstance(cls._instance, cls):
            print("Creating new TaskQueue instance")
            cls._instance = super(TaskQueue, cls).__new__(cls, *args, **kwargs)
            cls._instance._initialized = False

        return cls._instance

    def __init__(self):
        if not self._initialized:
            # Initialize instance variables only once
            self.queue: Dict[str, List[TaskEvent]] = defaultdict(list)
            self.locks = defaultdict(list)
            self._initialized = True  # Mark as initialized

    def _acquire_lock(self, user_id: int):
        if self.locks.get(user_id, None) is None:
            self.locks[user_id] = Lock()
        return self.locks.get(user_id)

    def put(self, user_id: int, task: str) -> TaskEvent:
        with self._acquire_lock(user_id):
            t = TaskEvent(task)
            self.queue[user_id].append(t)

            return t

    def complete(self, user_id: int, task_id: str, res):
        with self._acquire_lock(user_id):
            for i in range(len(self.queue[user_id])):
                if self.queue[user_id][i].task_id == task_id:
                    t = self.queue[user_id].pop(i)
                    t.complete(res)
                    break

    # def get(self, user_id: int) -> Task:
    #     """
    #     Returns the first PENDING task and changes its status to STARTED
    #     """
    #     with self._acquire_lock(user_id):
    #         if len(self.queue[user_id]) == 0:
    #             return None

    #         return self.queue[user_id].pop()

    def get_all(self, user_id: int) -> List[Task]:
        with self._acquire_lock(user_id):
            if len(self.queue[user_id]) == 0:
                return []

            tasks = []
            for t in filter(
                lambda t: t.task.status == TaskStatus.PENDING.value, self.queue[user_id]
            ):
                t.task.status = TaskStatus.STARTED.value
                tasks.append(t.task)

            return tasks

    def peak(self, user_id: int, n: int) -> List[Task]:
        """
        Get the first n tasks in queue without removing
        """
        with self._acquire_lock(user_id):
            if len(self.queue[user_id]) == 0:
                return []

            return [t.task for t in self.queue[user_id][:n]]


def get_queue(request: Request):
    return request.state.task_queue


def get_token_registry(request: Request):
    from src.token_registry import token_registry

    return token_registry


def get_token(request: Request):
    """
    Returns the user id
    """
    token = request.headers.get("x-task-auth", None)
    # need this or else we end up converting None to "None" **shakes fist @ python moment"
    return str(token) if token else None


--------------------------------------------------------------------------------

Chunk ID: queue/models.py::1
from cowboy_lib.api.runner.shared import Task


class CompleteTaskRequest(Task):
    pass


class GetTaskResponse(Task):
    pass


--------------------------------------------------------------------------------

Chunk ID: queue/permissions.py::1
from src.auth.permissions import BasePermission
from src.auth.service import get_current_user

from fastapi import HTTPException

from starlette.requests import Request
from starlette.responses import Response


class TaskGetPermissions(BasePermission):
    def __init__(self, request: Request):
        try:
            user = get_current_user(request=request)
            if not user:
                raise HTTPException(
                    status_code=self.user_error_code, detail=self.user_error_msg
                )

        # this happens when the db is not set
        except AttributeError:
            pass


--------------------------------------------------------------------------------

Chunk ID: queue/service.py::1
from cowboy_lib.api.runner.shared import Task

from typing import Optional, List, Dict

from src.queue.core import TaskQueue


def list_tasks(*, task_queue: TaskQueue, user_id: int, n: int) -> Optional[List[Task]]:
    """List all tasks in the queue."""

    return task_queue.peak(user_id, n)


def dequeue_task(*, task_queue: TaskQueue, user_id: int) -> Optional[List[Task]]:
    """Dequeue the first task in the queue: retrieve and delete it."""

    return task_queue.get_all(user_id)


def complete_task(
    *, task_queue: TaskQueue, user_id: int, task_id: str, result: Dict
) -> None:
    """Mark a task as completed."""
    task_queue.complete(user_id, task_id, result)


def enqueue_task_and_wait(*, task_queue: TaskQueue, task: Task, user_id: int):
    """Enqueue a task to the specified queue."""

    f = task_queue.put(user_id, task)
    return f


--------------------------------------------------------------------------------

Chunk ID: queue/views.py::1
from cowboy_lib.api.runner.shared import Task

from .service import list_tasks, dequeue_task, complete_task
from .models import CompleteTaskRequest
from .core import TaskQueue, get_queue, get_token_registry, get_token

from fastapi import APIRouter, Depends, HTTPException, Response

from src.database.core import get_db
from src.auth.service import get_current_user
from src.auth.models import CowboyUser

from typing import List, Set

task_queue_router = APIRouter()


@task_queue_router.get("/task/list", response_model=List[Task])
def list(
    task_queue: TaskQueue = Depends(get_queue),
    curr_user: CowboyUser = Depends(get_current_user),
):
    tasks = list_tasks(task_queue=task_queue, user_id=curr_user.id, n=3)
    return tasks


--------------------------------------------------------------------------------

Chunk ID: queue/views.py::2
# incredibly hacky, basically, to prevent db connections from being used up
# we exclude db connections for this endpoint, we do the following:
# 1. First request actually does get a db sess, which we use to auth the user
# 2. Grab user id and add it into a in-mem token_registry list
# 3. Return user id as "set-x-task-auth" header
# 4. When the client puts user id into x-task-auth header
# 5. Our DBMiddleware will check the header, and if token is in registry, will not
# add a db session to the request
@task_queue_router.get("/task/get", response_model=List[Task])
def get(
    response: Response,
    task_queue: TaskQueue = Depends(get_queue),
    # don't try to do anything with curr_user because most of the time
    # we only have user_token to work with
    curr_user: CowboyUser = Depends(get_current_user),
    token_registry: Set[str] = Depends(get_token_registry),
    user_token: str = Depends(get_token),
    # perms: str = Depends(PermissionsDependency([TaskGetPermissions])),
):
    # at this point we have passed db user auth; test
    # catches if user sets random token
    if user_token and user_token not in token_registry:
        raise HTTPException(
            status_code=401,
            detail="Token not in registry, cannot proceed. \
            Are you sure you are logged in on the client?",
        )
    # issue token if it does not exist
    elif not user_token:
        print("Setting new token ..")
        response.headers["set-x-task-auth"] = str(curr_user.id)
        token_registry.add(str(curr_user.id))

    tasks = dequeue_task(
        task_queue=task_queue, user_id=curr_user.id if curr_user else int(user_token)
    )
    return tasks


@task_queue_router.post("/task/complete", response_model=CompleteTaskRequest)
def complete(
    task: CompleteTaskRequest,
    task_queue: TaskQueue = Depends(get_queue),
    curr_user: CowboyUser = Depends(get_current_user),
):

    task_queue = complete_task(
        task_queue=task_queue,
        user_id=curr_user.id,
        task_id=task.task_id,
        result=task.result,
    )
    return task


--------------------------------------------------------------------------------

Chunk ID: repo/models.py::1
from cowboy_lib.coverage import TestCoverage

from sqlalchemy import Column, Integer, String, JSON, ForeignKey, Boolean
from sqlalchemy.orm import relationship
from pydantic import Field

from src.models import CowboyBase
from src.database.core import Base
from src.config import Language

from typing import List, Any, Dict, Optional


class RepoConfig(Base):
    """
    Stores configuration for a repository
    """

    __tablename__ = "repo_config"

    id = Column(Integer, primary_key=True)
    repo_name = Column(String)
    url = Column(String)
    source_folder = Column(String)
    cloned_folders = Column(String)
    # git remote and git main branch (to merge into)
    remote = Column(String)
    main = Column(String)
    language = Column(String)

    # keep this argument fluid, may change
    python_conf = Column(JSON)
    user_id = Column(Integer, ForeignKey("cowboy_user.id"))
    is_experiment = Column(Boolean)

    # relations
    test_modules = relationship(
        "TestModuleModel", backref="repo_config", cascade="all, delete-orphan"
    )
    nodes = relationship(
        "NodeModel", backref="repo_config", cascade="all, delete-orphan"
    )
    cov_list = relationship(
        "CoverageModel", backref="repo_config", cascade="all, delete-orphan"
    )
    stats = relationship("RepoStats", uselist=False, cascade="all, delete-orphan")

    def __init__(
        self,
        repo_name,
        url,
        source_folder,
        cloned_folders,
        python_conf,
        user_id,
        remote,  # origin
        main,
        language,
        is_experiment=False,
    ):
        self.repo_name = repo_name
        self.url = url
        self.source_folder = source_folder
        self.cloned_folders = ",".join(cloned_folders)
        self.python_conf = python_conf
        self.user_id = user_id
        self.remote = remote
        self.main = main
        self.language = language
        self.is_experiment = is_experiment

    def to_dict(self):
        return {
            "repo_name": self.repo_name,
            "url": self.url,
            "source_folder": self.source_folder,
            "cloned_folders": self.cloned_folders.split(","),
            "python_conf": self.python_conf,
            "user_id": self.user_id,
            "remote": self.remote,
            "main": self.main,
            "language": self.language,
            "is_experiment": self.is_experiment,
        }

    @property
    def base_cov(self) -> TestCoverage:
        return TestCoverage([cov.deserialize() for cov in self.cov_list])


--------------------------------------------------------------------------------

Chunk ID: repo/models.py::2
class LangConf(CowboyBase):
    """
    Holds the language/framework specific settings
    for a repo
    """

    # currently I expect only an interpreter/compiler path that points
    # to the runtime for the targeted repo
    interp: str


class PythonConf(LangConf):
    language: str = "python"
    cov_folders: List[str]
    test_folder: str
    interp: str
    pythonpath: str

    def get(self, __name: str, default: Any = None) -> Any:
        return self.dict().get(__name, default)


class RepoConfigBase(CowboyBase):
    repo_name: str
    url: str
    source_folder: str
    cloned_folders: List[str]
    python_conf: PythonConf

    language: Optional[Language] = Field(default="python")
    is_experiment: Optional[bool] = Field(default=False)
    main: Optional[str] = Field(default="main")
    remote: Optional[str] = Field(default="origin")


class RepoConfigGet(RepoConfigBase):
    pass


class RepoConfigCreate(RepoConfigBase):
    repo_name: str


class RepoConfigList(CowboyBase):
    repo_list: List[RepoConfigBase]


class RepoConfigRemoteCommit(CowboyBase):
    sha: str


# class RepoConfigDelete(BaseModel):
#     repo_name: str


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::1
from cowboy_lib.repo import GitRepo, SourceRepo
from cowboy_lib.coverage import TestCoverage

from src.utils import gen_random_name
from src.auth.models import CowboyUser
from src.config import REPOS_ROOT
from src.runner.service import run_test, RunServiceArgs
from src.queue.core import TaskQueue
from src.coverage.service import upsert_coverage
from src.test_modules.service import create_all_tms
from src.exceptions import CowboyRunTimeException

from .models import RepoConfig, RepoConfigCreate
from ..test_modules.iter_tms import iter_test_modules

from pathlib import Path
from logging import getLogger
from fastapi import HTTPException


logger = getLogger(__name__)


def get(*, db_session, curr_user: CowboyUser, repo_name: str) -> RepoConfig:
    """Returns a repo based on the given repo name."""
    return (
        db_session.query(RepoConfig)
        .filter(RepoConfig.repo_name == repo_name, RepoConfig.user_id == curr_user.id)
        .one_or_none()
    )


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::2
def get_or_raise(*, db_session, curr_user: CowboyUser, repo_name: str) -> RepoConfig:
    """Returns a repo based on the given repo name."""
    repo = (
        db_session.query(RepoConfig)
        .filter(
            RepoConfig.repo_name == repo_name,
            RepoConfig.user_id == curr_user.id,
        )
        .one_or_none()
    )
    # TODO: consider raising pydantic Validation error here instead
    # seems to be what dispatch does
    if not repo:
        raise HTTPException(status_code=400, detail=f"Repo {repo_name} not found")

    return repo


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::3
def get_all(*, db_session) -> list[RepoConfig]:
    """Returns all repos."""
    return db_session.query(RepoConfig).all()


def get_by_id_or_raise(
    *, db_session, curr_user: CowboyUser, repo_id: int
) -> RepoConfig:
    """Returns a repo based on the given repo id."""
    repo = (
        db_session.query(RepoConfig)
        .filter(RepoConfig.id == repo_id, RepoConfig.user_id == curr_user.id)
        .one_or_none()
    )
    if not repo:
        raise HTTPException(status_code=400, detail=f"Repo {repo_id} not found")

    return repo


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::4
def get_experiment(*, db_session, curr_user: CowboyUser, repo_name: str) -> RepoConfig:
    """Returns a repo based on the given repo name."""

    return (
        db_session.query(RepoConfig)
        .filter(
            RepoConfig.repo_name == repo_name,
            RepoConfig.user_id == curr_user.id,
            RepoConfig.is_experiment == True,
        )
        .one_or_none()
    )


def delete(*, db_session, curr_user: CowboyUser, repo_name: str) -> RepoConfig:
    """Deletes a repo based on the given repo name."""

    repo = get(db_session=db_session, curr_user=curr_user, repo_name=repo_name)
    if repo:
        db_session.delete(repo)
        db_session.commit()

        GitRepo.delete_repo(Path(repo.source_folder))
        return repo

    return None


def clean(*, db_session, curr_user: CowboyUser, repo_name: str) -> RepoConfig:
    """Cleans repo branches."""

    repo = get(db_session=db_session, curr_user=curr_user, repo_name=repo_name)
    if repo:
        GitRepo.clean_branches(Path(repo.source_folder))
        return repo

    return None


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::5
async def create(
    *,
    db_session,
    curr_user: CowboyUser,
    repo_in: RepoConfigCreate,
    task_queue: TaskQueue,
) -> RepoConfig:
    """Creates a new repo."""

    repo_dst = None
    try:
        repo = RepoConfig(
            **repo_in.dict(),
            user_id=curr_user.id,
        )

        repo_dst = Path(REPOS_ROOT) / repo.repo_name / gen_random_name()
        GitRepo.clone_repo(repo_dst, repo.url)

        src_repo = SourceRepo(repo_dst)
        repo.source_folder = str(repo_dst)
        db_session.add(repo)
        # have to commit here or because run_test depends on existing RepoConfig
        db_session.commit()

        logger.info(f"Running coverage for repo {repo.repo_name}")

        # get base coverage for repo
        service_args = RunServiceArgs(user_id=curr_user.id, task_queue=task_queue)
        cov_res = await run_test(repo.repo_name, service_args)

        base_cov = TestCoverage(cov_list=cov_res.coverage.cov_list)
        if base_cov.total_cov.covered == 0:
            raise CowboyRunTimeException("No coverage found for repo, probably due to misconfigured test runner config")

        upsert_coverage(
            db_session=db_session,
            repo_id=repo.id,
            cov_list=cov_res.coverage.cov_list,
        )

        # create test modules
        create_all_tms(db_session=db_session, repo_conf=repo, src_repo=src_repo)

        db_session.commit()
        return repo

    except Exception as e:
        db_session.rollback()
        if repo:
            delete(db_session=db_session, curr_user=curr_user, repo_name=repo.repo_name)

        if repo_dst:
            GitRepo.delete_repo(repo_dst)

        logger.error(f"Failed to create repo configuration: {e}")
        raise


--------------------------------------------------------------------------------

Chunk ID: repo/service.py::6
def update(
    *, db_session, curr_user: CowboyUser, repo_name: int, repo_in: RepoConfigCreate
) -> RepoConfig:
    """Updates a repo."""

    repo = get(db_session=db_session, curr_user=curr_user, repo_name=repo_name)
    if not repo:
        return None

    repo.update(repo_in)
    db_session.commit()

    return repo


async def create_or_update(
    *,
    db_session,
    curr_user: CowboyUser,
    repo_in: RepoConfigCreate,
    task_queue: TaskQueue,
) -> RepoConfig:
    """Create or update a repo"""
    repo_conf = get(
        db_session=db_session, curr_user=curr_user, repo_name=repo_in.repo_name
    )

    if not repo_conf:
        return await create(
            db_session=db_session,
            curr_user=curr_user,
            repo_in=repo_in,
            task_queue=task_queue,
        )

    return update(
        db_session=db_session,
        curr_user=curr_user,
        repo_name=repo_in.repo_name,
        repo_in=repo_in,
    )


def list(*, db_session, curr_user: CowboyUser) -> RepoConfig:
    """Lists all repos for a user."""

    return db_session.query(RepoConfig).filter(RepoConfig.user_id == curr_user.id).all()


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::1
from cowboy_lib.repo import GitRepo

from src.database.core import get_db
from src.models import HTTPSuccess
from src.auth.service import get_current_user, CowboyUser
from src.queue.core import get_queue, TaskQueue
from src.runner.service import RunServiceArgs, shutdown_client

from .service import create_or_update, get, delete, list, clean
from .models import (
    RepoConfigCreate,
    RepoConfigList,
    RepoConfigGet,
    RepoConfigRemoteCommit,
)

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from pathlib import Path


repo_router = APIRouter()


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::2
@repo_router.post("/repo/create", response_model=RepoConfigCreate)
async def create_repo(
    repo_in: RepoConfigCreate,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
    task_queue: TaskQueue = Depends(get_queue),
):
    repo = get(
        db_session=db_session, repo_name=repo_in.repo_name, curr_user=current_user
    )
    if repo:
        raise HTTPException(
            status_code=400, detail="A repo with this name already exists."
        )

    repo_config = await create_or_update(
        db_session=db_session,
        repo_in=repo_in,
        curr_user=current_user,
        task_queue=task_queue,
    )
    # need as_dict to convert cloned_folders to list
    return repo_config.to_dict()


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::3
@repo_router.delete("/repo/delete/{repo_name}", response_model=HTTPSuccess)
async def delete_repo(
    repo_name: str,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
    task_queue: TaskQueue = Depends(get_queue),
):
    deleted = delete(db_session=db_session, repo_name=repo_name, curr_user=current_user)
    if not deleted:
        raise HTTPException(
            status_code=400, detail="A repo with this name does not exists."
        )

    # need this to shut down the client after a repo is deleted, or else
    # it will use old cloned_folders to execute the runner
    args = RunServiceArgs(user_id=current_user.id, task_queue=task_queue)
    await shutdown_client(args)

    return HTTPSuccess()


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::4
@repo_router.delete("/repo/clean/{repo_name}", response_model=HTTPSuccess)
def clean_repo(
    repo_name: str,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    cleaned = clean(db_session=db_session, repo_name=repo_name, curr_user=current_user)

    if not cleaned:
        raise HTTPException(
            status_code=400, detail="A repo with this name does not exists."
        )
    return HTTPSuccess()


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::5
@repo_router.get("/repo/get/{repo_name}", response_model=RepoConfigGet)
def get_repo(
    repo_name: str,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    repo = get(db_session=db_session, repo_name=repo_name, curr_user=current_user)
    if not repo:
        raise HTTPException(
            status_code=400, detail="A repo with this name does not exists."
        )
    return repo.to_dict()


@repo_router.get("/repo/list", response_model=RepoConfigList)
def list_repos(
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    repos = list(db_session=db_session, curr_user=current_user)
    return RepoConfigList(repo_list=repos)


--------------------------------------------------------------------------------

Chunk ID: repo/views.py::6
# TODO: this should return HEAD of repo.source_folder rather than the remote repo
# once we finish our task refactor
@repo_router.get("/repo/get_head/{repo_name}", response_model=RepoConfigRemoteCommit)
def get_head(
    repo_name: str,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    repo = get(db_session=db_session, repo_name=repo_name, curr_user=current_user)
    if not repo:
        raise HTTPException(
            status_code=400, detail="A repo with this name does not exists."
        )

    git_repo = GitRepo(Path(repo.source_folder))

    # return RepoConfigRemoteCommit(sha=git_repo.local_commit)
    return RepoConfigRemoteCommit(sha=git_repo.remote_commit)


--------------------------------------------------------------------------------

Chunk ID: local/cache.py::1
import sqlite3
import hashlib
import json
import pickle
from functools import wraps
from pathlib import Path
from typing import Any, Optional

from cowboy_lib.coverage import CoverageResult
from cowboy_lib.repo.repository import PatchFile

CACHE_DIR = Path("cache")
CACHE_DIR.mkdir(exist_ok=True)
DB_PATH = CACHE_DIR / "test_cache.db"

def init_cache():
    """Initialize the SQLite cache database"""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS test_cache (
            hash TEXT PRIMARY KEY,
            result BLOB
        )
    """)
    conn.commit()
    conn.close()


--------------------------------------------------------------------------------

Chunk ID: local/cache.py::2
def compute_hash(repo_name: str, exclude_tests: list, include_tests: list, patch_file: Optional[PatchFile]) -> str:
    """Compute a hash of the input arguments"""
    hash_input = {
        "repo_name": repo_name,
        "exclude_tests": exclude_tests,
        "include_tests": include_tests,
        "patch": patch_file.patch if patch_file else None
    }
    hash_str = json.dumps(hash_input, sort_keys=True)
    return hashlib.sha256(hash_str.encode()).hexdigest()


--------------------------------------------------------------------------------

Chunk ID: local/cache.py::3
def read_cache(hash_key: str) -> Optional[CoverageResult]:
    """Read a result from the cache"""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT result FROM test_cache WHERE hash = ?", (hash_key,))
    row = cursor.fetchone()
    conn.close()

    if row:
        return pickle.loads(row[0])
    return None

def save_cache(hash_key: str, result: CoverageResult):
    """Save a result to the cache"""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute(
        "INSERT OR REPLACE INTO test_cache (hash, result) VALUES (?, ?)",
        (hash_key, pickle.dumps(result))
    )
    conn.commit()
    conn.close()


--------------------------------------------------------------------------------

Chunk ID: local/cache.py::4
def cache_test_run(func):
    """Decorator to cache test run results"""
    @wraps(func)
    async def wrapper(
        repo_name: str,
        service_args: Any,
        exclude_tests: list = [],
        include_tests: list = [],
        patch_file: Optional[PatchFile] = None,
    ) -> CoverageResult:
        init_cache()
        cache_key = compute_hash(repo_name, exclude_tests, include_tests, patch_file)

        # Try to get from cache
        print(f"Returning from cache: {(repo_name, exclude_tests, include_tests)}")
        cached_result = read_cache(cache_key)
        if cached_result is not None:
            return cached_result

        # Run the actual function if not in cache
        result = await func(repo_name, service_args, exclude_tests, include_tests, patch_file)
        if not cached_result:
            print(f"Saving to cache: {(repo_name, exclude_tests, include_tests)}")
            save_cache(cache_key, result)
        return result

    return wrapper


--------------------------------------------------------------------------------

Chunk ID: local/models.py::1
from cowboy_lib.repo.repository import PatchFileContext, GitRepo
from cowboy_lib.coverage import CoverageResult
from cowboy_lib.api.runner.shared import RunTestTaskArgs, FunctionArg

from pydantic import BaseModel, validator
from typing import List, Optional


class PythonConf(BaseModel):
    cov_folders: List[str]
    interp: str
    test_folder: Optional[str]
    pythonpath: Optional[str]

    # @validator("interp")
    # def validate_interp(cls, v):
    #     import os

    #     if not os.path.exists(v):
    #         raise ValueError(f"Interpreter path {v} does not exist")
    #     return v


--------------------------------------------------------------------------------

Chunk ID: local/models.py::2
class RepoConfig(BaseModel):
    repo_name: str  # of form owner_repo
    url: str
    cloned_folders: List[str]  # list of cloned folders used for parallelizing run_test
    source_folder: str  # ???
    # pytest specific confs (although they could be generally applicable)
    python_conf: "PythonConf"
    is_experiment: Optional[bool] = False

    @validator("url")
    def validate_url(cls, v):
        import re

        if not re.match(r"^https:\/\/github\.com\/[\w-]+\/[\w-]+(\.git)?$", v):
            raise ValueError(
                "URL must be a valid GitHub HTTPS URL and may end with .git"
            )
        # if v.endswith(".git"):
        #     raise ValueError("URL should not end with .git")
        if re.match(r"^git@github\.com:[\w-]+\/[\w-]+\.git$", v):
            raise ValueError("SSH URL format is not allowed")
        return v

    def __post_init__(self):
        if isinstance(self.python_conf, dict):
            self.python_conf = PythonConf(**self.python_conf)

    def serialize(self):
        return {
            "repo_name": self.repo_name,
            "url": self.url,
            "cloned_folders": self.cloned_folders,
            "source_folder": self.source_folder,
            "python_conf": self.python_conf.__dict__,
            "is_experiment": self.is_experiment,
        }


--------------------------------------------------------------------------------

Chunk ID: local/python.py::1
from cowboy_lib.repo.repository import PatchFileContext, GitRepo
from cowboy_lib.coverage import CoverageResult
from cowboy_lib.api.runner.shared import RunTestTaskArgs, FunctionArg

from .models import RepoConfig

import os
import subprocess
from typing import List, Tuple, NewType, Dict
import json
import hashlib
from pathlib import Path
import time
import sys
from contextlib import contextmanager
import queue
from logging import getLogger

log = getLogger(__name__)


COVERAGE_FILE = "coverage.json"
TestError = NewType("TestError", str)


class RunnerError(Exception):
    pass

class DiffFileCreation(Exception):
    pass

class CowboyClientError(Exception):
    pass

class TestSuiteError(Exception):
    pass

def hash_str(s: str) -> str:
    return hashlib.sha256(s.encode()).hexdigest()


def hash_file(filepath):
    """Compute SHA-256 hash of the specified file"""
    with open(filepath, "r", encoding="utf-8") as f:
        buf = f.read()

    return hash_str(buf)


--------------------------------------------------------------------------------

Chunk ID: local/python.py::2
def hash_coverage_inputs(directory: Path, cmd_str: str) -> str:
    """Compute SHA-256 for the curr dir and cmd_str"""
    hashes = []
    for f in directory.iterdir():
        if f.is_file() and f.name.endswith(".py"):
            file_hash = hash_file(f)
            hashes.append((str(f), file_hash))

    # Sort based on file path and then combine the file hashes
    hashes.sort()
    combined_hash = hashlib.sha256()
    for _, file_hash in hashes:
        combined_hash.update(file_hash.encode())

    combined_hash.update(cmd_str.encode())
    return combined_hash.hexdigest()


--------------------------------------------------------------------------------

Chunk ID: local/python.py::3
class LockedRepos:
    """
    A list of available repos for concurrent run_test invocations, managed as a FIFO queue
    """

    def __init__(self, cloned_folders: List[Path]):
        self.capacity = len(cloned_folders)
        self.queue = queue.Queue()

        for path in cloned_folders:
            if not path.exists():
                raise RunnerError(
                    f"Cloned folder: {str(path)} does not exist, something went wrong. Try deleting and re-creating the repo"
                )

            self.queue.put(GitRepo(path))

    @contextmanager
    def acquire_one(self) -> GitRepo:
        git_repo: GitRepo = self.queue.get()  # This will block if the queue is empty

        log.info(f"Acquired repo: {git_repo.repo_folder}")
        try:
            yield git_repo
        finally:
            self.release(git_repo)

            log.info(f"Released repo: {git_repo.repo_folder}")

    def release(self, git_repo: GitRepo):
        self.queue.put(git_repo)  # Return the repo back to the queue

    def __len__(self):
        return self.queue.qsize()


def get_exclude_path(
    func: FunctionArg,
    rel_fp: Path,
):
    """
    Converts a Function path
    """
    excl_name = (
        (func.name.split(".")[0] + "::" + func.name.split(".")[1])
        if func.is_meth
        else func.name
    )

    # need to do this on windows
    return str(rel_fp).replace("\\", "/") + "::" + excl_name


--------------------------------------------------------------------------------

Chunk ID: local/python.py::4
class PytestDiffRunner:
    """
    Executes the test suite
    """
    def __init__(
        self,
        # assume to be a test file for now
        repo_conf: RepoConfig,
        test_suite: str = "",
    ):
        self.test_folder = Path(repo_conf.python_conf.test_folder)
        self.interpreter = Path(repo_conf.python_conf.interp)
        self.python_path = Path(repo_conf.python_conf.pythonpath)

        # if not self.interpreter.exists():
        #     raise RunnerError(f"Runtime path {self.interpreter} does not exist")

        # missing = self.check_missing_deps(self.interpreter)
        # if missing:
        #     raise RunnerError(f"{missing} for {self.interpreter}")

        self.cov_folders = [Path(p) for p in repo_conf.python_conf.cov_folders]
        cloned_folders = [Path(p) for p in repo_conf.cloned_folders]
        self.test_repos = LockedRepos(cloned_folders)

        if len(self.test_repos) == 0:
            raise CowboyClientError("No cloned repos created, perhaps run init again?")

        self.test_suite = test_suite


--------------------------------------------------------------------------------

Chunk ID: local/python.py::5
class PytestDiffRunner:

    def check_missing_deps(self, interp) -> List[bool]:
        """
        Check if test depdencies are installed in the interpreter
        """
        not_installed = []
        deps = ["pytest-cov", "pytest"]
        try:
            for dep in deps:
                result = subprocess.run(
                    [interp, "-m", "pip", "show", dep],
                    capture_output=True,
                    text=True,
                )
                if result.returncode != 0:
                    not_installed.append(dep)

        except (subprocess.CalledProcessError, FileNotFoundError):
            not_installed.append(dep)

        return (
            f"The following deps are not installed in the env: {not_installed}\nPlease install these"
            if not_installed
            else ""
        )


--------------------------------------------------------------------------------

Chunk ID: local/python.py::6
class PytestDiffRunner:

    def verify_clone_dirs(self, cloned_dirs: List[Path]):
        """
        Verifies that the hash of all *.py files are the same for each cloned dir
        """
        import hashlib

        hashes = []
        for clone in cloned_dirs:
            f_buf = ""
            for py_file in clone.glob("test*.py"):
                with open(py_file, "r") as f:
                    f_buf += f.read()

            f_buf_hash = hashlib.md5(f_buf.encode()).hexdigest()
            hashes.append(f_buf_hash)

        if any(h != hashes[0] for h in hashes):
            raise CowboyClientError("Cloned directories are not the same")


--------------------------------------------------------------------------------

Chunk ID: local/python.py::7
class PytestDiffRunner:

    def set_test_repos(self, repo_paths: List[Path]):
        self.test_repos = LockedRepos(
            list(zip(repo_paths, [GitRepo(p) for p in repo_paths]))
        )

        self.verify_clone_dirs(repo_paths)

    def _get_exclude_tests_arg_str(
        self, excluded_tests: List[Tuple[FunctionArg, Path]], cloned_path: Path
    ):
        """
        Convert the excluded tests into Pytest deselect args
        """
        if not excluded_tests:
            return ""

        # PATCH: just concatenate the path with base folder
        tranf_paths = []
        for test, test_fp in excluded_tests:
            tranf_paths.append(get_exclude_path(test, test_fp))

        return "--deselect=" + " --deselect=".join(tranf_paths)

    def _get_include_tests_arg_str(self, include_tests: []):
        if not include_tests:
            return ""

        arg_str = ""
        AND = " and"
        for test in include_tests:
            arg_str += f"{test}{AND}"

        arg_str = arg_str[: -len(AND)]
        # return "-k " + '"' + arg_str + '"'
        return "-k " + arg_str


--------------------------------------------------------------------------------

Chunk ID: local/python.py::8
class PytestDiffRunner:

    def _construct_cmd(
        self, 
        repo_path: Path,  
        selected_tests: str, 
        deselected_tests: str,
        custom_cmd: str
    ):
        """
        Constructs the cmdstr for running pytest
        """
        cd_cmd = [
            "cd",
            str(repo_path),
        ]

        if custom_cmd:
            cmd = custom_cmd
        else:
            cmd = [
                str(self.interpreter),
                "-m",
                "pytest",
                str(self.test_folder),
                "--tb",
                "short",
                selected_tests,
                deselected_tests,
                # "-v",
                "--color",
                "no",
                f"--cov={'--cov='.join([str(folder) + ' ' for folder in self.cov_folders])}",
                "--cov-report",
                "json",
                # "--cov-report",
                # "term",
                "--continue-on-collection-errors",
                "--disable-warnings",
            ]

        print("CMDER:: ", " ".join(cd_cmd + ["&&"] + cmd))
        return " ".join(cd_cmd + ["&&"] + cmd)


--------------------------------------------------------------------------------

Chunk ID: local/python.py::9
class PytestDiffRunner:

    def _stream_and_capture_output(self, process: subprocess.Popen) -> Tuple[str, str]:
        """
        Streams subprocess output to terminal while also capturing it in variables.
        
        Args:
            process: A subprocess.Popen instance with stdout and stderr as PIPE
            
        Returns:
            Tuple containing the captured stdout and stderr as strings
        """
        stdout_chunks = []
        stderr_chunks = []

        # Get file objects for stdout and stderr
        stdout_fd = process.stdout.fileno()
        stderr_fd = process.stderr.fileno()

        # Make stdout and stderr non-blocking
        os.set_blocking(stdout_fd, False)
        os.set_blocking(stderr_fd, False)

        while True:
            # Check if process has finished
            retcode = process.poll()

            # Try reading from stdout
            try:
                stdout_chunk = os.read(stdout_fd, 1024).decode('utf-8')
                if stdout_chunk:
                    print(stdout_chunk, end='', flush=True)
                    stdout_chunks.append(stdout_chunk)
            except (OSError, BlockingIOError):
                pass

            # Try reading from stderr
            try:
                stderr_chunk = os.read(stderr_fd, 1024).decode('utf-8')
                if stderr_chunk:
                    print(stderr_chunk, end='', flush=True, file=sys.stderr)
                    stderr_chunks.append(stderr_chunk)
            except (OSError, BlockingIOError):
                pass

            # If process has finished and no more output, break
            if retcode is not None and not stdout_chunk and not stderr_chunk:
                break

            # Small sleep to prevent CPU spinning
            time.sleep(0.1)

        return ''.join(stdout_chunks), ''.join(stderr_chunks)


--------------------------------------------------------------------------------

Chunk ID: local/python.py::10
class PytestDiffRunner:

    def run_testsuite(self, 
                      args: RunTestTaskArgs, 
                      stream: bool = False,
                      custom_cmd: str = "") -> Tuple[CoverageResult, str, str]:
        with self.test_repos.acquire_one() as repo_inst:
            git_repo: GitRepo = repo_inst

            patch_file = args.patch_file
            if patch_file:
                patch_file.path = git_repo.repo_folder / patch_file.path
                log.info(f"Using patch file: {patch_file.path}")

            exclude_tests = args.exclude_tests
            include_tests = args.include_tests

            env = os.environ.copy()
            if self.python_path:
                env["PYTHONPATH"] = self.python_path

            exclude_tests = self._get_exclude_tests_arg_str(
                exclude_tests, git_repo.repo_folder
            )
            include_tests = self._get_include_tests_arg_str(include_tests)
            cmd_str = self._construct_cmd(
                git_repo.repo_folder, include_tests, exclude_tests, custom_cmd
            )

            log.info(f"Running with command: {cmd_str}")

            with PatchFileContext(git_repo, patch_file):
                proc = subprocess.Popen(
                    cmd_str,
                    # env=env,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    shell=True,
                    text=True,
                )
                if stream:
                    stdout, stderr = self._stream_and_capture_output(proc)
                else:
                    stdout, stderr = proc.communicate()

                # raise exception only no coverage
                try:
                    with open(git_repo.repo_folder / COVERAGE_FILE, "r") as f:
                        coverage_json = json.loads(f.read())
                        cov = CoverageResult(stdout, stderr, coverage_json)
                        if not cov.coverage:
                            print("No coverage!")
                            raise Exception()
                except FileNotFoundError:
                    raise TestSuiteError(stderr)

            log.info(f"GET COVEAGE: {cov.get_coverage()}")
        return (
            cov,
            stdout,
            stderr,
        )


--------------------------------------------------------------------------------

Chunk ID: local/run_test.py::1
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.api.runner.shared import RunTestTaskArgs
from cowboy_lib.coverage import CoverageResult
from cowboy_lib.ast.code import Function

from ..models import RunServiceArgs
from .python import PytestDiffRunner
from .models import RepoConfig
from .cache import cache_test_run

from src.config import TESTCONFIG_ROOT
from pathlib import Path
from typing import List, Tuple
import json


def get_repo_config(repo_name: str) -> RepoConfig:
    with open(Path(TESTCONFIG_ROOT) / f"{repo_name}.json", "r") as f:
        return RepoConfig(**json.load(f))


--------------------------------------------------------------------------------

Chunk ID: local/run_test.py::2
# @cache_test_run
async def run_test(
    repo_name: str,
    service_args: RunServiceArgs, # dont actually need this argument here, just for compatability
    exclude_tests: List[Tuple[Function, str]] = [],
    include_tests: List[str] = [],
    patch_file: PatchFile = None,
) -> CoverageResult:
    repo_config = get_repo_config(repo_name)
    args = RunTestTaskArgs(
        repo_name=repo_name,
        patch_file=patch_file,
        exclude_tests=[((f.name, f.is_meth()), path) for f, path in exclude_tests],
        include_tests=include_tests,
    )
    runner = PytestDiffRunner(repo_config)
    cov, _, _ = runner.run_testsuite(args)

    return cov


--------------------------------------------------------------------------------

Chunk ID: runner/models.py::1
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.coverage import CoverageResult, TestCoverage
from cowboy_lib.api.runner.shared import TaskResult

from fastapi import HTTPException
from pydantic import BaseModel, Field, validator
from dataclasses import dataclass

from src.queue.core import TaskQueue

class ClientRunnerException(HTTPException):
    def __init__(self, msg):
        # TODO: make pytest a variable for diff testing frameworks
        self.detail = f"Local pytest runner error: {msg}"
        self.status_code = 400


class RunnerExceptionResponse(BaseModel):
    exception: str


def json_to_coverage_result(res: TaskResult):
    if res.exception:
        raise ClientRunnerException(res.exception)

    cov_results = CoverageResult("", "", {})
    cov_results.coverage = TestCoverage.deserialize(res.coverage)
    cov_results.failed = res.failed

    return cov_results

@dataclass
class RunServiceArgs:
    user_id: int
    task_queue: TaskQueue
    # NEWTODO: not sure if we want to stick this here will probably move it somewhere
    # else
    custom_cmd: str


--------------------------------------------------------------------------------

Chunk ID: runner/service.py::1
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.coverage import CoverageResult
from cowboy_lib.ast.code import Function
from cowboy_lib.api.runner.shared import (
    Task,
    TaskType,
    RunTestTaskArgs,
    FunctionArg,
)
from src.queue.service import enqueue_task_and_wait
from src.queue.core import TaskQueue

from .models import RunServiceArgs, json_to_coverage_result

from typing import List, Tuple
from dataclasses import dataclass

async def run_test(
    repo_name: str,
    service_args: RunServiceArgs,
    exclude_tests: List[Tuple[Function, str]] = [],
    include_tests: List[str] = [],
    patch_file: PatchFile = None,
) -> CoverageResult:
    task = Task(
        type=TaskType.RUN_TEST,
        task_args=RunTestTaskArgs(
            repo_name=repo_name,
            patch_file=patch_file,
            exclude_tests=[((f.name, f.is_meth()), path) for f, path in exclude_tests],
            include_tests=include_tests,
        ),
    )

    future = enqueue_task_and_wait(
        task_queue=service_args.task_queue,
        user_id=service_args.user_id,
        task=task,
    )
    res = await future.wait()

    cov_res = json_to_coverage_result(res)
    return cov_res


async def shutdown_client(
    service_args: RunServiceArgs,
) -> CoverageResult:
    task = Task(type=TaskType.SHUTDOWN)

    future = enqueue_task_and_wait(
        task_queue=service_args.task_queue,
        user_id=service_args.user_id,
        task=task,
    )

    await future.wait(timeout=3)


--------------------------------------------------------------------------------

Chunk ID: src/runtest_conf.py::1
from src.config import RUN_TEST
if RUN_TEST == "app":
    from src.runner.service import run_test
elif RUN_TEST == "test":
    from src.runner.local.run_test import run_test


--------------------------------------------------------------------------------

Chunk ID: scripts/drop_db.py::1
import src.config as config
import click
from src.database.core import engine


@click.group()
def cowboy_database():
    pass


@cowboy_database.command("drop")
def drop_database():
    """Drops all data in database."""
    from sqlalchemy_utils import database_exists, drop_database

    if database_exists(str(config.SQLALCHEMY_DATABASE_URI)):
        drop_database(str(config.SQLALCHEMY_DATABASE_URI))


drop_database()


--------------------------------------------------------------------------------

Chunk ID: scripts/neuter_repo.py::1
from src.test_modules.iter_tms import iter_test_modules
from cowboy_lib.repo import SourceRepo
from cowboy_lib.test_modules import TestModule
from cowboy_lib.ast import NodeType

from typing import List

import sys
from pathlib import Path

def num_delete(tm: TestModule, to_keep: int = 1, to_delete: int = 1) -> int:
    if to_keep and to_delete:
        raise Exception("Cannot have both values > 0")

    # always leave at least one test
    if to_keep:
        num_to_del = max(0, len(tm.tests) - to_keep)
        return num_to_del
    elif to_delete:
        num_to_del = min(len(tm.tests) - 1, to_delete)
        return num_to_del
    else:
        raise Exception("Must provide either to_keep or to_delete value")


--------------------------------------------------------------------------------

Chunk ID: scripts/neuter_repo.py::2
def neuter_tests(
    test_modules: List[TestModule], src_repo: SourceRepo, to_keep, to_delete=0
):
    total_deleted = 0
    failed_mod = 0
    for tm in test_modules:
        try:
            print("Deleting tm: ", tm.name)
            to_exclude = []
            # BUG: tm.tests gets modified somehow
            num_to_del = num_delete(tm, to_keep=to_keep, to_delete=to_delete)
            total_tests = len(tm.tests)

            for func in tm.tests[:num_to_del]:
                to_exclude.append((func, tm.test_file.path))
                # CARE: this operation has changes state of src_repo,
                # which is then propagated to strategy below
                src_repo.find_file(tm.path).delete(
                    func.name, node_type=NodeType.Function
                )
                # tm.test_file.delete(func.name, node_type=NodeType.Function)

                with open(src_repo.repo_path / tm.test_file.path, "w") as f:
                    # print(tm.test_file.to_code())
                    f.write(src_repo.find_file(tm.path).to_code())

                total_deleted += 1
        except Exception as e:
            failed_mod += 1

    print("Total failed:", failed_mod)


--------------------------------------------------------------------------------

Chunk ID: scripts/neuter_repo.py::3
def neuter_tests2(
    test_modules: List[TestModule], 
    src_repo: SourceRepo, 
    to_keep, 
    to_delete=0, 
    new_repo: Optional[Path] = None
):
    total_deleted = 0
    failed_mod = 0
    for test_file in src_repo.test_files:
        for tm in [tm for tm in test_modules if tm.test_file == test_file]:
            og_bytes = len(open(src_repo.repo_path / test_file.path, "r").read())
            try:
                print("Deleting tm: ", tm.name)
                to_exclude = []
                # BUG: tm.tests gets modified somehow
                num_to_del = num_delete(tm, to_keep=to_keep, to_delete=to_delete)
                total_tests = len(tm.tests)

                for func in tm.tests[:num_to_del]:
                    to_exclude.append((func, tm.test_file.path))
                    # CARE: this operation has changes state of src_repo,
                    # which is then propagated to strategy below
                    src_repo.find_file(tm.path).delete(
                        func.name, node_type=NodeType.Function
                    )
                    # tm.test_file.delete(func.name, node_type=NodeType.Function)

                    total_deleted += 1

            except Exception as e:
                failed_mod += 1

        output_path = (new_repo if new_repo else src_repo.repo_path) / test_file.path
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w") as f:
            f.write(src_repo.find_file(test_file.path).to_code())

        new_bytes = len(src_repo.find_file(test_file.path).to_code())
        print(f"Reduced {test_file.path} by {og_bytes - new_bytes} bytes and {total_deleted} tests")

    print("Total failed:", failed_mod)


--------------------------------------------------------------------------------

Chunk ID: scripts/neuter_repo.py::4
if __name__ == "__main__":
    """
    python -m neuter_repo <repo_path> [new_repo_path]
    """
    repo = Path(sys.argv[1])
    if not repo.exists():
        print("Repo does not exist")
        sys.exit()

    new_repo = Path(sys.argv[2]) if len(sys.argv) > 2 else None

    src_repo = SourceRepo(repo)
    test_modules = iter_test_modules(src_repo)

    neuter_tests(test_modules, src_repo, to_keep=2, to_delete=0)


--------------------------------------------------------------------------------

Chunk ID: scripts/show_tables.py::1
from cowboy_lib.repo.source_repo import SourceRepo
from src.test_modules.iter_tms import iter_test_modules
from src.test_modules.models import TestModuleModel
from pathlib import Path

from sqlalchemy.sql import text


from src.database.core import engine
from sqlalchemy.orm import sessionmaker

Session = sessionmaker(bind=engine)
query = text(
    """
    SELECT schema_name FROM information_schema.schemata
    WHERE schema_name NOT IN ('pg_catalog', 'information_schema')
    AND schema_name NOT LIKE 'pg_toast%'
    AND schema_name NOT LIKE 'pg_temp_%'
"""
)

# Execute the query
with engine.connect() as connection:
    # Query to get all tables from the public schema
    table_query = text(
        """
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'public'
    """
    )

    # Execute the query to get all table names
    with engine.connect() as connection:
        tables = connection.execute(table_query).fetchall()

        # Iterate through each table and get its schema details
        for table in tables:
            table_name = table[0]
            print(f"\nSchema for table '{table_name}':")

            # Query to get schema details for each table
            schema_query = text(
                f"""
                SELECT column_name, data_type, is_nullable, column_default
                FROM information_schema.columns
                WHERE table_schema = 'public' AND table_name = :table_name
                ORDER BY ordinal_position
            """
            )

            # Fetch and print the schema details for each table
            schema_details = connection.execute(
                schema_query, {"table_name": table_name}
            ).fetchall()
            for detail in schema_details:
                print(detail)
                # print(f"Column Name: {detail['column_name']}, "
                #       f"Type: {detail['data_type']}, "
                #       f"Nullable: {detail['is_nullable']}, "
                #       f"Default: {detail['column_default']}")


--------------------------------------------------------------------------------

Chunk ID: stats/models.py::1
from sqlalchemy import Column, Integer, String, JSON, ForeignKey, Boolean
from src.database.core import Base


# TODO: make this track module level stats later
class RepoStats(Base):
    """
    Tracks repo level stats
    """

    __tablename__ = "repo_stats"

    id = Column(Integer, primary_key=True)
    total_tests = Column(Integer, default=0)
    accepted_tests = Column(Integer, default=0)
    rejected_tests = Column(Integer, default=0)

    repo_id = Column(Integer, ForeignKey("repo_config.id", ondelete="CASCADE"))


--------------------------------------------------------------------------------

Chunk ID: stats/service.py::1
from src.repo.models import RepoConfig
from .models import RepoStats

from contextlib import contextmanager
from sqlalchemy.orm import Session


@contextmanager
def update_repo_stats(*, db_session: Session, repo: RepoConfig):
    try:
        stats = db_session.query(RepoStats).filter_by(repo_id=repo.id).one_or_none()
        if not stats:
            stats = RepoStats(
                repo_id=repo.id, total_tests=0, accepted_tests=0, rejected_tests=0
            )
            db_session.add(stats)
        yield stats

        db_session.commit()
    except Exception as e:
        db_session.rollback()
        raise e


--------------------------------------------------------------------------------

Chunk ID: src/sync_repos.py::1
from cowboy_lib.repo.repository import GitRepo
from cowboy_lib.repo.diff import DiffMode

from src.queue.core import TaskQueue
from src.auth.service import get_user_token
from src.repo.service import get_all, get_or_raise
from src.test_modules.service import get_all_tms
from src.test_modules.models import TestModuleModel
from src.config import API_ENDPOINT

from src.logger import sync_repo as log

from pathlib import Path
import asyncio
from sqlalchemy.orm import Session
from typing import List
import threading
import requests
from urllib.parse import urljoin


class ChangedFiles:
    code_files: List[str]
    test_files: List[str]
    repo: int


def start_sync_thread(db_session: Session, task_queue: TaskQueue):
    def run_async_thread(loop, func, *args):
        """
        Helper func to run async functions in a new thread
        """
        asyncio.set_event_loop(loop)
        loop.run_until_complete(func(*args))

    threading.Thread(
        target=run_async_thread,
        args=(
            asyncio.new_event_loop(),
            check_for_changed_files,
            db_session,
            task_queue,
        ),
        daemon=True,
    ).start()


def testfile_to_tm(db: Session, repo_id: str, file: str) -> TestModuleModel:
    """
    Maps the file back to the test_module. This is a one-to-one mapping
    """
    tms = get_all_tms(db_session=db, repo_id=repo_id)
    for tm in tms:
        if tm.testfilepath == file:
            return tm

    return None


def srcfile_to_tm(db: Session, repo_id: str, file: str) -> TestModuleModel:
    """
    Maps the src file back to test_module. This is n-to-one mapping
    """
    tms = get_all_tms(db_session=db, repo_id=repo_id)
    for tm in tms:
        if file in tm.get_covered_files():
            return tm

    return None


def is_test_file(path: str):
    return Path(path).name.startswith("test_")


--------------------------------------------------------------------------------

Chunk ID: src/sync_repos.py::2
class APIClient:
    def __init__(self, endpoint, token):
        self.server = endpoint
        self.headers = {"Authorization": f"Bearer {token}"}

    async def get(self, uri: str):
        url = urljoin(self.server, uri)
        res = requests.get(url, headers=self.headers)

        return res.json()

    async def post(self, uri: str, data: dict):
        url = urljoin(self.server, uri)
        res = requests.post(url, json=data, headers=self.headers)

        return res.json()

    async def build_mapping(self, repo_name, mode, tms):
        """
        Builds the test module to source file mapping for each selected
        test module
        """
        await self.post(
            "/tm/build-mapping",
            {
                "repo_name": repo_name,
                "mode": mode,
                "tms": tms,
                "files": [],
                # TODO: change this to False
                "overwrite": True,
            },
        )


--------------------------------------------------------------------------------

Chunk ID: src/sync_repos.py::3
async def check_for_changed_files(
    db_session: Session, task_queue: TaskQueue
) -> ChangedFiles:
    """
    Checks for repo update
    """
    while True:
        repos = get_all(db_session=db_session)
        for repo in repos:
            git_repo = GitRepo(Path(repo.source_folder))
            user_token = get_user_token(db_session=db_session, user_id=repo.user_id)
            api_client = APIClient(API_ENDPOINT, user_token)

            commit = git_repo.diff_remote()
            if commit:
                for diff in commit.diffs:
                    if diff.attrs.mode == DiffMode.MODIFIED:
                        modified_file = diff.attrs.a_path
                        impacted_tm = (
                            testfile_to_tm(db_session, repo.id, modified_file)
                            if is_test_file(modified_file)
                            else srcfile_to_tm(db_session, repo.id, modified_file)
                        )
                        log.info(f"Impacted tm: {impacted_tm.name}")

                        if not impacted_tm:
                            # TODO: really bad, because it implies that we are
                            # missing some kind of update to the repo
                            return

                        await api_client.build_mapping(
                            repo.repo_name, "module", [impacted_tm.name]
                        )

        await asyncio.sleep(10)


--------------------------------------------------------------------------------

Chunk ID: target_code/models.py::1
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey
from sqlalchemy.orm import relationship
from pydantic import BaseModel
from pathlib import Path
from typing import List

from cowboy_lib.test_modules.test_module import TestModule
from cowboy_lib.test_modules.target_code import TargetCode
from cowboy_lib.repo.source_repo import SourceRepo
from src.database.core import Base
from src.ast.models import NodeModel


class TargetCodeModel(Base):
    """
    A chunk of code that is covered by the lines in a TestModule
    """

    __tablename__ = "target_code"
    id = Column(Integer, primary_key=True)
    start = Column(Integer)
    end = Column(Integer)
    lines = Column(String)
    filepath = Column(String)

    func_scope = relationship(
        "NodeModel",
        foreign_keys=[NodeModel.target_code_id],
        cascade="all",
        uselist=False,
        single_parent=True,
    )
    class_scope = relationship(
        "NodeModel",
        foreign_keys=[NodeModel.target_code_id],
        cascade="all",
        uselist=False,
        single_parent=True,
    )
    test_module_id = Column(Integer, ForeignKey("test_modules.id", ondelete="CASCADE"))
    coverage_id = Column(Integer, ForeignKey("coverage.id", ondelete="CASCADE"))
    coverage = relationship("CoverageModel")

    def __init__(
        self,
        start,
        end,
        lines,
        filepath,
        func_scope,
        class_scope,
        test_module_id,
        coverage_id,
    ):
        self.start = start
        self.end = end
        self.lines = "\n".join(lines)
        self.filepath = str(filepath)
        self.func_scope = func_scope
        self.class_scope = class_scope
        self.test_module_id = test_module_id
        self.coverage_id = coverage_id

    def get_lines(self) -> List[str]:
        return self.lines.split("\n")

    def to_str(self) -> str:
        repr = ""
        repr += f"Chunk: {self.filepath}\n"
        repr += self.lines

        return repr

    def serialize(self, src_repo: SourceRepo):
        return TargetCode(
            range=(self.start, self.end),
            lines=self.get_lines(),
            filepath=Path(self.filepath),
            func_scope=(
                self.func_scope.to_astnode(src_repo) if self.func_scope else None
            ),
            class_scope=(
                self.class_scope.to_astnode(src_repo) if self.class_scope else None
            ),
        )


class TgtCodeDeleteRequest(BaseModel):
    repo_name: str
    tm_name: str


--------------------------------------------------------------------------------

Chunk ID: target_code/service.py::1
from src.ast.models import NodeModel
from src.target_code.models import TargetCode, TargetCodeModel
from src.test_modules.models import TestModuleModel
from src.coverage.models import CoverageModel

from sqlalchemy.orm import Session

def delete_target_code(db_session: Session, tm_id: int):
    """Delete all target code for a test module."""

    deleted = (
        db_session.query(TargetCodeModel)
        .filter(TargetCodeModel.test_module_id == tm_id)
        .delete()
    )

    db_session.commit()

    return deleted


--------------------------------------------------------------------------------

Chunk ID: target_code/views.py::1
from fastapi import APIRouter, Depends, HTTPException
from src.database.core import get_db
from src.repo.service import get_or_raise
from src.auth.service import get_current_user

from src.test_modules.service import get_tm_by_name

from .models import TgtCodeDeleteRequest
from .service import delete_target_code


tgtcode_router = APIRouter()


@tgtcode_router.post("/tgt_code/delete/")
async def delete_tgt_code(
    tgt_delete_req: TgtCodeDeleteRequest,
    user=Depends(get_current_user),
    db=Depends(get_db),
):
    """
    Delete all target code for a test module
    """
    try:
        repo = get_or_raise(
            db_session=db, curr_user=user, repo_name=tgt_delete_req.repo_name
        )
        tm_model = get_tm_by_name(
            db_session=db, repo_id=repo.id, tm_name=tgt_delete_req.tm_name
        )
        deleted = delete_target_code(db_session=db, tm_id=tm_model.id)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    return {"detail": "Target code deleted"}


--------------------------------------------------------------------------------

Chunk ID: tasks/__init__.py::1
from .fork_repo import fork_repo


--------------------------------------------------------------------------------

Chunk ID: tasks/create_tgt_coverage.py::1
from cowboy_lib.repo import SourceRepo
from cowboy_lib.test_modules import TargetCode

# from .models import TestModule
# # Long term tasks represent tasks that we potentially want to offload to celery
from src.tasks.get_baseline_parallel import get_tm_target_coverage

from src.queue.core import TaskQueue
from src.repo.models import RepoConfig
from src.test_modules.models import TestModuleModel
from src.ast.models import NodeModel
from src.coverage.models import CoverageModel, coverage_to_model

from src.runner.service import RunServiceArgs
from src.target_code.models import TargetCodeModel
from src.utils import async_timed
from src.runtest_conf import run_test

from src.logger import buildtm_logger as log

from sqlalchemy.orm import Session

from pathlib import Path
from typing import List, Callable


--------------------------------------------------------------------------------

Chunk ID: tasks/create_tgt_coverage.py::2
def tgtcode_to_model(
    tc: TargetCode,
    cov_model: CoverageModel,
    repo_id: int,
    tm_model: TestModuleModel,
) -> TargetCodeModel:
    """
    Create target code models with associated nodes in a single transaction
    """
    func_scope = (
        NodeModel(
            repo_id=repo_id,
            testfilepath=str(tc.filepath),
            name=tc.func_scope.name,
            node_type=tc.func_scope.node_type,
        )
        if tc.func_scope
        else None
    )
    class_scope = (
        NodeModel(
            repo_id=repo_id,
            testfilepath=str(tc.filepath), 
            name=tc.class_scope.name,
            node_type=tc.class_scope.node_type,
        )
        if tc.class_scope
        else None
    )

    return TargetCodeModel(
        start=tc.range[0],
        end=tc.range[1],
        lines=tc.lines,
        filepath=str(tc.filepath),
        func_scope=func_scope,
        class_scope=class_scope,
        test_module_id=tm_model.id,
        coverage_id=cov_model.id,
    )


--------------------------------------------------------------------------------

Chunk ID: tasks/create_tgt_coverage.py::3
@async_timed
async def create_tgt_coverage(
    *,
    db_session: Session,
    task_queue: TaskQueue,
    repo: RepoConfig,
    tm_models: List[TestModuleModel],
    run_test: Callable,
    overwrite: bool = True,
):
    """
    Important function that sets up relationships between TestModule, TargetCode and
    Coverage
    """
    repo_path = Path(repo.source_folder)
    src_repo = SourceRepo(repo_path)
    run_args = RunServiceArgs(repo.user_id, task_queue)
    base_cov = repo.base_cov

    # NEWDESIGN: in future, consider putting detecting TM updates due to mapped
    # module changes here

    # NEWTODO: 
    # if overwrite or not base_cov:
    #     log.info("Overwriting base coverage ... th ")
    #     cov_res = await run_test(repo.repo_name, run_args)
    #     base_cov = cov_res.coverage
    #     upsert_coverage(
    #         db_session=db_session, repo_id=repo.id, cov_list=base_cov.cov_list
    #     )

    for tm_model in tm_models:
        log.info(f"Building target coverage for: {tm_model.name}")
        # only overwrite existing target_chunks if overwrite flag is set
        if not overwrite and tm_model.target_chunks:
            log.info(f"Skipping {tm_model.name} cuz overwrite: {overwrite}, has_target_chunks: {(bool(tm_model.target_chunks))}")
            continue

        tm = tm_model.serialize(src_repo)

        # generate src th 
        tgt_chunks = await get_tm_target_coverage(
            repo.repo_name, src_repo, tm, base_cov, run_test, run_args
        )
        target_models = []
        for tc in tgt_chunks:
            file_cov = base_cov.get_file_cov(tc.filepath)
            # NEWDESIGN: we are okay with generating new coverage model here for now
            # UPDATE: actually there is issue with null ID here so coverage model is not being created properly
            file_cov_model = coverage_to_model(file_cov)
            target_models.append(tgtcode_to_model(tc, file_cov_model, repo.id, tm_model))

        # TODO: convert this to a Logfire log?
        # log.info(f"{tm_model.name} Chunks: ")
        # for t in target_models:
        #     log.info(t.to_str())

        tm_model.target_chunks = target_models
        db_session.merge(tm_model)
        db_session.commit()


--------------------------------------------------------------------------------

Chunk ID: tasks/fork_repo.py::1
import requests
from typing import Tuple, List, Optional
import time
from logging import getLogger
import os

from dataclasses import dataclass, field

from git import Diff

# is this the right way? Feel like proper way is to load config
# at single point upstream and then pass it down .. but idk

# from dotenv import load_dotenv

# load_dotenv()
# GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN")

logger = getLogger("test_results")


class GithubAPI:
    access_token = "1234"

    def _make_github_api_request(self, url: str):
        headers = {
            "Authorization": f"token {self.access_token}",
            "Accept": "application/vnd.github.v3+json",
        }

        response = requests.get(url, headers=headers)
        remaining_rate_limit = int(response.headers.get("X-RateLimit-Remaining", 0))
        if remaining_rate_limit < 10:
            print("Rate limit exceeded, sleeping for 15 seconds")
            time.sleep(15)

        return response


--------------------------------------------------------------------------------

Chunk ID: tasks/fork_repo.py::2
class GithubAPI:

    @classmethod
    def fork_repository(cls, repo_full_name: str) -> str:
        """
        Forks the given repository.

        Parameters:
        - repo_full_name: The full name of the repository to fork (e.g., "owner/repo_name").

        Returns:
        The URL of the forked repository or None if the request fails.
        """
        url = f"https://api.github.com/repos/{repo_full_name}/forks"
        headers = {
            "Authorization": f"token {cls.access_token}",
            "Accept": "application/vnd.github.v3+json",
        }

        response = requests.post(url, headers=headers)
        if response.status_code in [202, 201]:  # 202 Accepted or 201 Created
            forked_repo_url = response.json().get("html_url")
            return forked_repo_url
        else:
            print(
                f"Failed to fork the repository {repo_full_name}. Status code: {response.status_code}"
            )
            return None


def fork_repo(url) -> str:
    return GithubAPI.fork_repository(url)


--------------------------------------------------------------------------------

Chunk ID: tasks/get_baseline_parallel.py::1
from cowboy_lib.repo.source_repo import SourceRepo
from cowboy_lib.coverage import Coverage, TestCoverage, CoverageResult
from cowboy_lib.test_modules.test_module import TestModule, TargetCode
from cowboy_lib.utils import testfiles_in_coverage

from src.queue.core import TaskQueue
from src.runner.service import RunServiceArgs
from src.logger import buildtm_logger as log

from typing import List, Tuple, Callable
from pathlib import Path
import asyncio
from collections import defaultdict

class TestInCoverageException(Exception):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.msg = "Test files are included in coverage report"


--------------------------------------------------------------------------------

Chunk ID: tasks/get_baseline_parallel.py::2
def set_chunks(
    changed_coverage: List[Coverage],
    source_repo: "SourceRepo",
    base_path: Path = None,
) -> List[TargetCode]:
    """
    Gets the missing/covered lines of each of the coverage differences
    """
    chunks = []
    for cov in changed_coverage:
        if cov.filename == "TOTAL":
            raise Exception("TOTAL COVERAGE FILE FOUND")

        cov.read_line_contents(base_path)
        for l_group in cov.get_contiguous_lines():
            start = l_group[0][0]
            end = l_group[-1][0]
            range = (start, end)

            src_file = source_repo.find_file(cov.filename)
            func, cls = src_file.map_line_to_node(start, end)

            lines = [g[1] for g in l_group]

            # print("Setting chunk with filepath: ", str(cov.filename))

            chunk = TargetCode(
                range=range,
                lines=lines,
                # could also just move the logic into TestModuleMixin
                filepath=str(cov.filename),
                func_scope=func if func else "",
                class_scope=cls if cls else "",
            )
            chunks.append(chunk)

    return chunks

def compare(base_cov: TestCoverage, module: TestCoverage):
    for cov1, cov2 in zip(base_cov.cov_list, module.cov_list):
        print(cov1.filename, cov2.filename)
        if cov1.stmts != cov2.stmts:
            print("Different stmts ")


--------------------------------------------------------------------------------

Chunk ID: tasks/get_baseline_parallel.py::3
async def get_tm_target_coverage(
    repo_name: str,
    src_repo: SourceRepo,
    tm: TestModule,
    base_cov: TestCoverage,
    run_test: Callable,
    run_args: RunServiceArgs,
) -> List[TargetCode]:
    """
    Test augmenting existing test classes by deleting random test methods, and then
    having LLM strategy generate them. Coverage is taken:
    1. After the deletion
    2. After the deletion with newly generated LLM testcases

    The diff measures how well we are able to supplant the coverage of the deleted methods
    """

    if testfiles_in_coverage(base_cov, src_repo):
        raise TestInCoverageException

    # First loop we find the total coverage of each test by itself
    only_module = [tm.name]
    # coverage with ONLY the current test module turned on
    log.info(f"Collecting target chunks for {tm.name}")

    # TODO: should be storing this as well
    module_cov = await run_test(
        repo_name,
        run_args,
        # part1: collect the coverage of a single module only
        include_tests=only_module,
    )
    log.info(f"BaseCov: {base_cov}")
    log.info(f"ModuleCov: {module_cov.coverage}")

    module_diff = base_cov - module_cov.coverage
    total_cov_diff = module_diff.total_cov.covered
    if total_cov_diff > 0:
        chg_cov = []
        coroutines = []

        for test in tm.tests:
            print("Running test ... ", test.name)
            task = run_test(
                repo_name,
                run_args,
                # part 2:
                # holds the coverage diff of individual tests after they have
                # been selectively turned off
                exclude_tests=[(test, tm.test_file.path)],
                include_tests=only_module,
            )
            coroutines.append(task)

        test_coverage = defaultdict(int)
        cov_res = await asyncio.gather(*[t for t in coroutines])
        for test, cov_res in zip(tm.tests, cov_res):                
            # log.info(f"Test results: {cov_res.coverage.total_cov.covered}")
            # log.info(
            #     f"Module cov: {module_cov.coverage.total_cov.covered}, Single cov: {cov_res.coverage.total_cov.covered}"
            # )

            # part 3: we subtract the module from the 
            single_diff: TestCoverage = module_cov.coverage - cov_res.coverage
            test_coverage[test.name] += single_diff.total_cov.covered

            # dont think we actually need this here .. confirm
            chg_cov.extend(single_diff.cov_list)

        # NEWTODO: potentially we should store this on a TestModel object
        # although we would first have to create that ...
        for test, coverage in test_coverage.items():
            log.info(f"{test}:{tm.name} covered by {coverage} lines")

        # re-init the chunks according to the aggregated individual test coverages
        chunks = set_chunks(
            chg_cov,
            source_repo=src_repo,
            base_path=src_repo.repo_path,
        )

    # Find out what's the reason for the missed tests
    else:
        log.info(f"No coverage difference found for {tm.name}")
        return []

    return chunks


--------------------------------------------------------------------------------

Chunk ID: test_gen/augment.py::1
from cowboy_lib.repo import SourceRepo, GitRepo
from cowboy_lib.test_modules.test_module import TestModule

from .augment_test.composer import Composer
from .models import AugmentTestResult

from src.stats.service import update_repo_stats
from src.auth.models import CowboyUser
from src.auth.service import retrieve_oai_key
from src.runner.service import RunServiceArgs
from src.queue.core import TaskQueue
from src.repo.models import RepoConfig
from src.test_modules.models import TestModuleModel
from src.config import AUGMENT_ROUNDS

from .service import create_test_result

from sqlalchemy.orm import Session
from pathlib import Path
from typing import List


# FEATURE: add lines additional lines covered here
def commit_message(test_names: List[str], cov_plus: int):
    names_str = "\n".join(test_names)
    return "Added test cases: \n" + names_str


--------------------------------------------------------------------------------

Chunk ID: test_gen/augment.py::2
async def augment_test(
    *,
    db_session: Session,
    task_queue: TaskQueue,
    repo: RepoConfig,
    tm_model: TestModuleModel,
    curr_user: CowboyUser,
    session_id: str
) -> List[AugmentTestResult]:
    """
    Generate test cases for the given test module using the specified strategy and evaluator
    """
    src_repo = SourceRepo(Path(repo.source_folder))
    git_repo = GitRepo(Path(repo.source_folder), remote=repo.remote, main=repo.main)
    tm = tm_model.serialize(src_repo)
    run_args = RunServiceArgs(user_id=curr_user.id, task_queue=task_queue)

    base_cov = repo.base_cov
    composer = Composer(
        repo_name=repo.repo_name,
        strat="WITH_CTXT",
        evaluator="ADDITIVE",
        src_repo=src_repo,
        test_input=tm,
        run_args=run_args,
        base_cov=base_cov,
        api_key=retrieve_oai_key(curr_user.id),
    )

    improved_tests, failed_tests, no_improve_tests = await composer.generate_test(
        n_times=AUGMENT_ROUNDS
    )

    # NEWTODO: should add some options for other output formats
    # write all improved test to source file and check out merge on repo
    # serialize tm first
    test_results = []
    test_file = tm.test_file
    for improved, cov in improved_tests:
        test_result = create_test_result(
            db_session=db_session,
            repo_id=repo.id,
            name=improved.name,
            test_case=improved.to_code(),
            cov_list=cov.cov_list,
            tm_id=tm_model.id,
            commit_hash=git_repo.get_curr_commit(),
            testfile=str(test_file.path),
            session_id=session_id,
            classname=None,
        )
        test_results.append(test_result)

    # update repo stats
    with update_repo_stats(db_session=db_session, repo=repo) as repo_stats:
        repo_stats.total_tests += (
            len(improved_tests) + len(failed_tests) + len(no_improve_tests)
        )

    return test_results


--------------------------------------------------------------------------------

Chunk ID: augment_test/base_strat.py::1
from cowboy_lib.repo.source_repo import SourceRepo

from abc import ABC, abstractmethod
from pathlib import Path
from dataclasses import dataclass

from src.runner.service import RunServiceArgs


@dataclass
class TestCaseInput(ABC):

    @property
    @abstractmethod
    def path(self) -> Path:
        raise NotImplementedError


class BaseStrategy(ABC):
    def __init__(self, src_repo: SourceRepo, test_input: TestCaseInput):
        self.src_repo = src_repo
        self.test_input = test_input

    @abstractmethod
    def build_prompt(self) -> str:
        """
        Builds the base prompt according to the strategy
        """

        raise NotImplementedError

    @abstractmethod
    def parse_llm_res(self):
        """
        Parses the LLM response to get the generated code
        """
        raise NotImplementedError


--------------------------------------------------------------------------------

Chunk ID: augment_test/composer.py::1
from .base_strat import BaseStrategy, TestCaseInput
from .types import StratResult
from .evaluators import (
    Evaluator,
    AugmentAdditiveEvaluator,
    AugmentParallelEvaluator,
    EvaluatorType,
    AUGMENT_EVALS,
)

from cowboy_lib.llm.invoke_llm import invoke_llm_async
from cowboy_lib.llm.models import OpenAIModel, ModelArguments
from cowboy_lib.repo.source_repo import SourceRepo
from cowboy_lib.repo.source_file import Function, LintException
from cowboy_lib.coverage import TestCoverage, TestError

from src.test_gen.augment_test.strats import AugmentStratType, AUGMENT_STRATS
from src.runner.service import RunServiceArgs
from src.exceptions import CowboyRunTimeException
from src.logger import testgen_logger


from src.config import LLM_RETRIES

from typing import Tuple, List


--------------------------------------------------------------------------------

Chunk ID: augment_test/composer.py::2
class Composer:
    """
    Used to instantiate different combinations of strategies for generating test cases
    """

    def __init__(
        self,
        repo_name: str,
        strat: AugmentStratType,
        evaluator: EvaluatorType,
        src_repo: SourceRepo,
        test_input: TestCaseInput,
        run_args: RunServiceArgs,
        base_cov: TestCoverage,
        api_key: str,
        verify: bool = False,
    ):
        self.repo_name = repo_name
        self.src_repo = src_repo
        self.test_input = test_input
        self.verify = verify
        self.base_cov = base_cov
        self.run_args = run_args

        self.strat: BaseStrategy = AUGMENT_STRATS[strat](self.src_repo, self.test_input)
        self.evaluator: Evaluator = AUGMENT_EVALS[evaluator](
            self.repo_name, self.src_repo, self.run_args
        )

        model_name = "gpt4"
        self.model = OpenAIModel(ModelArguments(model_name=model_name, api_key=api_key))

    def get_strat_name(self) -> str:
        return self.__class__.__name__

    def filter_overlap_improvements(
        self, tests: List[Tuple[Function, TestCoverage]]
    ) -> List[Tuple[Function, TestCoverage]]:
        no_overlap = []
        overlap_cov = self.base_cov
        for test, cov in tests:
            new_cov = overlap_cov + cov
            if new_cov.total_cov.covered > overlap_cov.total_cov.covered:
                no_overlap.append((test, cov))
                overlap_cov = new_cov

        return no_overlap

    # TODO: this function name is a lie, we should parallelize this


--------------------------------------------------------------------------------

Chunk ID: augment_test/composer.py::3
class Composer:
    async def gen_test_parallel(self, n_times: int) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:

        improved_tests = []
        failed_tests = []
        no_improve_tests = []

        # TODO: here is whee we initialize the StartCodeTx
        prompt = self.strat.build_prompt()
        print(f"Prompt: {prompt}")

        model_res = await invoke_llm_async(prompt, self.model, n_times)

        llm_results = [self.strat.parse_llm_res(res) for res in model_res]
        test_results = [StratResult(res, self.test_input.path) for res in llm_results]

        improved, failed, no_improve = await self.evaluator(
            test_results,
            self.test_input,
            self.base_cov,
            n_times=n_times,
        )

        improved_tests.extend(improved)
        filtered_improved = self.filter_overlap_improvements(improved_tests)
        improved_tests = filtered_improved

        failed_tests.extend(failed)
        no_improve_tests.extend(no_improve)

        return improved_tests, failed_tests, no_improve_tests


--------------------------------------------------------------------------------

Chunk ID: augment_test/composer.py::4
class Composer:

    async def gen_test_serial_additive(self, n_times: int) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        if not isinstance(self.evaluator, AugmentAdditiveEvaluator):
            raise Exception(
                f"Expected AugmentAdditiveEvaluator, got {self.evaluator.__class__}"
            )

        improved_tests = []
        failed_tests = []
        no_improve_tests = []
        prompt = self.strat.build_prompt()

        print("Prompt: ", prompt)

        for _ in range(n_times):
            retries = LLM_RETRIES
            src_file = None
            while retries > 0 and not src_file:
                try:
                    llm_res = await invoke_llm_async(
                        prompt,
                        model=self.model,
                        n_times=1,
                    )
                    src_file = self.strat.parse_llm_res(llm_res[0])
                except (SyntaxError, ValueError, LintException):
                    testgen_logger.info(f"LLM syntax error ... {retries} left")
                    retries -= 1
                    continue

            if not src_file:
                raise CowboyRunTimeException(
                    f"LLM generation failed for {self.test_input}"
                )

            test_result = [StratResult(src_file, self.test_input.path)]
            improved, failed, no_improve = await self.evaluator(
                test_result,
                self.test_input,
                self.base_cov,
                n_times=n_times,
            )
            improved_tests.extend(improved)
            filtered_improved = self.filter_overlap_improvements(improved_tests)
            improved_tests = filtered_improved

            # update test input with new functions that improved coverage
            for new_func in [
                func
                for func, _ in improved
                if func in [f[0] for f in filtered_improved]
            ]:
                self.test_input.test_file.append(
                    new_func.to_code(),
                    # wrong too, we need to check the
                    class_name=new_func.scope.name if new_func.scope else "",
                )

            failed_tests.extend(failed)
            no_improve_tests.extend(no_improve)

        return improved_tests, failed_tests, no_improve_tests

    async def generate_test(self, n_times: int) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        if isinstance(self.evaluator, AugmentAdditiveEvaluator):
            return await self.gen_test_serial_additive(n_times)
        elif isinstance(self.evaluator, AugmentParallelEvaluator):
            return await self.gen_test_parallel(n_times)


--------------------------------------------------------------------------------

Chunk ID: evaluators/__init__.py::1
from .eval_base import Evaluator

from .augment_parallel import AugmentParallelEvaluator
from .augment_additive import AugmentAdditiveEvaluator

from enum import Enum


class EvaluatorType(Enum):
    PARALLEL = "PARALLEL"
    ADDITIVE = "ADDITIVE"


AUGMENT_EVALS = {
    "PARALLEL": AugmentParallelEvaluator,
    "ADDITIVE": AugmentAdditiveEvaluator,
}


--------------------------------------------------------------------------------

Chunk ID: evaluators/augment_additive.py::1
from cowboy_lib.coverage import CoverageResult, TestError, TestCoverage
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.repo.source_file import Function

from typing import Tuple, List, TYPE_CHECKING

if TYPE_CHECKING:
    from test_gen.augment_test.types import StratResult
    from cowboy_lib.test_modules import TestModule

from .eval_base import Evaluator

from src.runner.service import run_test
from src.logger import testgen_logger


class AugmentAdditiveEvaluator(Evaluator):
    """
    Iteratively evals test results and re-prompts with partially successful
    test file to **attempt** to get additive coverage
    """

    async def __call__(
        self,
        llm_results: List["StratResult"],
        tm: "TestModule",
        base_cov: TestCoverage,
        n_times: int = 1,
    ) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
        "TestModule",
    ]:
        """
        Main eval method, accepts a list of results from the strategy and the
        targeted test module, and a baseline coverage to compare against
        """
        test_fp = tm.test_file.path
        test_results = await self.gen_test_and_diff_coverage(
            llm_results, base_cov, test_fp, n_times
        )
        improved, failed, no_improve = await self.process_test_results(
            test_results, tm, base_cov
        )

        return improved, failed, no_improve

    # questionable decision to make non-existent func Functions ..


--------------------------------------------------------------------------------

Chunk ID: evaluators/augment_additive.py::2
class AugmentAdditiveEvaluator(Evaluator):
    async def process_test_results(
        self,
        test_results: List[Tuple[CoverageResult, str]],
        tm: "TestModule",
        base_cov: TestCoverage,
    ) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        """
        Sequentially build a set of coverage improving testcases, discarding any
        generated tests that dont contribute coverage improvements
        """
        improved_tests: List[Tuple[Function, TestCoverage]] = []
        failed_tests: List[Tuple[Function, TestError]] = []
        noimprov_tests: List[Function] = []

        for cov_res, cov_diff, test_file in test_results:
            if cov_diff:
                new_funcs = self.get_new_funcs(test_file, tm.path)
                # iterate each generated function and measure if it has coverage
                # improvement against the base
                for func in new_funcs:
                    test_error = cov_res.get_failed(func.name)
                    if test_error:
                        testgen_logger.info(f"[FAILED] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        failed_tests.append((func, test_error))
                        continue

                    # TODO: make sure that this works for filename TMs as well
                    og_testfile = self.src_repo.find_file(tm.path).clone()
                    og_testfile.append(
                        func.to_code(), class_name=func.scope.name if func.scope else ""
                    )

                    patch_file = PatchFile(
                        path=str(tm.path), patch=og_testfile.to_code()
                    )
                    indvtest_cov = await run_test(
                        self.repo_name,
                        self.run_args,
                        patch_file=patch_file,
                    )

                    indv_improve = indvtest_cov.coverage - base_cov
                    if indv_improve.total_cov.covered > 0:
                        testgen_logger.info(f"[IMPROVE] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        improved_tests.append((func, indv_improve))
                    else:
                        testgen_logger.info(f"[NOIMPROVE] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        noimprov_tests.append((func, TestCoverage([])))

        return improved_tests, failed_tests, noimprov_tests


--------------------------------------------------------------------------------

Chunk ID: evaluators/augment_parallel.py::1
from cowboy_lib.coverage import CoverageResult, TestError, TestCoverage
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.repo.source_file import Function

from typing import Tuple, List, TYPE_CHECKING

if TYPE_CHECKING:
    from test_gen.augment_test.types import StratResult
    from cowboy_lib.test_modules import TestModule

from .eval_base import Evaluator

from concurrent.futures import ThreadPoolExecutor
from src.logger import testgen_logger


class AugmentParallelEvaluator(Evaluator):
    """
    Used to evaluate the results of a test strategy
    """

    async def __call__(
        self,
        llm_results: List["StratResult"],
        tm: "TestModule",
        base_cov: CoverageResult,
        n_times: int = 1,
    ) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        """
        Main eval method, accepts a list of results from the strategy and the
        targeted test module, and a baseline coverage to compare against
        """
        test_fp = tm.test_file.path
        test_results = await self.gen_test_and_diff_coverage(
            llm_results, base_cov, test_fp, n_times
        )
        improved, failed, no_improve = await self.process_test_results(
            test_results, tm, base_cov
        )

        return improved, failed, no_improve

    # questionable decision to make non-existent func Functions ..


--------------------------------------------------------------------------------

Chunk ID: evaluators/augment_parallel.py::2
class AugmentParallelEvaluator(Evaluator):
    async def process_test_results(
        self,
        test_results: List[Tuple[CoverageResult, str]],
        tm: "TestModule",
        del_cov: CoverageResult,
    ) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        """
        Takes the results from run_test and processes it into a list of
        coverage enhancing test cases and failed tests
        """
        improved_tests: List[Tuple[Function, TestCoverage]] = []
        failed_tests: List[Tuple[Function, TestError]] = []
        noimprov_tests: List[Function] = []

        for cov_res, cov_diff, test_file in test_results:
            if cov_diff:
                new_funcs = self.get_new_funcs(test_file, tm.path)
                # 1. create args for each runner thread
                runner_args = {}
                for func in new_funcs:
                    test_error = cov_res.get_failed(func.name)
                    if test_error:
                        testgen_logger.info(f"[FAILED] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        failed_tests.append((func, test_error))
                        continue

                    # TODO: make sure that this works for filename TMs as well
                    og_testfile = self.src_repo.find_file(tm.path).clone()
                    og_testfile.append(
                        func.to_code(), class_name=func.scope.name if func.scope else ""
                    )

                    patch_file = PatchFile(path=tm.path, patch=og_testfile.to_code())
                    runner_args[func] = {"patch_file": patch_file}

                # 2. run tests in parallel
                results = []
                with ThreadPoolExecutor(max_workers=4) as executor:
                    futures = [
                        (func, executor.submit(self.runner.run_test, **args))
                        for func, args in runner_args.items()
                    ]
                    testgen_logger.info(
                        f"{len(futures)} tests submitted for evaluation"
                    )
                    for func, future in futures:
                        results.append((func, future.result()))

                # 3. loop through test_results and update
                for func, res in results:
                    indvtest_cov, *_ = res

                    indv_improve = indvtest_cov.coverage - del_cov.coverage
                    if indv_improve.total_cov.covered > 0:
                        testgen_logger.info(f"[IMPROVE:{indv_improve.total_cov.covered}] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        improved_tests.append((func, indv_improve))
                    else:
                        testgen_logger.info(f"[NOIMPROVE] Generated Func: {func.name}")
                        testgen_logger.info(f"Code: \n{func.to_code()}")

                        noimprov_tests.append((func, TestCoverage([])))

                # MOVE into post analysis
                # log improvements or something
                # tgt_file = tm.targeted_files(base_path=False)[0]
                # tgt_file_after = cov_res.coverage.get_file_cov(tgt_file)
                # if tgt_file_after:
                #     logger.info(
                #         f"Target file improvement: {tgt_file_after.covered if tgt_file_after else 0}/{tgt_cov}"
                #     )
                # logger.info(
                #     f"Improvement: {cov_res.coverage.total_cov.covered}/{total_cov}"
                # )

        return improved_tests, failed_tests, noimprov_tests


--------------------------------------------------------------------------------

Chunk ID: evaluators/eval_base.py::1
from cowboy_lib.coverage import CoverageResult, TestError, TestCoverage
from cowboy_lib.repo.repository import PatchFile
from cowboy_lib.repo.source_file import Function, TestFile

from src.runner.service import run_test, RunServiceArgs
from src.logger import testgen_logger
from typing import Tuple, List, TYPE_CHECKING

if TYPE_CHECKING:
    from test_gen.augment_test.types import StratResult
    from cowboy_lib.test_modules import TestModule
    from cowboy_lib.repo.source_repo import SourceRepo

from pathlib import Path
from abc import ABC, abstractmethod


--------------------------------------------------------------------------------

Chunk ID: evaluators/eval_base.py::2
class Evaluator(ABC):
    def __init__(
        self, repo_name: str, src_repo: "SourceRepo", run_args: RunServiceArgs
    ):
        self.repo_name = repo_name
        self.src_repo = src_repo
        self.run_args = run_args

    async def gen_test_and_diff_coverage(
        self,
        strat_results: List["StratResult"],
        base_cov: TestCoverage,
        test_fp: Path,
        n_times: int = 1,
    ) -> List[Tuple[CoverageResult, TestCoverage]]:
        """
        Does two runs:
        1. Run to get coverage baseline
        2. Run with generated test case
        Return diff in coverage, and generated test case
        """
        test_results = []
        total_cost = 0

        # WARNING: for some reason failures here are not recorded fully for every new individual
        # that is generate. Possibly due to failures cascading?
        for i, (test_file, test_funcs) in enumerate(strat_results, start=1):
            patch_file = PatchFile(path=test_fp, patch=test_file)
            cov_ptched = await run_test(
                self.repo_name, self.run_args, patch_file=patch_file
            )
            cov_diff = cov_ptched.coverage - base_cov
            # TODO: this covered number is off check
            # NEWTODO: why can this number be negative?
            testgen_logger.info(
                f"New coverage from generated tests: {cov_diff.total_cov.covered}"
            )
            test_results.append((cov_ptched, cov_diff, test_file))

        return test_results


--------------------------------------------------------------------------------

Chunk ID: evaluators/eval_base.py::3
class Evaluator(ABC):

    def get_new_funcs(
        self,
        new_testfile: str,
        test_fp: Path,
    ) -> List[Function]:
        """
        Get newly generated functions
        """
        new_testfile = TestFile(lines=new_testfile.split("\n"), path=str(test_fp))
        old_testfile: TestFile = self.src_repo.find_file(test_fp)
        new_funcs = old_testfile.new_test_funcs(new_testfile)

        return new_funcs

    @abstractmethod
    async def __call__(
        self,
        llm_results: List["StratResult"],
        tm: "TestModule",
        base_cov: CoverageResult,
        n_times: int = 1,
    ):
        raise NotImplementedError

    @abstractmethod
    async def process_test_results(
        self,
        test_results: List[Tuple[CoverageResult, str]],
        tm: "TestModule",
        del_cov: CoverageResult,
    ) -> Tuple[
        List[Tuple[Function, TestCoverage]],
        List[Tuple[Function, TestError]],
        List[Function],
    ]:
        raise NotImplementedError


--------------------------------------------------------------------------------

Chunk ID: evaluators/types.py::1
from enum import Enum


class EvaluatorType(Enum):
    PARALLEL = "PARALLEL"
    ADDITIVE = "ADDITIVE"


--------------------------------------------------------------------------------

Chunk ID: strats/__init__.py::1
from .augment_strat import AugmentClassStrat
from .augment_with_ctxt_file import AugmentClassWithCtxtStrat
from .augment_base import AugmentTestStrategy
from .augment_with_missing import AugmentModuleMissing

from enum import Enum

AUGMENT_STRATS = {
    "VANILLA": AugmentClassStrat,
    "WITH_CTXT": AugmentClassWithCtxtStrat,
    "MODULE_MISSING": AugmentModuleMissing,
}


class AugmentStratType(Enum):
    VANILLA = "VANILLA"
    WITH_CTXT = "WITH_CTXT"
    MODULE_MISSING = "MODULE_MISSING"


--------------------------------------------------------------------------------

Chunk ID: strats/augment_base.py::1
from ..base_strat import BaseStrategy

from dataclasses import dataclass
from typing import List, TYPE_CHECKING

if TYPE_CHECKING:
    from cowboy_lib.test_modules import TestModule

from logging import getLogger

logger = getLogger("test_results")


@dataclass
class LLMResAppend:
    start: int
    lines: List[str]

    def __post_init__(self):
        self.start = int(self.start)

    def __repr__(self):
        return "\n".join(self.lines)


class AugmentTestStrategy(BaseStrategy):
    """
    Generates a test case for a test module
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # "cast" TestCaseInput to TestModule
        self.test_module: "TestModule" = self.test_input

        self.test_file = self.src_repo.find_file(
            self.test_module.test_file.path,
        )
        # this is weird .. we should just make it an arg on BaseStrategy
        self.target_cov = kwargs.get("target_cov", None)
        self.failed_index = 0


--------------------------------------------------------------------------------

Chunk ID: strats/augment_base.py::2
class AugmentTestStrategy(BaseStrategy):

    def parse_llm_res(self, llm_res: str) -> str:
        lines = llm_res.split("\n")
        llm_out = LLMResAppend(int(lines[0]), lines[1:])

        logger.debug(f"Generated code: \n{llm_out.__repr__()}")

        new_test_file = self.test_module.test_file.clone()
        new_test_file.append(
            "\n".join(llm_out.lines),
            class_name=self.test_module.name if self.test_module._isclass else "",
        )

        return new_test_file.to_code()


--------------------------------------------------------------------------------

Chunk ID: strats/augment_strat.py::1
from .augment_base import AugmentTestStrategy
from .prompt import AugmentTestPrompt
from ..utils import gen_enumerated_code_str, get_current_git_commit


from logging import getLogger


logger = getLogger("test_results")


class AugmentClassStrat(AugmentTestStrategy):
    """
    Just simply executes the LLM prompt without providing additional
    context
    """

    def build_prompt(self) -> str:
        curr_commit = get_current_git_commit(self.src_repo.repo_path)

        prompt = AugmentTestPrompt()

        test_code = self.test_module.get_test_code(curr_commit)
        # test_file = self.repo_ctxt.src_repo.get_file(self.test_module.test_file.path)
        # test_code = self.get_test_code(test_file, self.test_module.nodes)

        logger.info(f"ADDITIVE TEST CODE: {test_code}")

        test_code = gen_enumerated_code_str(test_code.split("\n"))
        prompt.insert_line("test_code", test_code)

        return prompt.get_prompt()

    def get_test_code(self, test_file, nodes):
        test_code = ""
        for node in nodes:
            try:
                test_code += test_file.find_by_nodetype(
                    node.name, node_type=node.node_type
                ).to_code()
            except Exception as e:
                logger.error(f"Error: {e}")
                continue

        return test_code


--------------------------------------------------------------------------------

Chunk ID: strats/augment_with_ctxt_file.py::1
from .augment_base import AugmentTestStrategy
from .prompt import AugmentTestPromptWithCtxt
from ..types import CtxtWindowExceeded
from ..utils import get_current_git_commit

from typing import TYPE_CHECKING

from src.logger import testgen_logger


class AugmentClassWithCtxtStrat(AugmentTestStrategy):
    """
    Provides the src code context from the source code file targeted by the
    test_module
    """

    def build_prompt(self) -> AugmentTestPromptWithCtxt:
        prompt = AugmentTestPromptWithCtxt()
        curr_commit = get_current_git_commit(self.src_repo.repo_path)

        test_code = self.test_module.get_test_code(curr_commit)
        test_fit = prompt.insert_line("test_code", test_code)
        if not test_fit:
            raise CtxtWindowExceeded("Test code too large to fit in prompt")

        # TODO: how to narrow the scope of this to class or even function
        # have ref to func/class node in TargetCode
        for fp in self.test_module.targeted_files():
            testgen_logger.info(f"Src file ctxt of {self.test_module.name}: {fp}")
            file = self.src_repo.find_file(fp)
            code_fit = prompt.insert_line("file_contents", file.to_code())

            if not code_fit:
                testgen_logger.warn(f"File {fp} too large to fit in prompt")
                continue

        if not self.test_module.targeted_files():
            testgen_logger.warn(
                f"No src files targeted by {self.test_module.name} found")

        return prompt.get_prompt()

    def get_test_code(self, test_file, nodes):
        test_code = ""
        for node in nodes:
            try:
                test_code += test_file.find_by_nodetype(
                    node.name, node_type=node.node_type
                ).to_code()
            except Exception as e:
                testgen_logger.error(f"Error: {e}")
                continue

        return test_code


--------------------------------------------------------------------------------

Chunk ID: strats/augment_with_missing.py::1
from .augment_base import AugmentTestStrategy
from .prompt import AugmentTestPromptMiss

from ..types import CtxtWindowExceeded
from ..utils import gen_enumerated_code_str, get_current_git_commit

from logging import getLogger

logger = getLogger("test_results")


class AugmentModuleMissing(AugmentTestStrategy):
    """
    Augment test by focusing on missing/uncovered lines
    """

    def build_prompt(self) -> AugmentTestPromptMiss:
        if not self.target_cov:
            raise ValueError("Target coverage not set")

        prompt = AugmentTestPromptMiss()

        curr_commit = get_current_git_commit(self.src_repo.repo_path)
        test_code = self.test_module.get_test_code(curr_commit).split("\n")
        test_code = gen_enumerated_code_str(test_code)
        missing_lines = self.target_cov.print_lines(line_type="missing")

        logger.info(f"Missing lines: {missing_lines}")

        test_fit = prompt.insert_line("test_code", test_code)
        if not test_fit:
            raise CtxtWindowExceeded("Test code too large to fit in prompt")

        for fp in self.test_module.targeted_files():
            file = self.src_repo.find_file(fp)
            code_fit = prompt.insert_line("file_contents", file.to_code())
            if not code_fit:
                logger.warn(f"File {fp} too large to fit in prompt")
                continue

        prompt.insert_line("missing_lines", missing_lines)
        return prompt.get_prompt()


--------------------------------------------------------------------------------

Chunk ID: strats/prompt.py::1
from ..types import Prompt, LMModelSpec
from typing import List


gpt4_spec = LMModelSpec("gpt-4", 0.01, 8192)


class AugmentTestPrompt(Prompt):
    def __init__(self):
        prompt = """The following is a unit test case written for a code feature
{{test_code}}

Extend the unit test class to increase coverage, by appending your new code into the existing code above.
Here is the structure of your output.
On first line, write the insertion line that your new code is going to be appended to
Then on the following lines, generate the code to be inserted. Keep the indentation consistent
For example, given the following input

0. class FakeClass:
1.    def fake_method(self):
2.       pass

The following output was generated
2

    def test_fake_method(self):
        pass

Now, generate your response
"""

        super().__init__(prompt, gpt4_spec, ["test_code"])


--------------------------------------------------------------------------------

Chunk ID: strats/prompt.py::2
class AugmentTestPromptWithCtxt(Prompt):
    def __init__(self):
        prompt = """The following is a unit test case written for a code feature. 
{{test_code}}

{% if file_contents %}
Here is the source code file that the test intends to cover:
{{file_contents}}
{% endif %}

Extend the unit test class to increase coverage, by appending your new code into the existing code above.
Here is the structure of your output.
On first line, write the insertion line that your new code is going to be appended to
Then on the following lines, generate the code to be inserted. Keep the indentation consistent
For example, given the following input

0. class FakeClass:
1.    def fake_method(self):
2.       pass

The following output was generated
2

    def test_fake_method(self):
        pass

Now, generate your response
"""
        super().__init__(prompt, gpt4_spec, ["test_code", "file_contents"])


--------------------------------------------------------------------------------

Chunk ID: strats/prompt.py::3
class AugmentTestPromptMiss(Prompt):
    def __init__(self):
        prompt = """The following is a unit test case written for a code feature. 
{{test_code}}

{% if file_contents %}
Here is the source code file that the test intends to cover:
{{file_contents}}
{% endif %}

{% if missing_lines %}
Here are the missing lines that the augmented test should cover:
{{missing_lines}}
{% endif %}

Extend the unit test class to increase coverage, by appending your new code into the existing code above.
Here is the structure of your output.
On first line, write the insertion line that your new code is going to be appended to
Then on the following lines, generate the code to be inserted. Keep the indentation consistent
For example, given the following input

0. class FakeClass:
1.    def fake_method(self):
2.       pass

The following output was generated
2

    def test_fake_method(self):
        pass

Now, generate your response
"""
        super().__init__(
            prompt, gpt4_spec, ["test_code", "file_contents", "missing_lines"]
        )


--------------------------------------------------------------------------------

Chunk ID: strats/types.py::1
from enum import Enum


class AugmentStratType(Enum):
    CLASS = "CLASS"
    CLASS_WITH_CTXT = "CLASS_WITH_CTXT"
    MODULE_MISSING = "MODULE_MISSING"


--------------------------------------------------------------------------------

Chunk ID: augment_test/types.py::1
from dataclasses import dataclass
from collections import defaultdict
from jinja2 import Template

from typing import List, NamedTuple
import tiktoken

from enum import Enum, auto


class StratResult(NamedTuple):
    contents: str
    test_path: str


class StrategyInitError(Exception):
    pass


class StrategyMode(Enum):
    GEN_NEW_TEST = auto()
    AUGMENT_TEST = auto()


@dataclass
class LMModelSpec:
    model: str
    cost: float
    ctxt_window: int


class CtxtWindowExceeded(Exception):
    pass


def num_tokens_from_string(string: str, encoding_name: str) -> int:
    """Returns the number of tokens in a text string."""
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens


--------------------------------------------------------------------------------

Chunk ID: augment_test/types.py::2
class Prompt:
    def __init__(self, prompt: str, model_spec: LMModelSpec, keywords: List[str]):
        self._prompt = prompt
        self._keyword_dict = defaultdict(lambda: "")
        self._keywords = keywords
        self._max_len = model_spec.ctxt_window
        self._model_name = model_spec.model
        self._current_len = self._tokenize(prompt)

    def _tokenize(self, prompt: str):
        # Fails with: ValueError: Unknown encoding <Encoding 'cl100k_base'>
        # encoding = tiktoken.encoding_for_model(self._model_name)
        return num_tokens_from_string(prompt, "cl100k_base")

    def insert_line(self, keyword: str, line: str, preamble: str = "") -> bool:
        if keyword not in self._keywords:
            raise ValueError(f"Keyword {keyword} not found in prompt")

        str_len = self._tokenize(line)
        if self._current_len + str_len > self._max_len:
            return False
        else:
            self._keyword_dict[keyword] += preamble + "\n" + line + "\n"

        return True

    def get_prompt(self) -> str:
        template = Template(self._prompt)
        return template.render(**self._keyword_dict)


--------------------------------------------------------------------------------

Chunk ID: augment_test/utils.py::1
from typing import List
from pathlib import Path
import subprocess


def gen_enumerated_code_str(code: List[str]) -> str:
    """
    Adds line numbers infront of source code like so:

    1: def test_something():
    2:     assert True

    """

    code.append("\n")
    # strip out newlines from code, idk how they got there
    code = [line.rstrip() for line in code]
    return "\n".join([f"{i}: {line}" for i, line in enumerate(code, start=1)])


--------------------------------------------------------------------------------

Chunk ID: augment_test/utils.py::2
def get_current_git_commit(repo_path: Path) -> str:
    """
    Uses subprocess to get the current git commit hash.

    Returns:
        str: The current git commit hash.
    """
    try:
        commit_hash = (
            subprocess.check_output(
                ["cd", str(repo_path.resolve()), "&&", "git", "rev-parse", "HEAD"],
                shell=True,
            )
            .strip()
            .decode("utf-8")
        )
        return commit_hash
    except subprocess.CalledProcessError as e:
        print(f"Error getting current git commit: {e}")
        return ""


--------------------------------------------------------------------------------

Chunk ID: test_gen/models.py::1
from src.database.core import Base
from src.models import CowboyBase

from sqlalchemy import Column, Integer, String, ForeignKey
from sqlalchemy.orm import relationship
from pydantic import BaseModel, validator, model_validator
from enum import Enum
from typing import List
from typing import Optional


class TMSelectMode(str, Enum):
    """
    Used to select the TestModules to be augmented
    """

    AUTO = "auto"
    FILE = "file"
    TM = "module"
    ALL = "all"


class Decision(int, Enum):
    YES = 1
    NO = 0
    UNDECIDED = -1


--------------------------------------------------------------------------------

Chunk ID: test_gen/models.py::2
class TMSelectModeBase(BaseModel):
    mode: TMSelectMode
    files: Optional[List[str]]
    tms: Optional[List[str]]

    @validator("files", always=True)
    def check_files(cls, v, values):
        if values.get("mode") == TMSelectMode.FILE and not v:
            raise ValueError("files must be specified if mode is FILE")
        return v

    @validator("tms", always=True)
    def check_tms(cls, v, values):
        if values.get("mode") == TMSelectMode.TM and not v:
            raise ValueError("tms must be specfied if mode is TM")
        return v


class AugmentTestRequest(TMSelectModeBase):
    repo_name: str


class AugmentTestResponse(CowboyBase):
    session_id: str


--------------------------------------------------------------------------------

Chunk ID: test_gen/models.py::3
class AugmentTestResult(Base):
    __tablename__ = "augment_test_results"
    id = Column(Integer, primary_key=True)
    name = Column(String)
    test_case = Column(String)
    decide = Column(Integer, default=-1)
    commit_hash = Column(String)

    # used in TestFile.append() to construct the modified test file
    testfile = Column(String)
    classname = Column(String, nullable=True)
    session_id = Column(String)

    repo_id = Column(Integer, ForeignKey("repo_config.id"))
    # TODO: not tested yet
    test_module_id = Column(Integer, ForeignKey("test_modules.id"))
    cov_list = relationship("CoverageModel", cascade="all, delete-orphan")

    def set_decision(self, decision: Decision):
        self.decision = decision.value

    def coverage_improve(self):
        return sum([cov.covered for cov in self.cov_list])


--------------------------------------------------------------------------------

Chunk ID: test_gen/models.py::4
class TestResultResponse(BaseModel):
    id: str
    name: str
    test_case: str
    test_file: str
    cov_improved: int
    decided: int


class UserDecision(BaseModel):
    id: int
    decision: Decision


class UserDecisionRequest(BaseModel):
    user_decision: List[UserDecision]

    @model_validator(mode="before")
    def check_user_decision(cls, values):
        if not values.get("user_decision"):
            raise ValueError("user_decision must not be empty")
        return values


class UserDecisionResponse(BaseModel):
    compare_url: str


--------------------------------------------------------------------------------

Chunk ID: test_gen/service.py::1
from cowboy_lib.repo import SourceRepo
from src.test_modules.service import (
    get_all_tms,
    get_tms_by_filename,
    get_tms_by_names,
)
from src.config import AUTO_GEN_SIZE

from .models import AugmentTestResult, TMSelectModeBase, TMSelectMode

from starlette.requests import Request
from typing import List


def save_all(*, db_session, test_results: List[AugmentTestResult]):
    for tr in test_results:
        db_session.add(tr)
    db_session.commit()


--------------------------------------------------------------------------------

Chunk ID: test_gen/service.py::2
def create_test_result(
    *,
    db_session,
    repo_id,
    name,
    test_case,
    cov_list,
    tm_id,
    commit_hash,
    testfile,
    session_id,
    classname=None,
):
    tr_model = AugmentTestResult(
        name=name,
        test_case=test_case,
        test_module_id=tm_id,
        commit_hash=commit_hash,
        testfile=testfile,
        classname=classname,
        session_id=session_id,
        repo_id=repo_id,
    )

    db_session.add(tr_model)
    db_session.commit()

    return tr_model


--------------------------------------------------------------------------------

Chunk ID: test_gen/service.py::3
def get_test_result_by_id_or_raise(*, db_session, test_id) -> AugmentTestResult:
    return (
        db_session.query(AugmentTestResult)
        .filter(AugmentTestResult.id == test_id)
        .one_or_none()
    )


def get_test_results_by_sessionid(*, db_session, session_id) -> AugmentTestResult:
    return (
        db_session.query(AugmentTestResult)
        # only return undecided sessions
        .filter(
            AugmentTestResult.session_id == session_id, AugmentTestResult.decide == -1
        ).all()
    )


def update_test_result_decision(*, db_session, test_id, decision: int):
    test_result = get_test_result_by_id_or_raise(db_session=db_session, test_id=test_id)
    test_result.decide = decision
    db_session.commit()


def delete_test_results_by_sessionid(*, db_session, session_id):
    db_session.query(AugmentTestResult).filter(
        AugmentTestResult.session_id == session_id
    ).delete()

    db_session.commit()


def get_session_id(request: Request):
    return request.state.session_id


--------------------------------------------------------------------------------

Chunk ID: test_gen/service.py::4
def select_tms(*, db_session, repo_id, request: TMSelectModeBase, src_repo: SourceRepo):
    if request.mode == TMSelectMode.AUTO.value:
        tm_models = get_all_tms(
            db_session=db_session, repo_id=repo_id, n=AUTO_GEN_SIZE
        )
    elif request.mode == TMSelectMode.FILE.value:
        tm_models = get_tms_by_filename(
            db_session=db_session, repo_id=repo_id, src_file=request.files
        )
    elif request.mode == TMSelectMode.TM.value:
        tm_models = get_tms_by_names(
            db_session=db_session, repo_id=repo_id, tm_names=request.tms
        )
    elif request.mode == TMSelectMode.ALL.value:
        tm_models = get_tms_by_names(
            db_session=db_session, repo_id=repo_id, tm_names=[]
        )

    return tm_models


--------------------------------------------------------------------------------

Chunk ID: test_gen/utils.py::1
from .models import AugmentTestResult

from collections import defaultdict
from typing import List


def gen_commit_msg(test_results: List[AugmentTestResult]):
    summary = defaultdict(int)
    for tr in test_results:
        summary[tr.testfile] += 1

    return f"COWBOY generated {','.join([f"{f}: {tests}\n" for f, tests in summary.items()])}"


--------------------------------------------------------------------------------

Chunk ID: test_gen/views.py::1
from cowboy_lib.repo import SourceRepo, GitRepo
from src.database.core import get_db
from src.auth.service import get_current_user
from src.stats.service import update_repo_stats
from src.repo.service import get_or_raise, get_by_id_or_raise
from src.queue.core import get_queue

# from src.logger import accepted_count, failed_count, total_count

from .models import (
    AugmentTestRequest,
    AugmentTestResponse,
    UserDecisionRequest,
    UserDecisionResponse,
    TestResultResponse,
)
from .augment import augment_test
from .service import (
    save_all,
    get_test_results_by_sessionid,
    get_test_result_by_id_or_raise,
    update_test_result_decision,
)
from .service import get_session_id, select_tms
from .utils import gen_commit_msg

from sqlalchemy.orm import Session
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException
from functools import reduce
import asyncio

test_gen_router = APIRouter()


MODE_AUTO_ERROR_MSG = """
No test modules found for auto mode. Consider using mode=all instead, if you
do not want to specify individual test modules or files to augment 
"""

MODE_FILES_ERROR_MSG = """
No test files found for files mode. This could be because we dont have the test 
module mapping generated for {filenames}. Consider first build mapping on these
files first. If that still doesnt work, it means that we were unable to recover
the test module to source file mapping for {filenames}
"""


--------------------------------------------------------------------------------

Chunk ID: test_gen/views.py::2
@test_gen_router.post("/test-gen/augment")
async def augment_test_route(
    request: AugmentTestRequest,
    db_session=Depends(get_db),
    session_id=Depends(get_session_id),
    curr_user=Depends(get_current_user),
    task_queue=Depends(get_queue),
):
    """
    Augment tests for a test module
    """
    repo = get_or_raise(
        db_session=db_session, curr_user=curr_user, repo_name=request.repo_name
    )
    src_repo = SourceRepo(Path(repo.source_folder))
    tm_models = select_tms(
        db_session=db_session, repo_id=repo.id, request=request, src_repo=src_repo
    )

    if not tm_models:
        detail = (
            MODE_AUTO_ERROR_MSG
            if request.mode == "auto"
            else MODE_FILES_ERROR_MSG.format(filenames=request.files)
        )
        return HTTPException(status_code=400, detail=detail)

    # TODO: put this into a service
    coroutines = []
    for tm_model in tm_models:
        coroutine = augment_test(
            db_session=db_session,
            task_queue=task_queue,
            repo=repo,
            tm_model=tm_model,
            curr_user=curr_user,
            session_id=session_id,
        )
        coroutines.append(coroutine)

    test_results = await asyncio.gather(*coroutines)
    test_results = reduce(lambda x, y: x + y, test_results)

    # we save here after async ops have finished running
    save_all(db_session=db_session, test_results=test_results)

    return AugmentTestResponse(session_id=session_id)


--------------------------------------------------------------------------------

Chunk ID: test_gen/views.py::3
@test_gen_router.get("/test-gen/results/{session_id}")
def get_results(
    session_id: str,
    db_session: Session = Depends(get_db),
):
    trs = get_test_results_by_sessionid(db_session=db_session, session_id=session_id)
    return [
        TestResultResponse(
            id=str(tr.id),
            name=tr.name,
            test_case=tr.test_case,
            test_file=tr.testfile,
            cov_improved=tr.coverage_improve(),
            decided=tr.decide,
        )
        for tr in trs
    ]


--------------------------------------------------------------------------------

Chunk ID: test_gen/views.py::4
@test_gen_router.post("/test-gen/results/decide/{sesssion_id}")
def accept_user_decision(
    request: UserDecisionRequest,
    curr_user=Depends(get_current_user),
    db_session=Depends(get_db),
):
    """
    Takes the result of the selected tests and appends all of the selected
    tests to TestModule (testfile/test class). Then check out a new branch against
    the remote repo with the changed files
    """

    repo_id = get_test_result_by_id_or_raise(
        db_session=db_session, test_id=request.user_decision[0].id
    ).repo_id
    repo = get_by_id_or_raise(
        db_session=db_session, curr_user=curr_user, repo_id=repo_id
    )
    src_repo = SourceRepo(Path(repo.source_folder))
    git_repo = GitRepo(Path(repo.source_folder))

    # NOTE: LintExceptions at this step should not happen because they would have occurred
    # earlier during the Evaluation phase
    changed_files = set()
    accepted_trs = []
    for decision in request.user_decision:
        tr = get_test_result_by_id_or_raise(db_session=db_session, test_id=decision.id)
        test_file = src_repo.find_file(tr.testfile)
        if decision.decision:
            test_file.append(tr.test_case, class_name=tr.classname)
            src_repo.write_file(test_file.path)
            changed_files.add(str(test_file.path))
            accepted_trs.append(tr)

        update_test_result_decision(
            db_session=db_session, test_id=decision.id, decision=decision.decision
        )

    # update stats
    with update_repo_stats(db_session=db_session, repo=repo) as repo_stats:
        repo_stats.accepted_tests += len(accepted_trs)
        repo_stats.rejected_tests += len(request.user_decision) - len(accepted_trs)

        # update logfire metrics
        # accepted_count = logfire.metric_counter("accepted_results", unit="1")
        # failed_count = logfire.metric_counter("failed_results", unit="1")
        # total_count = logfire.metric_counter("total_results", unit="1")

        # accepted_count.add(len(accepted_trs))
        # failed_count.add(len(request.user_decision) - len(accepted_trs))
        # total_count.add(len(request.user_decision))

    msg = gen_commit_msg(accepted_trs)
    compare_url = git_repo.checkout_and_push(
        "cowboy-augmented-tests", msg, list(changed_files)
    )

    return UserDecisionResponse(compare_url=compare_url)


--------------------------------------------------------------------------------

Chunk ID: test_modules/iter_tms.py::1
from cowboy_lib.utils import get_current_git_commit
from cowboy_lib.repo.source_repo import SourceRepo

from src.test_modules.models import TestModule

from typing import List


def iter_test_modules(src_repo: SourceRepo, filter_fn=None) -> List[TestModule]:
    """
    Generator for TestModules
    TestModules can be either be:
    1. All the individual functions inside a TestFile
    2. All the functions inside a class inside a TestFile
    3. Some of the individual functions inside a TestFile
    4. Some of the functions inside a class inside a TestFile
    
    Args:
        src_repo: Source repository to scan for test modules
        filter_fn: Optional function to filter test modules, takes TestModule and returns bool
    """
    test_modules: List[TestModule] = []
    for test_file in src_repo.test_files:
        ind_funcs = [f for f in test_file.test_funcs() if not f.scope]
        if ind_funcs:
            func_module = TestModule(
                test_file, ind_funcs, get_current_git_commit(src_repo.repo_path)
            )
            test_modules.append(func_module)

        for test_class in test_file.test_classes():
            class_module = TestModule(
                test_file, [test_class], get_current_git_commit(src_repo.repo_path)
            )
            test_modules.append(class_module)

    # NOTE: literally dont work and literally dont need
    # name_counter = Counter([tm.name for tm in test_modules])
    # while any([count > 1 for count in name_counter.values()]):
    #     for tm in test_modules:
    #         if name_counter[tm.name] > 1:
    #             tm.name = get_new_name(tm.name, tm.test_file.path)
    #             name_counter = Counter([tm.name for tm in test_modules])

    if filter_fn:
        test_modules = [tm for tm in test_modules if filter_fn(tm)]

    return test_modules


--------------------------------------------------------------------------------

Chunk ID: test_modules/models.py::1
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
from sqlalchemy.orm import relationship
from pathlib import Path
from typing import List, Optional
from pydantic import BaseModel

from cowboy_lib.test_modules.test_module import TestModule
from cowboy_lib.repo.source_repo import SourceRepo

from src.database.core import Base
from src.ast.models import NodeModel
from src.target_code.models import TargetCodeModel
from src.test_gen.models import TMSelectModeBase, AugmentTestResult


class IncompatibleCommit(Exception):
    pass


class TestModuleModel(Base):
    __tablename__ = "test_modules"
    id = Column(Integer, primary_key=True)
    name = Column(String)
    testfilepath = Column(String)
    # srcpaths = Column(String)
    commit_sha = Column(String)
    experiment_id = Column(String, nullable=True)
    # use this flag to track test_modules that have already gone through
    # auto-test augmentation
    auto_gen = Column(Boolean, default=False)

    repo_id = Column(Integer, ForeignKey("repo_config.id"))
    nodes = relationship(
        "NodeModel",
        backref="test_module",
        foreign_keys=[NodeModel.test_module_id],
        cascade="all, delete-orphan",
    )
    target_chunks = relationship(
        "TargetCodeModel",
        backref="test_module",
        foreign_keys=[TargetCodeModel.test_module_id],
        cascade="all",
    )
    test_results = relationship(
        "AugmentTestResult",
        backref="test_module",
        foreign_keys=[AugmentTestResult.test_module_id],
        cascade="all, delete-orphan",
    )

    # def __init__(
    #         self,
    #         name: str,
    #         testfilepath: str,
    #         srcpaths: List[str],
    #         commit_sha: str,
    #         repo_id: int,
    #         experiment_id: Optional[str] = None,
    #         auto_gen: bool = False,
    # ):
    #     self.name = name
    #     self.testfilepath = testfilepath
    #     self.srcpaths = ",".join(srcpaths)
    #     self.commit_sha = commit_sha
    #     self.repo_id = repo_id
    #     self.experiment_id = experiment_id
    #     self.auto_gen = auto_gen

    # @property
    # def srcpaths_list(self) -> List[str]:
    #     return self.srcpaths.split(",") if self.srcpaths else []


    def serialize(self, src_repo: SourceRepo) -> TestModule:
        """
        Convert model back to TestModule
        """
        return TestModule(
            test_file=src_repo.find_file(Path(self.testfilepath)),
            commit_sha=self.commit_sha,
            nodes=[NodeModel.to_astnode(n, src_repo) for n in self.nodes],
            chunks=[c.serialize(src_repo) for c in self.target_chunks],
        )

    def get_covered_files(self) -> List[str]:
        """
        Returns the source files that are covered by this test module
        """

        # there must be a better way of doing this ...
        return list(set([chunk.filepath for chunk in self.target_chunks]))


--------------------------------------------------------------------------------

Chunk ID: test_modules/models.py::2
class TestModuleModel(Base):

    def score(self, filename: str, src_repo: SourceRepo) -> int:
        """
        Get score for a single file
        """
        if filename not in self.get_covered_files():
            raise Exception("Filename is not covered by TM")

        tgt_code_chunks = [
            chunk for chunk in self.target_chunks if chunk.filepath == filename
        ]
        # coverage same for all tgt_code_chunks
        cov = tgt_code_chunks[0].coverage
        file = src_repo.find_file(tgt_code_chunks[0].filepath)
        total_lines = len(file.lines)

        total_covered = cov.covered
        total_missing = cov.misses
        chunk_covered = 0

        for chunk in tgt_code_chunks:
            chunk_covered += len(chunk.get_lines())

        return chunk_covered + total_missing / total_lines if total_lines else 0

    def agg_score(self, src_repo: SourceRepo) -> int:
        """
        Get aggregate score over all covered source files
        """
        agg_score = 0
        all_files = self.get_covered_files()

        for f in all_files:
            agg_score += self.score(f, src_repo)

        return agg_score / len(all_files) if len(all_files) > 0 else 0


class BuildMappingRequest(TMSelectModeBase):
    repo_name: str
    overwrite: Optional[bool] = False


class TestModuleReponse(BaseModel):
    filepath: str
    name: str
    unit_tests: List[str]


--------------------------------------------------------------------------------

Chunk ID: test_modules/service.py::1
from cowboy_lib.repo import SourceRepo

from src.repo.models import RepoConfig
from src.ast.service import create_node

from .iter_tms import iter_test_modules
from .models import TestModuleModel, TestModule

from sqlalchemy.orm import Session
from typing import List


def create_all_tms(*, db_session: Session, repo_conf: RepoConfig, src_repo: SourceRepo):
    """Create all test modules for a repo."""
    test_modules = iter_test_modules(src_repo)

    for tm in test_modules:
        create_tm(db_session=db_session, repo_id=repo_conf.id, tm=tm)


--------------------------------------------------------------------------------

Chunk ID: test_modules/service.py::2
def create_tm(*, db_session: Session, repo_id: str, tm: TestModule):
    """Create a test module and the nodes"""

    tm_model = TestModuleModel(
        name=tm.name,
        testfilepath=str(tm.test_file.path),
        commit_sha=tm.commit_sha,
        repo_id=repo_id,
    )

    # need to commit before so node has access to tm_model.id
    db_session.add(tm_model)
    db_session.commit()

    for node in tm.nodes:
        create_node(
            db_session=db_session,
            node=node,
            repo_id=repo_id,
            filepath=tm_model.testfilepath,
            test_module_id=tm_model.id,
        )

    return tm_model


--------------------------------------------------------------------------------

Chunk ID: test_modules/service.py::3
def get_tm_by_name(
    *, db_session: Session, repo_id: str, tm_name: str
) -> TestModuleModel:
    """
    Query by name and return all if no names are provided
    """

    query = db_session.query(TestModuleModel).filter(TestModuleModel.repo_id == repo_id)
    if tm_name:
        query = query.filter(TestModuleModel.name == tm_name)

    return query.one_or_none()


def get_tms_by_names(
    *, db_session: Session, repo_id: str, tm_names: List[str]
) -> List[TestModuleModel]:
    """
    Query by name and return all if no names are provided
    """
    if tm_names == []:
        return get_all_tms(db_session=db_session, repo_id=repo_id, n=None)

    query = db_session.query(TestModuleModel).filter(TestModuleModel.repo_id == repo_id)
    if tm_names:
        query = query.filter(TestModuleModel.name.in_(tm_names))

    return query.all()


--------------------------------------------------------------------------------

Chunk ID: test_modules/service.py::4
def update_tm(*, db_session: Session, tm_model: TestModuleModel):
    """
    Updates an existing TM
    """
    db_session.merge(tm_model)
    db_session.commit()

    return tm_model


def get_tms_by_filename(
    *, db_session: Session, repo_id: str, src_file: str
) -> List[TestModuleModel]:
    """
    Query all TMs for a repo
    """
    all_tms = get_all_tms(db_session=db_session, repo_id=repo_id, n=None)
    return [tm for tm in all_tms if src_file in tm.get_covered_files()]


--------------------------------------------------------------------------------

Chunk ID: test_modules/service.py::5
# TODO: need to figure out when to set this flag to False
def get_all_tms(
    *, db_session: Session, repo_id: str, n: int = 2
) -> List[TestModuleModel]:
    """
    Query all TMs for a repo
    """
    all_tms = db_session.query(TestModuleModel).filter(TestModuleModel.repo_id == repo_id)
    # sort by id to get deterministic ordering
    sorted_tms = sorted(all_tms, key=lambda tm: tm.id, reverse=True)
    select_tms = sorted_tms[:n]

    # update auto_gen flag on TMs
    for tm in select_tms:
        tm.auto_gen = True
        update_tm(db_session=db_session, tm_model=tm)

    return select_tms


--------------------------------------------------------------------------------

Chunk ID: test_modules/views.py::1
from fastapi import APIRouter, Depends

from cowboy_lib.repo import SourceRepo

from src.database.core import get_db
from src.queue.core import TaskQueue, get_queue
from src.auth.service import get_current_user, CowboyUser
from src.repo.service import get_or_raise
from src.test_gen.service import select_tms

from src.models import HTTPSuccess
from src.tasks.create_tgt_coverage import create_tgt_coverage
from src.runner.models import ClientRunnerException
from src.runtest_conf import run_test

from .models import BuildMappingRequest, TestModuleReponse
from .service import get_all_tms

from sqlalchemy.orm import Session
from pathlib import Path
from typing import List

from src.logger import testgen_logger as log


tm_router = APIRouter()


--------------------------------------------------------------------------------

Chunk ID: test_modules/views.py::2
@tm_router.post("/tm/build-mapping")
async def get_tm_target_coverage(
    request: BuildMappingRequest,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
    task_queue: TaskQueue = Depends(get_queue),
):
    repo = get_or_raise(
        db_session=db_session, curr_user=current_user, repo_name=request.repo_name
    )
    src_repo = SourceRepo(Path(repo.source_folder))
    tm_models = select_tms(
        db_session=db_session, repo_id=repo.id, request=request, src_repo=src_repo
    )
    # tm_models = [tm_model for tm_model in tm_models if not tm_model.target_chunks]

    try:
        await create_tgt_coverage(
            db_session=db_session,
            task_queue=task_queue,
            repo=repo,
            tm_models=tm_models,
            run_test=run_test,
            overwrite=request.overwrite,
        )
        return HTTPSuccess()

    except ClientRunnerException as e:
        raise e


--------------------------------------------------------------------------------

Chunk ID: test_modules/views.py::3
@tm_router.get("/tm/{repo_name}", response_model=List[TestModuleReponse])
def get_tms(
    repo_name: str,
    db_session: Session = Depends(get_db),
    current_user: CowboyUser = Depends(get_current_user),
):
    repo = get_or_raise(
        db_session=db_session, curr_user=current_user, repo_name=repo_name
    )
    src_repo = SourceRepo(Path(repo.source_folder))
    tms = [
        tm.serialize(src_repo)
        for tm in get_all_tms(db_session=db_session, repo_id=repo.id)
    ]

    return [
        TestModuleReponse(
            filepath=str(tm.path),
            name=tm.name,
            unit_tests=[ut.name for ut in tm.tests],
        )
        for tm in tms
    ]


--------------------------------------------------------------------------------

Chunk ID: src/token_registry.py::1
token_registry = set()


--------------------------------------------------------------------------------

Chunk ID: src/utils.py::1
import functools
import random
import string
import uuid
import time
import functools
from src.logger import testgen_logger


# nested level get() function
def resolve_attr(obj, attr, default=None):
    """Attempts to access attr via dotted notation, returns none if attr does not exist."""
    try:
        return functools.reduce(getattr, attr.split("."), obj)
    except AttributeError:
        return default


def gen_random_name():
    """
    Generates a random name using ASCII, 8 characters in length
    """

    return "".join(random.choices(string.ascii_lowercase, k=8))


def generate_id():
    """
    Generates a random UUID
    """
    return str(uuid.uuid4())


def async_timed(func):
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        start_time = time.time()
        result = await func(*args, **kwargs)
        end_time = time.time()
        testgen_logger.info(
            f"Function {func.__name__} took {end_time - start_time:.4f} seconds"
        )
        return result

    return wrapper


def sync_timed(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(
            f"Function {func.__name__} took {end_time - start_time:.4f} seconds"
        )
        return result

    return wrapper


--------------------------------------------------------------------------------

Chunk ID: test.py::1
# import sys

# sys.path.append("/home/ubuntu/cowboy-server-good")

from cowboy_lib.repo import SourceRepo
from src.test_modules.service import get_tm_by_name, get_all_tms

from src.database.core import engine

from sqlalchemy.orm import sessionmaker
from pathlib import Path


repo_path = "/home/ubuntu/cowboy-server-good/repos/test2/qrjmnlxt"
src_repo = SourceRepo(Path(repo_path))
Session = sessionmaker(bind=engine)
db_session = Session()


# tm_model = get_tm_by_name(db_session=db_session, repo_id=17, tm_name="TestWoodpecker")
# tm = tm_model.serialize(src_repo)

tm_models = get_all_tms(db_session=db_session, repo_id=17)
tm_models = sorted(tm_models, key=lambda tm: tm.agg_score(src_repo), reverse=True)

for tm in tm_models:
    print(tm.name, tm.agg_score(src_repo))


--------------------------------------------------------------------------------

Chunk ID: test_compare.py::1
from cowboy_lib.coverage import TestCoverage
import json

BASECOV_PATH = "basecov_textual.json"
MODULECOV_PATH = "modulecov_textual.json"


def compare(a: TestCoverage, b: TestCoverage):
    for t1, t2 in zip(a.cov_list, b.cov_list):
        assert t1.filename == t2.filename
        if t1.stmts != t2.stmts:
            print(f"Diff in {t1.filename}: stmts")
            print(f"Base: {t1} || {t1.covered + t1.misses},{t1.stmts}")
            print(f"Module: {t2} || {t2.covered + t2.misses},{t2.stmts}")

            # CLUE #1: this always matches
            # if t1.covered + t1.misses != t1.stmts or t2.covered + t2.misses != t2.stmts:
            #     raise Exception(f"Total stmts mismatch for {t1.filename}")

            if t1.covered < t2.covered:
                raise Exception(f"Coverage decreased for {t1.filename}")


--------------------------------------------------------------------------------

Chunk ID: test_compare.py::2
def test_file_diff():
    with open(BASECOV_PATH) as f:
        base_cov = TestCoverage.from_coverage_file(json.load(f))
    with open(MODULECOV_PATH) as f:
        module_cov = TestCoverage.from_coverage_file(json.load(f))

    print("\nBase Cov: ", base_cov)
    print("Module Cov: ", module_cov)

    compare(base_cov, module_cov)

    cov_diff = base_cov - module_cov
    # print(cov_diff)

    # assert cov_diff.total_cov.stmts == 33
    # assert cov_diff.total_cov.misses == 11
    # assert cov_diff.total_cov.covered == 22

test_file_diff()


--------------------------------------------------------------------------------

Chunk ID: test_eval.py::1
from src.runner.local.run_test import get_repo_config
import asyncio
from pathlib import Path
from src.tasks.get_baseline_parallel import get_tm_target_coverage
from src.test_modules.iter_tms import iter_test_modules
from cowboy_lib.repo.source_repo import SourceRepo
from src.runner.local.run_test import run_test

async def test_get_coverage():
    repo_name = "codecovapi-neutered"
    repo_config = get_repo_config(repo_name)
    repo_path = Path(repo_config.source_folder)
    src_repo = SourceRepo(repo_path)
    tm = iter_test_modules(src_repo, lambda tm: tm.name == "TestGithubAppInstallationUsage")[0]

    base_cov = await run_test(repo_name, None)
    chunks = await get_tm_target_coverage(
        repo_name=repo_name,
        src_repo=src_repo,
        tm=tm,
        base_cov=base_cov.get_coverage(),
        run_test=run_test,
        run_args=None
    )

    return chunks

if __name__ == "__main__":
    asyncio.run(test_get_coverage())


--------------------------------------------------------------------------------