improve assignments controller

app/controllers/assignments_controller.rb

@@ -4,10 +4,18 @@ class AssignmentsController < ApplicationController
   before_action :set_assignment, only: %i[show update destroy]
 
   def index
-    project = Project.find(requirements_filter[:project_id])
-    statement_of_works_ids = project.statement_of_works.map(&:id)
+    if assignments_filter[:project_id].present?
+      project = Project.find(assignments_filter[:project_id])
+      statement_of_works_ids = project.statement_of_works.map(&:id)
+    elsif assignments_filter[:statement_of_work_id].present?
+      statement_of_works_ids = [assignments_filter[:statement_of_work_id]]
+    else
+      render json: { error: 'project_id or statement_of_work_id is required' }, status: :bad_request
+      return
+    end
+
     requirements = Requirement.where(statement_of_work_id: statement_of_works_ids).active_in_period(
-      requirements_filter[:start_date], requirements_filter[:end_date]
+      assignments_filter[:start_date], assignments_filter[:end_date]
     )
 
     @assignments = Assignment.where(requirement: requirements)
@@ -47,7 +55,7 @@ def assignment_params
     params.require(:assignment).permit(:coverage, :requirement_id, :start_date, :end_date, :user_id)
   end
 
-  def requirements_filter
-    params.permit(:project_id, :start_date, :end_date)
+  def assignments_filter
+    params.permit(:project_id, :statement_of_work_id, :start_date, :end_date)
   end
 end

spec/controllers/assignments_controller_spec.rb

@@ -13,15 +13,34 @@
   end
 
   describe 'GET #index' do
-    it 'returns a success response' do
-      requirement = create(:requirement)
-      project = requirement.statement_of_work.project
+    let(:project) { create(:project) }
+    let(:statement_of_work) { create(:statement_of_work, :with_maintenance, project:) }
+    let(:secondary_statement_of_work) { create(:statement_of_work, project:) }
+    let(:requirement) { create(:requirement, statement_of_work:) }
+    let(:secondary_requirement) { create(:requirement, statement_of_work: secondary_statement_of_work) }
 
+    before do
       create_list(:assignment, 2, coverage: 0.5, requirement:)
-      get :index,
-          params: { start_date: Time.zone.today - 6.days, end_date: Time.zone.today + 6.days, project_id: project.id }
+      create_list(:assignment, 5, coverage: 0.5, requirement: secondary_requirement)
+    end
+
+    context 'when given the project id' do
+      it 'returns all the assignments of that project' do
+        get :index,
+            params: { start_date: Time.zone.today - 6.days, end_date: Time.zone.today + 6.days, project_id: project.id }
+
+        expect(response.parsed_body.size).to eq(7)
+      end
+    end
 
-      expect(response.parsed_body.size).to eq(2)
+    context 'when given the statement of work id' do
+      it 'returns all the requirements of that statement of work' do
+        get :index,
+            params: { start_date: Time.zone.today - 6.days, end_date: Time.zone.today + 6.days,
+                      statement_of_work_id: secondary_statement_of_work.id }
+
+        expect(response.parsed_body.size).to eq(5)
+      end
     end
   end