two_factors_test.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. // Copyright 2020 The Gogs Authors. All rights reserved.
  2. // Use of this source code is governed by a MIT-style
  3. // license that can be found in the LICENSE file.
  4. package db
  5. import (
  6. "testing"
  7. "time"
  8. "github.com/jinzhu/gorm"
  9. "github.com/stretchr/testify/assert"
  10. "gogs.io/gogs/internal/errutil"
  11. )
  12. func Test_twoFactors(t *testing.T) {
  13. if testing.Short() {
  14. t.Skip()
  15. }
  16. t.Parallel()
  17. tables := []interface{}{new(TwoFactor), new(TwoFactorRecoveryCode)}
  18. db := &twoFactors{
  19. DB: initTestDB(t, "twoFactors", tables...),
  20. }
  21. for _, tc := range []struct {
  22. name string
  23. test func(*testing.T, *twoFactors)
  24. }{
  25. {"Create", test_twoFactors_Create},
  26. {"GetByUserID", test_twoFactors_GetByUserID},
  27. {"IsUserEnabled", test_twoFactors_IsUserEnabled},
  28. } {
  29. t.Run(tc.name, func(t *testing.T) {
  30. t.Cleanup(func() {
  31. err := clearTables(t, db.DB, tables...)
  32. if err != nil {
  33. t.Fatal(err)
  34. }
  35. })
  36. tc.test(t, db)
  37. })
  38. }
  39. }
  40. func test_twoFactors_Create(t *testing.T, db *twoFactors) {
  41. // Create a 2FA token
  42. err := db.Create(1, "secure-key", "secure-secret")
  43. if err != nil {
  44. t.Fatal(err)
  45. }
  46. // Get it back and check the Created field
  47. tf, err := db.GetByUserID(1)
  48. if err != nil {
  49. t.Fatal(err)
  50. }
  51. assert.Equal(t, gorm.NowFunc().Format(time.RFC3339), tf.Created.UTC().Format(time.RFC3339))
  52. // Verify there are 10 recover codes generated
  53. var count int64
  54. err = db.Model(new(TwoFactorRecoveryCode)).Count(&count).Error
  55. if err != nil {
  56. t.Fatal(err)
  57. }
  58. assert.Equal(t, int64(10), count)
  59. }
  60. func test_twoFactors_GetByUserID(t *testing.T, db *twoFactors) {
  61. // Create a 2FA token for user 1
  62. err := db.Create(1, "secure-key", "secure-secret")
  63. if err != nil {
  64. t.Fatal(err)
  65. }
  66. // We should be able to get it back
  67. _, err = db.GetByUserID(1)
  68. if err != nil {
  69. t.Fatal(err)
  70. }
  71. // Try to get a non-existent 2FA token
  72. _, err = db.GetByUserID(2)
  73. expErr := ErrTwoFactorNotFound{args: errutil.Args{"userID": int64(2)}}
  74. assert.Equal(t, expErr, err)
  75. }
  76. func test_twoFactors_IsUserEnabled(t *testing.T, db *twoFactors) {
  77. // Create a 2FA token for user 1
  78. err := db.Create(1, "secure-key", "secure-secret")
  79. if err != nil {
  80. t.Fatal(err)
  81. }
  82. assert.True(t, db.IsUserEnabled(1))
  83. assert.False(t, db.IsUserEnabled(2))
  84. }