diff --git a/tests/Feature/Controllers/BudgetControllerTest.php b/tests/Feature/Controllers/BudgetControllerTest.php
index 5db4645cd2..ea45bcbf0d 100644
--- a/tests/Feature/Controllers/BudgetControllerTest.php
+++ b/tests/Feature/Controllers/BudgetControllerTest.php
@@ -297,7 +297,6 @@ class BudgetControllerTest extends TestCase
         $collector->shouldReceive('setTypes')->andReturnSelf();
         $collector->shouldReceive('withoutBudget')->andReturnSelf();
         $collector->shouldReceive('withOpposingAccount')->andReturnSelf();
-        $collector->shouldReceive('withCategoryInformation')->andReturnSelf();
         $collector->shouldReceive('getPaginatedJournals')->andReturn(new LengthAwarePaginator([], 0, 10));
 
         $date = new Carbon();
@@ -332,7 +331,6 @@ class BudgetControllerTest extends TestCase
         $collector->shouldReceive('withOpposingAccount')->andReturnSelf();
         $collector->shouldReceive('withoutBudget')->andReturnSelf();
         $collector->shouldReceive('setTypes')->andReturnSelf();
-        $collector->shouldReceive('withCategoryInformation')->andReturnSelf();
         $collector->shouldReceive('getJournals')->andReturn(new Collection);
         $collector->shouldReceive('getPaginatedJournals')->andReturn(new LengthAwarePaginator([], 0, 10));
 
@@ -369,7 +367,6 @@ class BudgetControllerTest extends TestCase
         $collector->shouldReceive('withOpposingAccount')->andReturnSelf();
         $collector->shouldReceive('setTypes')->andReturnSelf();
         $collector->shouldReceive('withoutBudget')->andReturnSelf();
-        $collector->shouldReceive('withCategoryInformation')->andReturnSelf();
         $collector->shouldReceive('getPaginatedJournals')->andReturn(new LengthAwarePaginator([], 0, 10));
 
         $date = new Carbon();
@@ -422,8 +419,8 @@ class BudgetControllerTest extends TestCase
         $collector->shouldReceive('setLimit')->andReturnSelf();
         $collector->shouldReceive('setPage')->andReturnSelf();
         $collector->shouldReceive('setBudget')->andReturnSelf();
-        $collector->shouldReceive('withCategoryInformation')->andReturnSelf();
         $collector->shouldReceive('getPaginatedJournals')->andReturn(new LengthAwarePaginator([], 0, 10));
+        $collector->shouldReceive('withBudgetInformation')->andReturnSelf();
 
         $accountRepos = $this->mock(AccountRepositoryInterface::class);
         $accountRepos->shouldReceive('getAccountsByType')->andReturn(new Collection);
@@ -488,7 +485,7 @@ class BudgetControllerTest extends TestCase
         $collector->shouldReceive('setLimit')->andReturnSelf();
         $collector->shouldReceive('setPage')->andReturnSelf();
         $collector->shouldReceive('setBudget')->andReturnSelf();
-        $collector->shouldReceive('withCategoryInformation')->andReturnSelf();
+        $collector->shouldReceive('withBudgetInformation')->andReturnSelf();
         $collector->shouldReceive('getPaginatedJournals')->andReturn(new LengthAwarePaginator([], 0, 10));
 
 
diff --git a/tests/Feature/Controllers/Chart/BudgetControllerTest.php b/tests/Feature/Controllers/Chart/BudgetControllerTest.php
index a0235b978d..6b2eb15417 100644
--- a/tests/Feature/Controllers/Chart/BudgetControllerTest.php
+++ b/tests/Feature/Controllers/Chart/BudgetControllerTest.php
@@ -104,18 +104,13 @@ class BudgetControllerTest extends TestCase
     {
         $generator    = $this->mock(GeneratorInterface::class);
         $collector    = $this->mock(JournalCollectorInterface::class);
-        $accountRepos = $this->mock(AccountRepositoryInterface::class);
         $transactions = factory(Transaction::class, 10)->make();
-        $accounts     = factory(Account::class, 10)->make();
 
         $collector->shouldReceive('setAllAssetAccounts')->once()->andReturnSelf();
-        $collector->shouldReceive('setTypes')->withArgs([[TransactionType::WITHDRAWAL]])->once()->andReturnSelf();
         $collector->shouldReceive('setBudget')->andReturnSelf();
         $collector->shouldReceive('setRange')->andReturnSelf();
         $collector->shouldReceive('getJournals')->andReturn($transactions);
 
-        $accountRepos->shouldReceive('getAccountsByType')->andReturn($accounts)->once();
-
         $generator->shouldReceive('pieChart')->once()->andReturn([]);
 
         $this->be($this->user());
@@ -141,7 +136,6 @@ class BudgetControllerTest extends TestCase
         $categories   = factory(Category::class, 10)->make();
 
         $collector->shouldReceive('setAllAssetAccounts')->once()->andReturnSelf();
-        $collector->shouldReceive('setTypes')->withArgs([[TransactionType::WITHDRAWAL]])->once()->andReturnSelf();
         $collector->shouldReceive('setBudget')->andReturnSelf();
         $collector->shouldReceive('setRange')->andReturnSelf();
         $collector->shouldReceive('withCategoryInformation')->andReturnSelf();