1""" 2Unit tests for the Inventory class. 3""" 4__copyright__ = "Copyright (C) 2014-2016 Martin Blais" 5__license__ = "GNU GPLv2" 6 7import datetime 8import unittest 9import copy 10from datetime import date 11 12from beancount.core.number import D 13from beancount.core.amount import A 14from beancount.core import amount 15from beancount.core.position import Position 16from beancount.core.position import Cost 17from beancount.core.inventory import Inventory 18from beancount.core.inventory import Booking 19from beancount.core import convert 20from beancount.core import position 21from beancount.core import inventory 22from beancount.utils import invariants 23 24 25P = position.from_string 26I = inventory.from_string 27 28 29def setUp(module): 30 invariants.instrument_invariants(Inventory, 31 inventory.check_invariants, 32 inventory.check_invariants) 33 34def tearDown(module): 35 invariants.uninstrument_invariants(Inventory) 36 37 38class TestInventory(unittest.TestCase): 39 40 def checkAmount(self, inventory, number, currency): 41 if isinstance(number, str): 42 number = D(number) 43 amount_ = amount.Amount(number, currency) 44 inv_amount = inventory.get_currency_units(amount_.currency) 45 self.assertEqual(inv_amount, amount_) 46 47 def test_from_string(self): 48 inv = inventory.from_string('') 49 self.assertEqual(Inventory(), inv) 50 51 inv = inventory.from_string('10 USD') 52 self.assertEqual( 53 Inventory([Position(A("10 USD"))]), 54 inv) 55 56 inv = inventory.from_string(' 10.00 USD ') 57 self.assertEqual( 58 Inventory([Position(A("10 USD"))]), 59 inv) 60 61 inv = inventory.from_string('1 USD, 2 CAD') 62 self.assertEqual( 63 Inventory([Position(A("1 USD")), 64 Position(A("2 CAD"))]), 65 inv) 66 67 inv = inventory.from_string('2.2 HOOL {532.43 USD}, 3.413 EUR') 68 self.assertEqual( 69 Inventory([Position(A("2.2 HOOL"), Cost(D('532.43'), 'USD', None, None)), 70 Position(A("3.413 EUR"))]), 71 inv) 72 73 inv = inventory.from_string( 74 '2.2 HOOL {532.43 USD}, 2.3 HOOL {564.00 USD, 2015-07-14}, 3.413 EUR') 75 self.assertEqual( 76 Inventory([Position(A("2.2 HOOL"), 77 Cost(D('532.43'), 'USD', None, None)), 78 Position(A("2.3 HOOL"), 79 Cost(D('564.00'), 'USD', datetime.date(2015, 7, 14), None)), 80 Position(A("3.413 EUR"))]), 81 inv) 82 83 inv = inventory.from_string( 84 '1.1 HOOL {500.00 # 11.00 USD}, 100 CAD') 85 self.assertEqual( 86 Inventory([Position(A("1.1 HOOL"), 87 Cost(D('510.00'), 'USD', None, None)), 88 Position(A("100 CAD"))]), 89 inv) 90 91 def test_ctor_empty_len(self): 92 # Test regular constructor. 93 inv = Inventory() 94 self.assertTrue(inv.is_empty()) 95 self.assertEqual(0, len(inv)) 96 97 inv = Inventory([P('100.00 USD'), 98 P('101.00 USD')]) 99 self.assertFalse(inv.is_empty()) 100 self.assertEqual(1, len(inv)) 101 102 inv = Inventory([P('100.00 USD'), 103 P('100.00 CAD')]) 104 self.assertFalse(inv.is_empty()) 105 self.assertEqual(2, len(inv)) 106 107 inv = Inventory() 108 self.assertEqual(0, len(inv)) 109 inv.add_amount(A('100 USD')) 110 self.assertEqual(1, len(inv)) 111 inv.add_amount(A('100 CAD')) 112 self.assertEqual(2, len(inv)) 113 114 def test_str(self): 115 inv = I('100.00 USD, 101.00 CAD') 116 self.assertEqual('(100.00 USD, 101.00 CAD)', str(inv)) 117 118 def test_copy(self): 119 inv = Inventory() 120 inv.add_amount(A('100.00 USD')) 121 self.checkAmount(inv, '100', 'USD') 122 123 # Test copying. 124 inv2 = copy.copy(inv) 125 inv2.add_amount(A('50.00 USD')) 126 self.checkAmount(inv2, '150', 'USD') 127 128 # Check that the original object is not modified. 129 self.checkAmount(inv, '100', 'USD') 130 131 def test_op_eq(self): 132 inv1 = I('100 USD, 100 CAD') 133 inv2 = I('100 CAD, 100 USD') 134 self.assertEqual(inv1, inv2) 135 self.assertEqual(inv2, inv1) 136 137 inv3 = I('200 USD, 100 CAD') 138 self.assertNotEqual(inv1, inv3) 139 self.assertNotEqual(inv3, inv1) 140 141 inv4 = I('100 USD, 100 JPY') 142 self.assertNotEqual(inv1, inv4) 143 self.assertNotEqual(inv4, inv1) 144 145 inv5 = I('100 JPY, 100 USD') 146 self.assertEqual(inv4, inv5) 147 148 def test_op_lt(self): 149 inv1 = I('100 USD, 100 CAD') 150 inv2 = I('100 CAD, 100 USD') 151 self.assertFalse(inv1 < inv2) 152 self.assertFalse(inv2 < inv1) 153 154 inv3 = I('200 USD, 100 CAD') 155 self.assertTrue(inv1 < inv3) 156 self.assertTrue(inv2 < inv3) 157 self.assertFalse(inv3 < inv1) 158 159 inv4 = I('100 USD, 100 JPY') 160 self.assertTrue(inv1 > inv4) 161 162 inv5 = I('100 JPY, 100 USD') 163 self.assertTrue(inv1 > inv5) 164 self.assertFalse(inv4 < inv5) 165 166 def test_is_small__value(self): 167 test_inv = I('1.50 JPY, 1.51 USD, 1.52 CAD') 168 for inv in test_inv, -test_inv: 169 self.assertFalse(inv.is_small(D('1.49'))) 170 self.assertFalse(inv.is_small(D('1.50'))) 171 self.assertTrue(inv.is_small(D('1.53'))) 172 self.assertTrue(inv.is_small(D('1.52'))) 173 174 def test_is_small__dict(self): 175 test_inv = I('0.03 JPY, 0.003 USD') 176 for inv in test_inv, -test_inv: 177 # Test all four types of inequalities. 178 self.assertTrue(inv.is_small({'JPY': D('0.05'), 'USD': D('0.005')})) 179 self.assertFalse(inv.is_small({'JPY': D('0.005'), 'USD': D('0.0005')})) 180 self.assertTrue(inv.is_small({'JPY': D('0.05'), 'USD': D('0.5')})) 181 self.assertFalse(inv.is_small({'JPY': D('0.005'), 'USD': D('0.005')})) 182 183 # Test border case and an epsilon under. 184 self.assertTrue(inv.is_small({'JPY': D('0.03'), 'USD': D('0.003')})) 185 self.assertFalse(inv.is_small({'JPY': D('0.02999999999999'), 186 'USD': D('0.003')})) 187 self.assertFalse(inv.is_small({'JPY': D('0.03'), 'USD': D('0.00299999')})) 188 189 # Test missing precisions. 190 self.assertFalse(inv.is_small({'JPY': D('0.05')})) 191 self.assertFalse(inv.is_small({'USD': D('0.005')})) 192 193 # Test extra precisions. 194 self.assertTrue(inv.is_small({'JPY': D('0.05'), 195 'USD': D('0.005'), 196 'CAD': D('0.0005')})) 197 198 # Test no precisions. 199 self.assertFalse(inv.is_small({})) 200 201 def test_is_small__with_default(self): 202 inv = I('0.03 JPY') 203 self.assertTrue(inv.is_small({'JPY': D('0.05')})) 204 self.assertFalse(inv.is_small({'JPY': D('0.02')})) 205 206 def test_is_mixed(self): 207 inv = I('100 HOOL {250 USD}, 101 HOOL {251 USD}') 208 self.assertFalse(inv.is_mixed()) 209 210 inv = I('100 HOOL {250 USD}, -1 HOOL {251 USD}') 211 self.assertTrue(inv.is_mixed()) 212 213 inv = I('-2 HOOL {250 USD}, -1 HOOL {251 USD}') 214 self.assertFalse(inv.is_mixed()) 215 216 def test_is_reduced_by(self): 217 # Test with regular all position inventory. 218 inv = I('100 HOOL {250 USD}, 101 HOOL {251 USD}') 219 self.assertFalse(inv.is_reduced_by(A('2 HOOL'))) 220 self.assertFalse(inv.is_reduced_by(A('0 HOOL'))) 221 self.assertTrue(inv.is_reduced_by(A('-2 HOOL'))) 222 223 # Test with a mixed-sign inventory. 224 inv = I('100 HOOL {250 USD}, -101 HOOL {251 USD}') 225 self.assertTrue(inv.is_reduced_by(A('2 HOOL'))) 226 self.assertFalse(inv.is_reduced_by(A('0 HOOL'))) 227 self.assertTrue(inv.is_reduced_by(A('-2 HOOL'))) 228 229 def test_op_neg(self): 230 inv = Inventory() 231 inv.add_amount(A('10 USD')) 232 ninv = -inv 233 self.checkAmount(ninv, '-10', 'USD') 234 235 pinv = I('1.50 JPY, 1.51 USD, 1.52 CAD') 236 ninv = I('-1.50 JPY, -1.51 USD, -1.52 CAD') 237 self.assertEqual(pinv, -ninv) 238 239 def test_op_abs(self): 240 inv = I('1.50 USD, 2.00 USD, -1.52 CAD') 241 self.assertEqual(abs(inv), I('3.50 USD, 1.52 CAD')) 242 243 def test_op_mul(self): 244 inv = I('10 HOOL {1.11 USD}, 2.22 CAD') 245 inv2 = inv * D('3') 246 self.assertEqual(I('30 HOOL {1.11 USD}, 6.66 CAD'), inv2) 247 248 def test_get_only_position(self): 249 inv = I('10 HOOL {1.11 USD}, 2.22 CAD') 250 with self.assertRaises(AssertionError): 251 inv.get_only_position() 252 inv = I('10 HOOL {1.11 USD}') 253 self.assertEqual(A('10 HOOL'), inv.get_only_position().units) 254 inv = I('') 255 self.assertIsNone(inv.get_only_position()) 256 257 def test_get_currency_units(self): 258 inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD') 259 self.assertEqual(inv.get_currency_units('JPY'), A('40.50 JPY')) 260 self.assertEqual(inv.get_currency_units('USD'), A('40.51 USD')) 261 self.assertEqual(inv.get_currency_units('CAD'), A('40.52 CAD')) 262 self.assertEqual(inv.get_currency_units('AUD'), A('0 AUD')) 263 self.assertEqual(inv.get_currency_units('NZD'), A('0 NZD')) 264 265 def test_segregate_units(self): 266 inv = I('2.2 HOOL {532.43 USD}, ' 267 '2.3 HOOL {564.00 USD, 2015-07-14}, ' 268 '3.41 CAD, 101.20 USD') 269 ccymap = inv.segregate_units(['HOOL', 'USD', 'EUR']) 270 self.assertEqual({ 271 None: I('3.41 CAD'), 272 'USD': I('101.20 USD'), 273 'EUR': inventory.Inventory(), 274 'HOOL': I('2.2 HOOL {532.43 USD}, ' 275 '2.3 HOOL {564.00 USD, 2015-07-14}')}, ccymap) 276 277 def test_units1(self): 278 inv = Inventory() 279 self.assertEqual(inv.reduce(convert.get_units), I('')) 280 281 inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD') 282 self.assertEqual(inv.reduce(convert.get_units), 283 I('40.50 JPY, 40.51 USD, 40.52 CAD')) 284 285 # Check that the same units coalesce. 286 inv = I('2 HOOL {400 USD}, 3 HOOL {410 USD}') 287 self.assertEqual(inv.reduce(convert.get_units), I('5 HOOL')) 288 289 inv = I('2 HOOL {400 USD}, -3 HOOL {410 USD}') 290 self.assertEqual(inv.reduce(convert.get_units), I('-1 HOOL')) 291 292 POSITIONS_ALL_KINDS = [ 293 P('40.50 USD'), 294 P('40.50 USD {1.10 CAD}'), 295 P('40.50 USD {1.10 CAD, 2012-01-01}')] 296 297 def test_units(self): 298 inv = Inventory(self.POSITIONS_ALL_KINDS + 299 [P('50.00 CAD')]) 300 inv_cost = inv.reduce(convert.get_units) 301 self.assertEqual(I('121.50 USD, 50.00 CAD'), inv_cost) 302 303 def test_cost(self): 304 inv = Inventory(self.POSITIONS_ALL_KINDS + 305 [P('50.00 CAD')]) 306 inv_cost = inv.reduce(convert.get_cost) 307 self.assertEqual(I('40.50 USD, 139.10 CAD'), inv_cost) 308 309 def test_average(self): 310 # Identity, no aggregation. 311 inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD') 312 self.assertEqual(inv.average(), inv) 313 314 # Identity, no aggregation, with a mix of lots at cost and without cost. 315 inv = I('40 USD {1.01 CAD}, 40 USD') 316 self.assertEqual(inv.average(), inv) 317 318 # Aggregation. 319 inv = I('40 USD {1.01 CAD}, 40 USD {1.02 CAD}') 320 self.assertEqual(inv.average(), I('80.00 USD {1.015 CAD}')) 321 322 # Aggregation, more units. 323 inv = I('2 HOOL {500 USD}, 3 HOOL {520 USD}, 4 HOOL {530 USD}') 324 self.assertEqual(inv.average(), I('9 HOOL {520 USD}')) 325 326 # Average on zero amount, same costs 327 inv = I('2 HOOL {500 USD}') 328 inv.add_amount(A('-2 HOOL'), Cost(D('500'), 'USD', None, None)) 329 self.assertEqual(inv.average(), I('')) 330 331 # Average on zero amount, different costs 332 inv = I('2 HOOL {500 USD}') 333 inv.add_amount(A('-2 HOOL'), 334 Cost(D('500'), 'USD', datetime.date(2000, 1, 1), None)) 335 self.assertEqual(inv.average(), I('')) 336 337 def test_currencies(self): 338 inv = Inventory() 339 self.assertEqual(set(), inv.currencies()) 340 341 inv = I('40 USD {1.01 CAD}, 40 USD') 342 self.assertEqual({'USD'}, inv.currencies()) 343 344 inv = I('40 AAPL {1.01 USD}, 10 HOOL {2.02 USD}') 345 self.assertEqual({'AAPL', 'HOOL'}, inv.currencies()) 346 347 def test_currency_pairs(self): 348 inv = Inventory() 349 self.assertEqual(set(), inv.currency_pairs()) 350 351 inv = I('40 USD {1.01 CAD}, 40 USD') 352 self.assertEqual(set([('USD', 'CAD'), ('USD', None)]), inv.currency_pairs()) 353 354 inv = I('40 AAPL {1.01 USD}, 10 HOOL {2.02 USD}') 355 self.assertEqual(set([('AAPL', 'USD'), ('HOOL', 'USD')]), inv.currency_pairs()) 356 357 def test_add_amount(self): 358 inv = Inventory() 359 inv.add_amount(A('100.00 USD')) 360 self.checkAmount(inv, '100', 'USD') 361 362 # Add some amount 363 inv.add_amount(A('25.01 USD')) 364 self.checkAmount(inv, '125.01', 'USD') 365 366 # Subtract some amount. 367 inv.add_amount(A('-12.73 USD')) 368 self.checkAmount(inv, '112.28', 'USD') 369 370 # Subtract some to be negative (should be allowed if no lot). 371 inv.add_amount(A('-120 USD')) 372 self.checkAmount(inv, '-7.72', 'USD') 373 374 # Subtract some more. 375 inv.add_amount(A('-1 USD')) 376 self.checkAmount(inv, '-8.72', 'USD') 377 378 # Add to above zero again 379 inv.add_amount(A('18.72 USD')) 380 self.checkAmount(inv, '10', 'USD') 381 382 def test_add_amount__zero(self): 383 inv = Inventory() 384 inv.add_amount(A('0 USD')) 385 self.assertEqual(0, len(inv)) 386 387 def test_add_amount__booking(self): 388 inv = Inventory() 389 _, booking = inv.add_amount(A('100.00 USD')) 390 self.assertEqual(Booking.CREATED, booking) 391 392 _, booking = inv.add_amount(A('20.00 USD')) 393 self.assertEqual(Booking.AUGMENTED, booking) 394 395 _, booking = inv.add_amount(A('-20 USD')) 396 self.assertEqual(Booking.REDUCED, booking) 397 398 _, booking = inv.add_amount(A('-100 USD')) 399 self.assertEqual(Booking.REDUCED, booking) 400 401 def test_add_amount__multi_currency(self): 402 inv = Inventory() 403 inv.add_amount(A('100 USD')) 404 inv.add_amount(A('100 CAD')) 405 self.checkAmount(inv, '100', 'USD') 406 self.checkAmount(inv, '100', 'CAD') 407 408 inv.add_amount(A('25 USD')) 409 self.checkAmount(inv, '125', 'USD') 410 self.checkAmount(inv, '100', 'CAD') 411 412 def test_add_amount__withlots(self): 413 # Testing the strict case where everything matches, with only a cost. 414 inv = Inventory() 415 inv.add_amount(A('50 HOOL'), Cost(D('700'), 'USD', None, None)) 416 self.checkAmount(inv, '50', 'HOOL') 417 418 inv.add_amount(A('-40 HOOL'), Cost(D('700'), 'USD', None, None)) 419 self.checkAmount(inv, '10', 'HOOL') 420 421 position_, _ = inv.add_amount(A('-12 HOOL'), 422 Cost(D('700'), 'USD', None, None)) 423 self.assertTrue(next(iter(inv)).is_negative_at_cost()) 424 425 # Testing the strict case where everything matches, a cost and a lot-date. 426 inv = Inventory() 427 inv.add_amount(A('50 HOOL'), Cost(D('700'), 'USD', date(2000, 1, 1), None)) 428 self.checkAmount(inv, '50', 'HOOL') 429 430 inv.add_amount(A('-40 HOOL'), Cost(D('700'), 'USD', date(2000, 1, 1), None)) 431 self.checkAmount(inv, '10', 'HOOL') 432 433 position_, _ = inv.add_amount(A('-12 HOOL'), Cost(D('700'), 'USD', 434 date(2000, 1, 1), None)) 435 self.assertTrue(next(iter(inv)).is_negative_at_cost()) 436 437 def test_add_amount__allow_negative(self): 438 inv = Inventory() 439 440 # Test adding positions of different types. 441 position_, _ = inv.add_amount(A('-11 USD')) 442 self.assertIsNone(position_) 443 position_, _ = inv.add_amount(A('-11 USD'), 444 Cost(D('1.10'), 'CAD', None, None)) 445 self.assertIsNone(position_) 446 position_, _ = inv.add_amount(A('-11 USD'), 447 Cost(D('1.10'), 'CAD', date(2012, 1, 1), None)) 448 self.assertIsNone(position_) 449 450 # Check for reductions. 451 invlist = list(inv) 452 self.assertTrue(invlist[1].is_negative_at_cost()) 453 self.assertTrue(invlist[2].is_negative_at_cost()) 454 inv.add_amount(A('-11 USD'), Cost(D('1.10'), 'CAD', None, None)) 455 inv.add_amount(A('-11 USD'), Cost(D('1.10'), 'CAD', date(2012, 1, 1), None)) 456 self.assertEqual(3, len(inv)) 457 458 # Test adding to a position that does exist. 459 inv = I('10 USD, 10 USD {1.10 CAD}, 10 USD {1.10 CAD, 2012-01-01}') 460 position_, _ = inv.add_amount(A('-11 USD')) 461 self.assertEqual(position_, position.from_string('10 USD')) 462 position_, _ = inv.add_amount(A('-11 USD'), 463 Cost(D('1.10'), 'CAD', None, None)) 464 self.assertEqual(position_, position.from_string('10 USD {1.10 CAD}')) 465 position_, _ = inv.add_amount(A('-11 USD'), 466 Cost(D('1.10'), 'CAD', date(2012, 1, 1), None)) 467 self.assertEqual(position_, position.from_string('10 USD {1.10 CAD, 2012-01-01}')) 468 469 def test_add_position(self): 470 inv = Inventory() 471 for pos in self.POSITIONS_ALL_KINDS: 472 inv.add_position(pos) 473 self.assertEqual(Inventory(self.POSITIONS_ALL_KINDS), inv) 474 475 def test_op_add(self): 476 inv1 = I('17.00 USD') 477 orig_inv1 = I('17.00 USD') 478 inv2 = I('21.00 CAD') 479 inv3 = inv1 + inv2 480 self.assertEqual(I('17.00 USD, 21.00 CAD'), inv3) 481 self.assertEqual(orig_inv1, inv1) 482 483 def test_update(self): 484 inv1 = I('11 USD') 485 inv2 = I('12 CAD') 486 inv_updated = inv1.add_inventory(inv2) 487 expect_updated = I('11 USD, 12 CAD') 488 self.assertEqual(expect_updated, inv_updated) 489 self.assertEqual(expect_updated, inv1) 490 491 def test_sum_inventories(self): 492 inv1 = Inventory() 493 inv1.add_amount(A('10 USD')) 494 495 inv2 = Inventory() 496 inv2.add_amount(A('20 CAD')) 497 inv2.add_amount(A('55 HOOL')) 498 499 _ = inv1 + inv2 500 501 def test_reduce(self): 502 inv = I('100.00 USD, 101.00 CAD, 100 HOOL {300.00 USD}') 503 inv_units = inv.reduce(lambda posting: posting.units) 504 self.assertEqual(I('100.00 USD, 101.00 CAD, 100 HOOL'), inv_units) 505 506 507if __name__ == '__main__': 508 unittest.main() 509