Fix errors

Author: rohitvinnakota-codecov

src/seer/automation/codegen/codegen_context.py

@@ -90,3 +90,14 @@ def get_unit_test_memory(
                 .one_or_none()
             )
         return pr_context
+
+    def store_unit_test_memory(
+        self, owner: str, repo: str, pr_id: int
+    ) -> DbPrContextToUnitTestGenerationRunIdMapping | None:
+        with Session() as session:
+            pr_context = (
+                session.query(DbPrContextToUnitTestGenerationRunIdMapping)
+                .filter_by(owner=owner, repo=repo, pr_id=pr_id)
+                .one_or_none()
+            )
+        return pr_context

src/seer/automation/codegen/unit_test_coding_component.py

@@ -34,4 +34,85 @@ class UnitTestCodingComponent(BaseComponent[CodeUnitTestRequest, CodeUnitTestOut
     def invoke(
         self, request: CodeUnitTestRequest, llm_client: LlmClient = injected
     ) -> CodeUnitTestOutput | None:
-        pass
+        with BaseTools(self.context, repo_client_type=RepoClientType.CODECOV_UNIT_TEST) as tools:
+            agent = LlmAgent(
+                tools=tools.get_tools(),
+                config=AgentConfig(interactive=False),
+            )
+
+            codecov_client_params = request.codecov_client_params
+
+            code_coverage_data = CodecovClient.fetch_coverage(
+                repo_name=codecov_client_params["repo_name"],
+                pullid=codecov_client_params["pullid"],
+                owner_username=codecov_client_params["owner_username"],
+            )
+
+            test_result_data = CodecovClient.fetch_test_results_for_commit(
+                repo_name=codecov_client_params["repo_name"],
+                owner_username=codecov_client_params["owner_username"],
+                latest_commit_sha=codecov_client_params["head_sha"],
+            )
+
+            existing_test_design_response = llm_client.generate_text(
+                model=AnthropicProvider.model("claude-3-7-sonnet@20250219"),
+                prompt=CodingUnitTestPrompts.format_find_unit_test_pattern_step_msg(
+                    diff_str=request.diff
+                ),
+            )
+
+            llm_client.generate_text(
+                model=AnthropicProvider.model("claude-3-7-sonnet@20250219"),
+                prompt=CodingUnitTestPrompts.format_plan_step_msg(
+                    diff_str=request.diff,
+                    has_coverage_info=code_coverage_data,
+                    has_test_result_info=test_result_data,
+                ),
+            )
+
+            final_response = agent.run(
+                run_config=RunConfig(
+                    prompt=CodingUnitTestPrompts.format_unit_test_msg(
+                        diff_str=request.diff, test_design_hint=existing_test_design_response
+                    ),
+                    system_prompt=CodingUnitTestPrompts.format_system_msg(),
+                    model=AnthropicProvider.model("claude-3-7-sonnet@20250219"),
+                    run_name="Generate Unit Tests",
+                ),
+            )
+
+            if not final_response:
+                return None
+            plan_steps_content = extract_text_inside_tags(final_response, "plan_steps")
+
+            if len(plan_steps_content) == 0:
+                raise ValueError("Failed to extract plan_steps from the planning step of LLM")
+
+            coding_output = PlanStepsPromptXml.from_xml(
+                f"<plan_steps>{escape_multi_xml(plan_steps_content, ['diff', 'description', 'commit_message'])}</plan_steps>"
+            ).to_model()
+
+        if not coding_output.tasks:
+            raise ValueError("No tasks found in coding output")
+        file_changes: list[FileChange] = []
+        for task in coding_output.tasks:
+            repo_client = self.context.get_repo_client(
+                task.repo_name, type=RepoClientType.CODECOV_UNIT_TEST
+            )
+            if task.type == "file_change":
+                file_content, _ = repo_client.get_file_content(task.file_path)
+                if not file_content:
+                    logger.warning(f"Failed to get content for {task.file_path}")
+                    continue
+
+                changes, _ = task_to_file_change(task, file_content)
+                file_changes += changes
+            elif task.type == "file_delete":
+                change = task_to_file_delete(task)
+                file_changes.append(change)
+            elif task.type == "file_create":
+                change = task_to_file_create(task)
+                file_changes.append(change)
+            else:
+                logger.warning(f"Unsupported task type: {task.type}")
+        return CodeUnitTestOutput(diffs=file_changes)

src/seer/automation/codegen/unit_test_github_pr_creator.py

@@ -60,9 +60,9 @@ def create_github_pull_request(self):
         new_pr_url = new_pr.html_url
         self.repo_client.post_unit_test_reference_to_original_pr(original_pr_url, new_pr_url)
 
-    def store_pr_context(self, new_pr: PullRequest, original_pr: PullRequest):
+    def store_pr_context(self, new_pr: PullRequest):
         with Session() as session:
-            pr_id_mapping = DbPrContextToUnitTestGenerationRunIdMapping(
+            run_info = DbPrContextToUnitTestGenerationRunIdMapping(
                 provider="github",
                 owner=self.repo_client.repo.owner.login,
                 repo=self.repo_client.repo.name,
@@ -71,6 +71,5 @@ def store_pr_context(self, new_pr: PullRequest, original_pr: PullRequest):
                 run_id=self.unit_test_run_id,
                 original_pr_url=self.pr.html_url,
             )
-            session.add(pr_id_mapping)
+            session.add(run_info)
             session.commit()
-            print("SAVED TO DB")