updated tests for new feature

Tests for your PR